通用YOCO：高效的深度缩放架构 Universal YOCO for Efficient Depth Scaling
YOCO-U结合递归计算与高效注意力，在保持常数级KV缓存的同时实现深度缩放
前置知识
KV缓存
Key-Value缓存是Transformer解码器在自回归生成时存储的历史token的键值对，用于加速后续token的注意力计算。标准Transformer中每个decoder层都有自己的KV缓存，总内存消耗为O(NL)，其中N是序列长度，L是层数。YOCO架构通过共享缓存将复杂度降低到O(N)。
本文的核心贡献之一是在引入递归计算的同时保持KV缓存为常数级，理解KV缓存的增长机制是理解YOCO-U优势的关键。
测试时缩放
测试时缩放指在推理阶段增加计算量来提升模型性能的技术，包括思维链推理、多轮采样等方法。这些方法通过分配更多FLOPs给推理过程，显著提升了LLM在复杂推理任务上的表现。然而，这种能力依赖于预训练阶段建立的内在能力，而非直接改善预训练本身。
本文针对的问题正是如何在架构层面更高效地支持测试时缩放，理解测试时缩放的动机和局限性有助于评估YOCO-U的贡献。
高效注意力
高效注意力是指具有亚二次复杂度的注意力变体，如滑动窗口注意力、线性注意力、RetNet、Mamba等。这些方法通过限制注意力范围或使用低秩近似，将时间复杂度从O(N^2)降低到O(N)或更低。SWA只关注固定窗口内的token，窗口外token通过递归传递信息。
YOCO-U将递归计算限制在使用高效注意力的浅层模块中，这是其能够在不显著增加计算和内存开销的情况下实现深度缩放的技术基础。
参数共享递归
参数共享递归是Universal Transformer的核心思想，通过在深度维度共享参数并多次迭代同一模块来增加模型的表达能力。这种方法在参数量不变的情况下增加了计算深度，理论上可以通过更多迭代来渐进改进表示。然而，在标准Transformer上应用递归会导致KV缓存线性增长且需要重复全局注意力计算。
YOCO-U采用了参数共享递归的思路，但将其限制在特定的架构位置，理解递归计算的优势和挑战是评估YOCO-U设计选择的基础。
研究动机
随着测试时缩放的兴起，大规模语言模型在推理和智能体能力方面取得了显著提升。然而，标准Transformer架构难以高效支持推理时的计算缩放。现有方法面临两个关键限制：首先，仅依赖后训练策略相对低效，无法充分利用预训练阶段的基础知识和表示深度。其次，虽然在标准Transformer中实现循环机制理论上可以扩展计算深度，但会导致过高的计算成本，Key-Value缓存的内存占用随模型深度线性增长。例如，标准Transformer的KV缓存复杂度为O(LND)，其中L是层数、N是序列长度、D是隐藏维度，这在长序列场景下成为严重的内存瓶颈。此外，标准的Universal Transformer将递归应用于整个网络，导致冗余开销和优化困难。
本文的目标是本文的目标是提出一种能够高效支持深度缩放的Transformer架构，在不增加KV缓存足迹的前提下通过参数共享和递归计算提升模型能力。具体而言，本文旨在设计一个架构，将递归计算的优势与YOCO架构的推理效率相结合，实现比两者单独使用更好的协同效应。目标包括保持YOCO原有的线性预填充复杂度O(N)和单次KV缓存优势，同时通过递归计算增强表示深度，最终在相同的FLOPs预算下获得更好的性能或在达到相同性能时需要更少的训练数据。
与已有工作不同的是，本文的独特切入角度在于将递归计算与YOCO的decoder-decoder架构相结合，但仅将递归限制在浅层的自解码器模块中。这与现有工作形成鲜明对比：Universal Transformer将递归应用于整个网络，导致高开销；RINS在标准Transformer的早期层应用递归，但仍然面临KV缓存线性增长的问题；而并行缩放方法如ParScale不增加建模深度，性能通常不如递归缩放。本文的核心洞察是递归计算与高效注意力模块的结合可以相互增强，YOCO-U通过限制递归到使用滑动窗口注意力的浅层模块，避免了全局注意力的重复计算，同时通过参数共享在固定FLOP预算内显著提升表达能力。
核心方法
YOCO-U基于YOCO（You Only Cache Once）框架，该框架将模型分为自解码器和交叉解码器两部分。自解码器使用高效注意力（如滑动窗口注意力）处理输入token，生成紧凑的全局KV缓存；交叉解码器通过交叉注意力在所有层重用这个共享缓存，避免了标准Transformer每层KV缓存的问题。YOCO-U的创新在于将静态的自解码器替换为通用自解码器，该模块通过共享参数进行T步递归计算，在不增加参数数量的情况下通过递归深度增强表示能力。递归仅限于浅层模块，而交叉解码器及其全局注意力保持不变。这种设计的关键优势是全局KV缓存只需生成一次并在迭代间重用，不引入额外的全局注意力计算。整体架构形成了Think Deeper, Cache Once的设计理念：通过在高效注意力模块中进行递归来think deeper，同时通过共享缓存机制实现cache once。
YOCO-U的核心创新是将递归计算限制在使用高效注意力的浅层模块中，而不是像Universal Transformer那样对整个网络进行递归。这个设计的本质区别在于：YOCO-U中的递归步骤只产生高效注意力的计算开销，而不是完整Transformer层的全部成本。全局KV缓存由通用自解码器生成一次，然后在所有交叉解码器层中被重用。递归只影响局部窗口的KV缓存，由于窗口大小W远小于序列长度N，对于长序列来说这个额外开销可以忽略不计。表1显示，YOCO-U的KV缓存复杂度为O((N+WTL)D)，其中全局部分O(ND)与迭代次数T无关，只有局部高效注意力缓存O(WTLD)随迭代增长。由于W通常为512而N可达32K-256K，这个开销在长序列场景下确实可以忽略。与标准Transformer循环的O(LTND)和RINS的O(LTND)相比，YOCO-U在保持递归优势的同时显著降低了内存占用。
方法步骤详情
YOCO-U的方法分为以下几个完整步骤。首先，输入序列被打包成嵌入矩阵。第二步，自解码器由L/2层模块组成，每层使用高效自注意力和SwiGLU前馈网络。第三步，通用自解码器对自解码器模块执行T次递归计算。第四步，USD的输出生成交叉解码器的全局KV缓存。第五步，交叉解码器同样由L/2层组成，但使用交叉注意力重用共享缓存，其中每层的查询是逐层计算的。最后，通过softmax分类器进行下一token预测。整个过程中，递归只影响前三个步骤的自解码器部分，而全局KV缓存和交叉解码器保持不变。
技术新颖性
YOCO-U的技术新颖性体现在三个关键方面。首先，它首次将递归计算与decoder-decoder架构相结合，证明了两种技术的协同效应优于各自单独使用。与Universal Transformer对整个网络递归导致冗余开销不同，YOCO-U将递归限制在浅层的自解码器模块，这种部分递归设计在固定FLOP预算内显著提升了表达能力。其次，YOCO-U证明了递归计算与高效注意力模块的结合可以相互增强，实验显示YOCO-U在性能上与RINS相当，但保持了显著更低的KV缓存和推理成本。第三，从架构演进的角度，YOCO-U展示了如何在不牺牲推理效率的前提下实现深度缩放，它保留了YOCO的线性预填充和单次KV缓存优势，同时通过参数共享的递归计算改善了token利用率和缩放行为。
实验结果
论文通过多个实验验证了YOCO-U的有效性。在语言建模缩放实验中（图2），当对齐训练FLOPs时，YOCO-U在相同的计算预算下实现了更低的损失，证明了递归计算带来的优越计算效率。更重要的是，当对齐训练token数量时，YOCO-U显示出更好的token利用率，使用约62%更少的token即可达到与非递归YOCO训练210B token相当的性能。在端任务评估中（表2），使用LM Eval-Harness在300B token上训练的模型显示，YOCO-U在等FLOPs设置下平均得分为46.23，显著超越YOCO基线的41.78（提升4.45个百分点），即使在等步数设置下也达到47.08。这证明收益不仅仅来自额外计算，而是架构设计本身。在数学推理任务上（图3），经过思维链SFT后，YOCO-U在11个数学基准测试中全部超越YOCO，平均准确率提升24.4%，显示显式推理和隐式推理的改进是正交的。在架构对比实验中（表3），YOCO-U与RINS性能相当（均为48.3平均分），但保持了显著更低的KV缓存和推理成本。长序列实验（图4）显示YOCO-U在Book和Code数据上随着输入长度增加持续获得更低的困惑度，证明递归计算与高效注意力的结合没有损害长上下文能力。Needle In-A-Haystack检索测试（表4）进一步确认YOCO-U保持了鲁棒的检索性能。推理效率实验（图7）表明YOCO-U在256K上下文长度下，预填充吞吐量比标准Transformer高10.2倍，比RINS高20倍；解码吞吐量比标准Transformer高2.21倍，仅比非递归YOCO损失5%；KV缓存内存消耗与YOCO几乎重叠，仅比RINS低18.9-17.9倍。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| ARC-C | 准确率(%) | 47.87 | 46.50 | +1.37个百分点 |
| Winogrande | 准确率(%) | 68.67 | 61.72 | +6.95个百分点 |
| HellaSwag | 准确率(%) | 66.80 | 63.44 | +3.36个百分点 |
| MMLU | 准确率(%) | 54.63 | 49.59 | +5.04个百分点 |
| BBH | 准确率(%) | 35.49 | 33.13 | +2.36个百分点 |
| GSM8K | 准确率(%) | 50.49 | 38.06 | +12.43个百分点 |
| Humaneval | 准确率(%) | 10.98 | 9.15 | +1.83个百分点 |
| DROP | 准确率(%) | 34.94 | 32.62 | +2.32个百分点 |
| 数学基准平均 | 准确率(%) | 约60% | 约48% | +24.4%相对提升 |
| Needle In-A-Haystack | 检索准确率 | 1.00/0.95 | 0.87/0.82(Transformer) | +15%/+13个百分点 |
局限与改进
作者在消融研究中提到了几个局限性。首先，循环位置实验显示在更深层进行递归计算会带来收益递减，这表明递归在浅层更有效，但也意味着在更深层的架构中应用递归可能效果有限。其次，表示分析（图8）显示随着循环次数增加，连续层之间的角距离逐渐减小，表明表示趋于固定点，这意味着存在边际收益递减的问题，过多的迭代可能不会带来显著收益。从我的观察来看，论文没有详细讨论YOCO-U在非常短序列（如窗口大小W以内）场景下的表现，在这种情况下高效注意力的优势可能不明显，而递归开销可能相对更显著。此外，论文主要关注语言建模和推理任务，对于生成质量、创造性任务等其他方面的表现尚未充分评估。最后，虽然YOCO-U在KV缓存方面有优势，但递归计算仍然增加了总FLOPs，在极度延迟敏感的实时应用中可能仍需权衡。
独立分析的弱点
从独立分析的角度，YOCO-U存在几个可以改进的弱点。首先，论文中默认自解码器循环3次，导致总FLOPs是非递归基线的2倍，这意味着在相同FLOPs预算下，YOCO-U实际上需要与更大参数量的模型进行比较，而非严格的相同参数量比较。改进方向可以包括自适应循环策略，根据输入复杂度动态调整循环次数，或者探索更高效的循环机制如早期退出。其次，YOCO-U在长序列下的优势依赖于W远小于N的假设，但在极短序列或单token预测场景下，局部窗口KV缓存的开销可能相对更显著。改进方向可以是设计混合策略，在短序列时切换到非递归模式。第三，论文中的表示分析显示表示在迭代中趋于固定点，这意味着过多迭代可能带来收益递减，改进方向可以是引入迭代间残差连接、归一化或正则化技术来保持表示的多样性和迭代效率。最后，YOCO-U目前主要与MoE模型结合使用，在密集模型上的表现和优势尚未充分探索，改进方向可以包括在密集模型上系统评估YOCO-U的性能收益。
未来方向
作者提出的未来研究方向包括探索YOCO-U在更大规模模型和更长训练周期下的表现，以及将递归计算与推理时缩放技术更深度地结合。基于本文成果可以延伸的方向包括：第一，探索不同的递归模式，如分层递归（不同层使用不同迭代次数）或跨块递归（跨多个decoder-decoder块进行递归），以进一步优化计算效率。第二，研究YOCO-U与其他高效注意力变体（如RetNet、Mamba、gated DeltaNet）的结合效果，论文中提到这些变体在混合架构中表现类似，但更系统的研究可能发现更优组合。第三，将YOCO-U的思想扩展到多模态场景，如在视觉-语言模型中应用递归计算到视觉编码器部分，探索跨模态递归的可能性。第四，开发针对YOCO-U架构的专用优化技术，如针对递归模块的编译器优化、量化策略或硬件加速器设计，以充分发挥其架构优势。第五，研究YOCO-U在持续学习、模型压缩和知识蒸馏等场景中的应用，利用其递归特性来改进这些任务的效率。
复现评估
论文提供了较为详细的复现信息。模型配置方面，隐藏维度为2560，层数为20，使用细粒度MoE（64个专家中激活8个加1个共享专家），专家维度为1024，总参数量为10B，激活参数为1.3B。训练使用AdamW优化器，beta值为(0.9, 0.95)，最大学习率为0.001，训练75k步（约300B token）。训练在AMD MI300X GPU上进行。代码和模型权重尚未开源，这增加了完整复现的难度。然而，论文提供了详细的超参数设置和训练配方，包括训练长度8192、批大小4M token、位置编码（Self-Decoder用RoPE，Cross-Decoder用NoPE）等。论文还提到训练过程高度稳定，在整个训练范围内损失轨迹平滑且没有显著峰值。从硬件需求来看，训练10B MoE模型300B token需要大量计算资源，但较小的模型配置（如论文中评估的1.3B密集模型）对于学术机构来说应该是可行的。总体而言，复现难度中等偏高，主要障碍在于大规模计算资源和代码可用性。
论文图表
