这篇论文把Transformer的预测和记忆分开了,实验证明效果更好,下游任务平均提了2-3个点,搞LLM架构的值得细看。
该论文提出状态-预测分离假设,认为Transformer中预测下一个token与存储有用状态可被分离。作者设计双流Transformer变体,将两个功能分配到不同计算流。在多个规模下的预训练实验中,该方法在数据和计算效率上持续优于标准Transformer,验证损失更低,下游任务平均提升2-3个百分点。额外实证分析排除了潜在混淆因素,揭示了设计带来的梯度差异。
The State-Prediction Separation Hypothesis
Transformers use the same forward computation stream to both predict the next token and store useful state for future token predictions. We formulate the \emph{state-prediction separation hypothesis}: disentangling the two roles yields better language modeling performance. We design a Transformer variant that uses two computation streams to separate the two functions, and conduct pretraining experiments across various scales. Our experiments show that state-prediction separation consistently offers better data and compute efficiencies, improving validation loss and outperforming standard Transformers by 2--3 percentage points on average on downstream tasks. We also conduct extensive empirical analysis that rules out potential confounders and demonstrates the fundamental difference in the gradients our design entails.