Gemma中RoPE代码详细讲解

本文详细解读了RoPE(RotaryPositionEmbedding)在Torch中的实现过程,重点介绍了预计算频率cis和应用旋转嵌入的函数,对比了与LLama代码的不同之处,指出其简洁性。

摘要生成于 C知道 ,由 DeepSeek-R1 满血版支持, 前往体验 >

最近在看Gemma代码感觉比LLama的代码看的方便点, 看到RoPE代码跟常规的方式不太一样(也不算常规,就是我理解的方式),特此记录一下。我的RoPE入门代码参考:Rotary Position Embedding (RoPE, 旋转式位置编码) | 原理讲解+torch代码实现
原理我就不讲了,直接贴一下图,图源自于上面的链接。
在这里插入图片描述
我们先粘贴一下代码,逐步讲解:

dim:单头维度信息
end:序列长度
theta:10000
def precompute_freqs_cis(dim: int,
                         end: int,
                         theta: float = 10000.0) -> torch.Tensor:
    """Precomputes the frequency cis."""
    freqs = 1.0 / (theta**(torch.arange(0, dim, 2)[:(dim // 2)].float() / dim))
    t = torch.arange(end, device=freqs.device)
    freqs = torch.outer(t, freqs).float()
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64
    return freqs_cis

x:输入特征维度[batch, end, num_head, dim]
freqs_cis:上个函数获取的结果
def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
    """Applies the rotary embedding to the query and key tensors."""
    x_ = torch.view_as_complex(
        torch.stack(torch.chunk(x.transpose(1, 2).float(), 2, dim=-1),
                    dim=-1))
    x_out = torch.view_as_real(x_ * freqs_cis).type_as(x)
    x_out = torch.cat(torch.chunk(x_out, 2, dim=-1), dim=-2)
    x_out = x_out.reshape(x_out.shape[0], x_out.shape[1], x_out.shape[2],
                          -1).transpose(1, 2)
    return x_out

precompute_freqs_cis

  • freqs = 1.0 / (theta**(torch.arange(0, dim, 2)[:(dim // 2)].float() / dim))
    这个代码主要是实现一下公式在这里插入图片描述
    torch.arange(0, dim, 2),生成列表 [0,2, …d_model//2]
  • t = torch.arange(end, device=freqs.device)
    生成序列长度, [0, 1, …, end(也就是序列长度)]
  • freqs = torch.outer(t, freqs).float()
    进行笛卡尔积,维度变成[end, dim//2]
    在这里插入图片描述
  • freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
    通过polar函数生成cos和sin值,为什么要使用torch.ones_like(freqs), 下面公式,abs为1,不就是cos值和sin值了
    在这里插入图片描述

apply_rotary_emb

  • x_ = torch.view_as_complex(
    torch.stack(torch.chunk(x.transpose(1, 2).float(), 2, dim=-1),
    dim=-1))
    x.transpose(1, 2).float()将输入维度变为[batch, num_head, end, dim]
    torch.chunk将数据前dim//2 和后dim//2分开,我理解的是[q0, q1, …qn]是奇偶分开,而不是前后分开,可能无所谓吧。
    torch.stack则是对维度进行合并,产生[batch, num_head, end, dim//2, 2]这种维度。
    我简单举例子验证一下:
import torch
a = torch.arange(10)
print(a)
b = torch.chunk(a, 2, dim=-1)
print(b)
c = torch.stack(b, dim = -1)
print(c)

在这里插入图片描述
并且用torch.view_as_complex转为复述的形式

  • x_out = torch.view_as_real(x_ * freqs_cis).type_as(x)
    x_ * freqs_cis实现复述的预算
    举例子解释一下:
    x_为:q0+q1 i
    freqs_cis:cos+sin i
    实部为:q0cos -q1sin
    虚部为:q1cos + q1sin
    torch.view_as_real函数则是把转为实数的形式,a+bi->[a, b]形式
  • x_out = torch.cat(torch.chunk(x_out, 2, dim=-1), dim=-2)
    则是把维度转为[batch, num_head, end, dim]的形式
  • x_out = x_out.reshape(x_out.shape[0], x_out.shape[1], x_out.shape[2],
    -1).transpose(1, 2)
    转为输入的形式

下面为LLama的RoPE实现:

def apply_rotary_emb(
    xq: torch.Tensor,
    xk: torch.Tensor,
    freqs_cis: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
    return xq_out.type_as(xq), xk_out.type_as(xk)

LLama我理解,他则是采用类似于奇偶分开的方式,我简单尝试了一下:

import torch
a = torch.arange(10)
print(a)
b = a.reshape(5,2)
print(b)

在这里插入图片描述

总结:

以上就是我对RoPE代码实现的理解,相比原来理解的方式,这种相对更加简洁,但是略有一些绕

<think>嗯,用户想了解在GWAS中使用GEMMA的方法。首先,我需要回忆一下GWAS和GEMMA的基本概念。GWAS是全基因组关联分析,用来寻找基因变异与表型之间的关联。GEMMA是一个用于进行混合模型分析的软件,特别适合处理样本间的亲缘关系,比如利用线性混合模型(LMM)来减少假阳性结果。 用户可能已经知道GWAS的基础,但需要具体操作GEMMA的步骤。我应该先介绍GEMMA的优势,比如处理群体结构、亲缘关系,以及它的快速算法。然后分步骤说明如何使用,包括数据准备、输入文件格式,比如基因型数据通常用PLINK格式,表型数据需要特定的格式。 接下来,运行GEMMA的关键步骤,比如先计算亲缘关系矩阵,然后进行关联分析。这里需要详细说明每个步骤的命令行参数,比如-bslmm选项或者-lmm的不同类型。同时,可能需要提醒用户注意输入文件的路径和格式是否正确,避免常见错误。 结果解读部分,用户可能需要知道输出文件的内容,比如每个SNP的p值、效应量等。这里要解释p值的显著性,以及如何确定阈值,比如Bonferroni校正。此外,GEMMA可能提供后验概率,需要说明其意义。 注意事项方面,数据质量控制很重要,比如MAF过滤、缺失数据处理。计算资源方面,GEMMA虽然高效,但大数据集可能需要较多内存和时间。最后,结合其他工具如GCTA或FUMA进行功能注释,帮助用户更全面地分析结果。 需要确保回答结构清晰,步骤明确,同时提醒可能遇到的问题和解决方法。语言要简洁,避免过于技术化,但保持专业性。可能还需要推荐查阅GEMMA的官方文档,或者提供示例命令,让用户更容易上手。</think>在基因组关联分析(GWAS)中,GEMMA(Genome-wide Efficient Mixed Model Association)是一款高效的混合线性模型工具,特别适用于校正群体结构和亲缘关系的影响。以下是使用GEMMA进行GWAS的分步说明和关键要点: --- ### **一、GEMMA的核心优势** 1. **混合模型**:通过加入随机效应项(如亲缘关系矩阵)控制复杂样本结构,减少假阳性。 2. **快速算法**:利用Eigen分解和稀疏矩阵优化,计算效率显著高于传统LMM(线性混合模型)。 3. **灵活分析**:支持连续型和二分类表型,提供LMM、BSLMM(贝叶斯稀疏线性混合模型)等多种模型。 --- ### **二、数据准备** #### 输入文件格式: - **基因型数据**:需转换为PLINK二进制格式(`.bed`, `.bim`, `.fam`)。 - **表型文件**:文本文件(`.txt`或`.csv`),每列对应一个表型,缺失值用`NA`表示。 - **协变量文件**(可选):如年龄、性别等固定效应,需标准化处理。 --- ### **三、运行GEMMA的关键步骤** #### 步骤1:计算亲缘关系矩阵 ```bash gemma -g [genotype.bed] -p [phenotype.txt] -gk 1 -o [output_prefix] ``` - `-gk 1`:生成标准化的亲缘关系矩阵(K矩阵)。 - 输出文件:`[output_prefix].cXX.txt`。 #### 步骤2:执行关联分析(以LMM为例) ```bash gemma -g [genotype.bed] -p [phenotype.txt] -k [output_prefix.cXX.txt] \ -lmm 4 -maf 0.01 -o [result_prefix] ``` - `-lmm 4`:指定LMM模型,数字对应不同检验方法(1-4为Wald/LRT/Score等)。 - `-maf 0.01`:过滤次要等位基因频率(MAF)低于1%的位点。 --- ### **四、结果解读** #### 主要输出文件: - `[result_prefix].assoc.txt`:包含每个SNP的关联统计量(如p值、效应量、标准误)。 - **关键列**: - `beta`:SNP效应值。 - `se`:标准误。 - `p_wald`或`p_lrt`:显著性p值。 - `p_lrt_bonf`:Bonferroni校正后p值。 #### 显著性阈值: - 通常采用全基因组显著性阈值 $$ p < 5 \times 10^{-8} $$。 - 对BSLMM结果,关注后验概率(如 `pip > 0.05`)。 --- ### **五、注意事项** 1. **数据质控**: - 过滤低质量SNP(缺失率>5%、哈迪-温伯格平衡p值<1e-6)。 - 标准化表型数据(如去除离群值)。 2. **计算资源**: - 10万样本量需约16GB内存,计算时间随样本数平方增长。 3. **模型选择**: - 简单结构用LMM,复杂遗传架构可尝试BSLMM。 4. **结果验证**: - 通过QQ图检查p值分布(理想情况:大部分点沿对角线分布)。 - 曼哈顿图中显著信号需在独立样本中复现。 --- ### **六、扩展应用** - **多表型分析**:逐一运行不同表型文件或使用并行脚本。 - **条件分析**:在显著位点处加入协变量,检验剩余信号。 - **功能注释**:结合FUMA或ANNOVAR对显著SNP进行功能注释。 --- ### **七、示例代码** ```bash # 计算亲缘矩阵 gemma -g chr1.bed -p pheno.txt -gk 1 -o chr1_kinship # 运行LMM分析(Wald检验) gemma -g chr1.bed -p pheno.txt -k chr1_kinship.cXX.txt \ -lmm 4 -maf 0.01 -o chr1_gwas ``` --- 如需进一步优化分析,可参考GEMMA官方文档(https://github.com/genetics-statistics/GEMMA)或结合GCTA、PLINK进行数据预处理。
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值