破解DeepSeek的上下文窗口

128K上下文窗口。这个数字在2023年初还属于科幻范畴,到2024年已经快成标配了。但同样是128K,体验差距可以天差地别。有的模型声称支持128K,实测到40K就开始胡言乱语;有的模型128K内表现稳定,但推理速度慢得让你想砸键盘。

DeepSeek-V2的128K属于第三种:既稳定又快。这就值得聊一聊了。

要让模型支持长上下文,首先要解决位置编码的问题。传统的位置编码方案——不管是绝对位置编码还是正弦位置编码——都有一个共同的毛病:训练时见过的最长序列多长,推理时就只能处理多长。你想让它处理更长的文本,位置编码直接就是随机值,模型自然就懵了。

于是有了各种相对位置编码:RoPE、ALiBi、xPos等等。DeepSeek用的是RoPE的改进版本,具体来说是在YaRN方法基础上做的定制。YaRN的思路很简单也很有效:训练的时候用较短的序列,推理的时候把位置频率外推到更长的范围。但外推不是简单线性延伸,需要对频率做某种重新缩放。

这里有一个关键的数学问题:RoPE的高频和低频成分对外推的敏感度不一样。高频成分(对应位置之间的小距离)对外推相对鲁棒,低频成分(对应大距离的位置关系)外推后很容易失真。YaRN的做法是对不同频率做不同的缩放,高频部分直接外推,低频部分做插值,中间部分渐进过渡。

DeepSeek的改进在于,他们把频率划分做得更细致,引入了多个过渡带。这样做的好处是让模型在128K范围内的位置感知更加平滑,不会出现某些中间距离的位置关系被“遗忘”的情况。

但从我的测试来看,DeepSeek的上下文窗口优势不仅仅来自位置编码。更关键的可能是他们在训练数据上下功夫了。

长上下文能力不是凭空产生的,得用长序列训练数据喂出来。什么样的数据天然具有长序列结构?书、论文、长对话、代码仓库。我猜DeepSeek在数据配比中大幅提高了这些长序列来源的比例。特别是代码,一个完整的代码文件动辄几千行,函数之间的调用关系跨越几百甚至几千行,这种天然的远程依赖关系对训练长上下文能力非常有帮助。

还有一个训练技巧:在预训练后期逐步增加序列长度。一开始用短的(比如4K),让模型先学会基本的语言规律;然后逐渐加长到8K、16K、32K,最终到128K。这种渐进式训练比一上来就用128K序列高效得多,因为短序列训练快,可以让模型快速收敛到不错的初始状态,长序列阶段只需要微调长距离依赖能力。

但渐进式训练有个陷阱:学习率调度。不同长度阶段的损失函数分布不一样,如果用同一个学习率,很容易在新长度阶段出现不稳定。DeepSeek大概率是每个阶段都有独立的学习率warmup和衰减策略,甚至可能对不同层用了不同的学习率。

说完了训练,说推理。长上下文推理的最大痛点是“中间丢失”现象。你丢给它一大段文本,让它回答关于中间部分的问题,很多模型会漏掉关键信息或者直接编造。这种现象被称为“lost in the middle”,是当前大模型的通病。

DeepSeek对这个问题的缓解来自于两个方面。一个是前面提到的MLA,因为压缩了KV缓存,模型被迫学习保留关键信息,某种程度上起到了信息瓶颈的正面作用。另一个是他们在训练时做了针对性的数据增强,故意把重要信息放在文本的不同位置,让模型学会从任意位置提取信息。

我做过一个实验:把一份长文档的关键信息分别放在开头、中间和末尾,测试不同模型的信息提取准确率。DeepSeek在中间位置的准确率确实比同期的几个竞品高出一截。虽然还做不到首尾一致,但差距缩小了很多。

还有一个实际问题:长上下文的显存管理。128K序列的KV缓存即使压缩了也还是不小。DeepSeek在这里用了一个叫“分页注意力”的技术,把KV缓存分成固定大小的页,不连续的序列可以共享物理存储。这个技术最初是vLLM提出的,DeepSeek做了适配MLA的改进,让分页管理跟压缩后的KV表示对齐。

工程实现上有意思的一个细节是,他们的分页大小不是固定的,而是根据序列长度动态调整。短序列用大页减少管理开销,长序列用小页提高内存利用率。这个自适应设计在小批量推理时提升很明显。

说到底,长上下文能力的核心挑战不是“能处理多长”,而是“在任意长度下都能保持稳定的理解和生成质量”。很多模型为了支持长上下文牺牲了短文本的表现,或者128K只是理论数字实际根本用不起来。DeepSeek在这两点上做得都比较扎实,128K不是营销数字,是实际可用、稳定可靠的窗口大小。

当然,128K远不是终点。已经有团队在研究百万token级别的上下文窗口了。那将是完全不同量级的挑战,位置编码、注意力计算、显存管理都需要重新思考。不知道DeepSeek在这一点上有没有什么新进展,期待他们的下一步。