杨植麟带队发论文,Kimi下一代模型架构曝光,推理之父惊呼:深度学习2.0要来了!
由知名华人AI学者何恺明等人提出的标准残差连接,最早在ResNet中系统化提出,并迅速成为现代深度神经网络的核心结构之一。
随着深度学习的发展,这一设计不仅在计算机视觉中取得巨大成功,也逐渐成为大型语言模型和各类Transformer系统中的基础构建模块。
其核心思想非常直观:让网络学习“残差”(即输入与输出之间的差值),而不是直接拟合完整映射。标准残差连接中,每一层的输入等于所有之前层输出的累积和。残差连接实际上是一种跨深度的信息聚合机制,所有层的贡献权重都是1,也就是均匀加权。
▲从左到右:传统残差、全注意力残差、块注意力残差
在训练阶段,Block AttnRes带来更多的流水线并行中通信开销。标准残差结构在流水线阶段之间只需传输固定大小的隐藏状态,而Block AttnRes需要在每个阶段访问此前所有块的表征。
如果采用朴素实现,就必须在阶段转换时重复传输完整的历史块表示,通信量会随着块数量累积呈平方增长。为了解决这一问题,月之暗面引入跨阶段缓存(cross-stage cache)机制,大幅减少冗余通信。
同时,每个块只需在虚拟阶段中存储一次,结合激活检查点机制后,每层激活内存占用与标准Transformer基本一致,因此整体训练的额外开销很小,在实际系统中端到端训练时间增加不到4%。
在推理阶段,Block AttnRes的主要挑战来自跨层块表征的重复访问和长上下文缓存带来的内存压力。为此他们采用了两阶段计算策略,首先在阶段一中对缓存的块表征执行一次批量查询,为同一块内的所有层同时计算块间注意力,从而将原本每层都需要访问历史块的操作合并为一次矩阵计算,显著降低内存读取次数。
随后在阶段二中按顺序计算块内各层的注意力,并通过在线softmax将结果与阶段一的输出合并,这一阶段可以与周围算子进行内核融合,从而进一步减少I/O开销。
此外,为缓解长上下文预填充时块表征缓存过大的问题,系统将这些表征沿序列维度在张量并行设备之间进行分片,使每个设备只存储部分序列。通过两阶段计算和序列分片的结合,BlockAttnRes在推理时仅带来不到2%的延迟开销,同时显著降低了长上下文场景下的显存需求。
在实验方面,研究团队将这一结构集成到Kimi Linear架构中进行验证。
实验表明,在相同参数规模下,引入Block AttnRes后,模型在多个基准测试中都取得了性能提升。其中提升最多的是多步推理基准测试GPQA-Diamond,采用AttnRes的模型得分提升幅度超过了20%。