如何在百度上搜到网站,广州网站优化方案,哪家企业网页制作好,管理咨询公司需要的能力前言
在自然语言处理领域#xff0c;随着大语言模型#xff08;LLMs#xff09;不断拓展其阅读、理解和生成文本的能力#xff0c;如何高效处理长文本成为一项关键挑战。近日#xff0c;Moonshot AI Research 联合清华大学、浙江大学的研究人员提出了一种创新方法 —— 混…前言
在自然语言处理领域随着大语言模型LLMs不断拓展其阅读、理解和生成文本的能力如何高效处理长文本成为一项关键挑战。近日Moonshot AI Research 联合清华大学、浙江大学的研究人员提出了一种创新方法 —— 混合块注意力机制Mixture of Block AttentionMoBA它将专家混合Mixture of ExpertsMoE原理应用于注意力机制为解决长文本处理难题带来了新的思路。 在 Transformer 架构广泛应用的当下其注意力机制存在明显弊端。在处理长文本时传统注意力机制需将每个 token 与其他所有 token 进行比较这使得计算成本随序列长度呈二次方增长。当模型处理长篇文档、多章书籍、法律简报或大型代码库等包含大量文本信息的任务时这种计算成本会变得难以承受。此前为解决这一问题研究人员尝试过多种方法。例如滑动窗口机制将 token 限制在局部邻域内虽降低了计算量但会忽略重要的全局关系而一些彻底改变基本架构的方法如用全新结构替代 softmax 注意力机制往往需要从头开始重新训练模型难以利用现有的预训练成果。 核心原理
MoBA 的出现有效弥补了上述方法的不足。它的核心在于将输入划分为易于管理的 “块”并借助可训练的门控系统来确定每个查询 token 相关的块。这种设计遵循 “少结构” 原则不预先定义哪些 token 应该相互作用而是由学习到的门控网络做出决策。与固定结构或近似处理的方法不同MoBA 能让模型自主学习注意力的聚焦点。而且MoBA 可与现有的基于 Transformer 的模型无缝协作它作为一种 “插件” 或替代方案保持与原模型相同的参数数量避免架构膨胀同时保留因果掩码确保自回归生成的准确性。在实际应用中MoBA 能在稀疏注意力和全注意力之间灵活切换。处理超长输入时稀疏注意力可提升速度而在训练的某些层或阶段若需要全注意力模型也能切换回标准模式。 从技术细节来看MoBA 将上下文划分为多个块每个块包含连续的 token 序列。门控机制通过比较查询 token 与块的池化键表示计算查询 token 与每个块之间的 “亲和度” 分数然后选择得分最高的块。这样只有最相关块中的 token 才会对最终的注意力分布产生影响。同时包含查询 token 本身的块始终被纳入以确保局部上下文信息可访问。并且MoBA 执行因果掩码防止 token 关注未来位置维持从左到右的自回归属性。这种基于块的方法大幅减少了 token 比较次数使计算规模低于二次方随着上下文长度增加到数十万甚至数百万个 token效率提升愈发显著。此外MoBA 与现代加速器和专用内核兼容性良好。研究人员将 MoBA 与 FlashAttention一种高性能的快速、内存高效的精确注意力库相结合根据所选块对查询 - 键 - 值操作进行精心分组进一步优化了计算流程。实验数据显示在处理一百万个 token 时MoBA 相比传统全注意力机制速度提升约 6 倍凸显了其在实际应用中的优势。 在性能测试方面MoBA 表现出色。技术报告显示在多种任务中MoBA 的性能与全注意力机制相当但在处理长序列时可显著节省计算资源。在语言建模数据测试中当序列长度为 8192 或 32768 个 token 时MoBA 的困惑度与全注意力 Transformer 相近。更为关键的是当研究人员将上下文长度逐渐扩展到 128000 及更长时MoBA 仍能保持强大的长上下文理解能力。在 “尾随 token” 评估中MoBA 能够有效处理长提示末尾附近的 token 预测任务且预测质量没有明显下降。研究人员还对 MoBA 的块大小和门控策略进行了敏感性探索。实验表明细化粒度使用更小的块但选择更多的块有助于模型更接近全注意力的效果。即使在忽略大部分上下文的情况下自适应门控也能识别与查询真正相关的块。此外“混合” 模式展现出一种平衡策略部分层继续使用 MoBA 提升速度少数层则恢复全注意力。这种混合方法在监督微调任务中尤为有益例如当输入中的某些位置在训练目标中被屏蔽时保留少数上层的全注意力可使模型保持广泛的上下文覆盖有助于需要全局视角的任务。
关键代码分析
以下是对 MoBA 库关键代码 MixedAttention 类的分析以及关键代码的摘录与注释
整体分析
MixedAttention 类是一个自定义的 torch.autograd.Function用于实现混合块注意力机制。这个类主要包含两个静态方法forward 和 backward分别用于前向传播和反向传播。
class MixedAttention(torch.autograd.Function):# 前向传播函数staticmethoddef forward(ctx,q, # 查询张量k, # 键张量v, # 值张量self_attn_cu_seqlen, # 自注意力累积序列长度moba_q, # MoBA 查询张量moba_kv, # MoBA 键值张量moba_cu_seqlen_q, # MoBA 查询累积序列长度moba_cu_seqlen_kv, # MoBA 键值累积序列长度max_seqlen, # 最大序列长度moba_chunk_size, # MoBA 块大小moba_q_sh_indices, # MoBA 查询块索引):# 保存一些参数用于后续的反向传播ctx.max_seqlen max_seqlenctx.moba_chunk_size moba_chunk_sizectx.softmax_scale softmax_scale q.shape[-1] ** (-0.5)# 自注意力计算_, _, _, _, self_attn_out_sh, self_attn_lse_hs, _, _ (_flash_attn_varlen_forward(qq,kk,vv,cu_seqlens_qself_attn_cu_seqlen,cu_seqlens_kself_attn_cu_seqlen,max_seqlen_qmax_seqlen,max_seqlen_kmax_seqlen,softmax_scalesoftmax_scale,causalTrue,dropout_p0.0,))# MoBA 注意力计算_, _, _, _, moba_attn_out, moba_attn_lse_hs, _, _ _flash_attn_varlen_forward(qmoba_q,kmoba_kv[:, 0],vmoba_kv[:, 1],cu_seqlens_qmoba_cu_seqlen_q,cu_seqlens_kmoba_cu_seqlen_kv,max_seqlen_qmax_seqlen,max_seqlen_kmoba_chunk_size,softmax_scalesoftmax_scale,causalFalse,dropout_p0.0,)# 转换 lse 形状从 hs 转换为 sh遵循传统混合注意力逻辑self_attn_lse_sh self_attn_lse_hs.t().contiguous()moba_attn_lse moba_attn_lse_hs.t().contiguous()# 初始化输出缓冲区形状与 q 相同output torch.zeros((q.shape[0], q.shape[1], q.shape[2]), deviceq.device, dtypetorch.float32)# 将输出张量展平为二维便于后续索引操作output_2d output.view(-1, q.shape[2])# 计算混合 lse# 减去最大 lse 以避免指数爆炸max_lse_1d self_attn_lse_sh.view(-1)max_lse_1d max_lse_1d.index_reduce(0, moba_q_sh_indices, moba_attn_lse.view(-1), amax)self_attn_lse_sh self_attn_lse_sh - max_lse_1d.view_as(self_attn_lse_sh)moba_attn_lse (moba_attn_lse.view(-1).sub(max_lse_1d.index_select(0, moba_q_sh_indices)).reshape_as(moba_attn_lse))# 计算自注意力和 MoBA 注意力的 softmax 结果mixed_attn_se_sh self_attn_lse_sh.exp()moba_attn_se moba_attn_lse.exp()# 将 MoBA 注意力结果累加到自注意力结果上mixed_attn_se_sh.view(-1).index_add_(0, moba_q_sh_indices, moba_attn_se.view(-1))mixed_attn_lse_sh mixed_attn_se_sh.log()# 加权自注意力输出factor (self_attn_lse_sh - mixed_attn_lse_sh).exp() # [ vS, H ]self_attn_out_sh self_attn_out_sh * factor.unsqueeze(-1)output_2d self_attn_out_sh.reshape_as(output_2d)# 加权 MoBA 输出mixed_attn_lse (mixed_attn_lse_sh.view(-1).index_select(0, moba_q_sh_indices).view_as(moba_attn_lse))factor (moba_attn_lse - mixed_attn_lse).exp() # [ vS, H ]moba_attn_out moba_attn_out * factor.unsqueeze(-1)raw_attn_out moba_attn_out.view(-1, moba_attn_out.shape[-1])output_2d.index_add_(0, moba_q_sh_indices, raw_attn_out)# 将输出转换为与输入相同的数据类型output output.to(q.dtype)# 恢复最大 lsemixed_attn_lse_sh mixed_attn_lse_sh max_lse_1d.view_as(mixed_attn_se_sh)# 保存中间结果用于反向传播ctx.save_for_backward(output,mixed_attn_lse_sh,q,k,v,self_attn_cu_seqlen,moba_q,moba_kv,moba_cu_seqlen_q,moba_cu_seqlen_kv,moba_q_sh_indices,)return output# 反向传播函数staticmethoddef backward(ctx, d_output):# 从上下文中获取保存的参数max_seqlen ctx.max_seqlenmoba_chunk_size ctx.moba_chunk_sizesoftmax_scale ctx.softmax_scale(output,mixed_attn_vlse_sh,q,k,v,self_attn_cu_seqlen,moba_q,moba_kv,moba_cu_seqlen_q,moba_cu_seqlen_kv,moba_q_sh_indices,) ctx.saved_tensors# 确保输入梯度连续d_output d_output.contiguous()# 计算自注意力的梯度dq, dk, dv, _ _flash_attn_varlen_backward(doutd_output,qq,kk,vv,outoutput,softmax_lsemixed_attn_vlse_sh.t().contiguous(),dqNone,dkNone,dvNone,cu_seqlens_qself_attn_cu_seqlen,cu_seqlens_kself_attn_cu_seqlen,max_seqlen_qmax_seqlen,max_seqlen_kmax_seqlen,softmax_scalesoftmax_scale,causalTrue,dropout_p0.0,window_size(-1, -1),softcap0.0,alibi_slopesNone,deterministicTrue,)# 计算 MoBA 注意力的梯度headdim q.shape[-1]d_moba_output (d_output.view(-1, headdim).index_select(0, moba_q_sh_indices).unsqueeze(1))moba_output (output.view(-1, headdim).index_select(0, moba_q_sh_indices).unsqueeze(1))mixed_attn_vlse (mixed_attn_vlse_sh.view(-1).index_select(0, moba_q_sh_indices).view(1, -1))dmq, dmk, dmv, _ _flash_attn_varlen_backward(doutd_moba_output,qmoba_q,kmoba_kv[:, 0],vmoba_kv[:, 1],outmoba_output,softmax_lsemixed_attn_vlse,dqNone,dkNone,dvNone,cu_seqlens_qmoba_cu_seqlen_q,cu_seqlens_kmoba_cu_seqlen_kv,max_seqlen_qmax_seqlen,max_seqlen_kmoba_chunk_size,softmax_scalesoftmax_scale,causalFalse,dropout_p0.0,window_size(-1, -1),softcap0.0,alibi_slopesNone,deterministicTrue,)# 合并 MoBA 的键和值的梯度dmkv torch.stack((dmk, dmv), dim1)return dq, dk, dv, None, dmq, dmkv, None, None, None, None, None
代码关键部分解释 前向传播 (forward) 分别计算自注意力和 MoBA 注意力的结果。对注意力分数进行处理包括形状转换、归一化等操作以避免指数爆炸。将自注意力和 MoBA 注意力的结果进行加权合并得到最终的输出。保存中间结果用于后续的反向传播。 反向传播 (backward) 根据前向传播保存的中间结果计算自注意力和 MoBA 注意力的梯度。最终返回各个输入张量的梯度。
小结 通过这种方式MixedAttention 类实现了 MoBA 混合块注意力机制通过将上下文划分为块并进行选择性的注意力计算有效减少了计算量提升了处理长文本的效率。
总结
总体而言MoBA 非常适合处理涉及大量上下文的任务如长篇文档阅读理解、大规模代码补全以及需要完整对话历史的多轮对话系统。它在提高效率的同时性能损失极小为大规模训练大语言模型提供了一种极具吸引力的方法。虽然目前 MoBA 主要应用于文本领域但研究人员认为其底层机制在其他数据模态中也具有应用潜力。只要序列长度足够长引发计算或内存问题将查询分配给块 “专家” 的思路就有望缓解瓶颈同时保持处理关键全局依赖关系的能力。随着语言应用中的序列长度持续增长像 MoBA 这样的方法可能会在推动神经语言建模的可扩展性和成本效益方面发挥关键作用为人工智能的发展注入新的活力。
