LightTransfer：将你的LLM轻松转为Hybrid model，增强o1-like长文本生成能力

最近hybrid model的有效性已被广泛验证（比如minimax-01和gemma2等）。该篇论文研究了一个非常有趣的主题：如何将预训练好的dense transformer（如qwq）转化为hybrid model。具体来说，提出了一个轻量转化算法——LightTransfer💡，将50%层的full attention替换为sliding window attention。

作者：Wzl
来自：深度学习自然语言处理
链接：https://arxiv.org/abs/2410.13846
主页: https://sites.google.com/view/lighttransfer

动机和观察

（1）从什么角度实现？

• 一个高效且自然的建模方式：大规模预训练的LLM+少部分的修改+少部分的训练 -> Hybrid model。
• 一种直接且有效的做法：将full attention layers中部分standard attn layer替换为efficient attn layer。

（2）efficient attn layer实现？

引入KV cache compression的代表做法：streaming attn+sink，实现sliding window attn（即保留每一层开头和结尾对应的KV对）。

（3）部分layer替换？

• 某些layer的attn是dense的，替换为sliding window attn对attn的破坏比较大。从图中可以看出，一些层的注意力并不具有高稀疏的特点，所以不适合用sliding window attn的方式。

方法：LightTransfer

（1）两种模式

• LightTransfer-test ：通过fig.3的框架，识别出lazy layers，直接在推理阶段替换这些lazy layer为efficient attn layer。
• LightTransfer-train ：为了适应o1类模型的复杂长文本生成场景，lazy layer被替换之后，进行SFT。

（2）test/train的设计符合观察

原文中的基于Figure2的第二点观察为：“Layer behavior remains consistent for a given input”（可以理解为pd过程lazy layer不变）。说明了在prefill阶段选择好的lazy layer在decoding阶段仍然适用，因此不需要在decoding阶段更换lazy layer，并且这也更有利于保证train模式下替换lazy layer之后，SFT的稳定性和有效性。

实验

• test模式下 ，在LongBench测试中，LightTransfer-test的指标相比baseline取得了不错的效果🆙
• train模式下 ，在三个常用的数学benchmark测试中，LightTransfer-train的结果在AIME24上提升显著🆙
• 两种模式下的推理速度提升 ，在16K的长文本推理场景，相比full-attention加速达到了2.17x🚀
• 消融实验 ，相比之前的Pyramid等方式，LightTransfer采用的Lazy layer的Hybrid方式更合理🧐