The Context-Ready Transformer

# 上下文就绪变换器（Context-Ready Transformer） 来源：https://arxiv.org/html/2606.27538 ###### 摘要 我们提出*上下文就绪变换器*，这是一种新型循环神经网络架构，由DD层变换器块构建而成，该块在每个token进入之前对其进行预上下文化。在从左到右的生成过程中，一个校正网络将前一个位置的块输出（即对过去上下文的缓存摘要）与当前token嵌入相结合，使得token在进入块时已是上下文化的，而非原始嵌入。在序列推理时，校正链使该架构成为循环神经网络。训练时，我们将校正过程在整个序列上展开KK次，每一步并行处理所有位置。预训练变换器也可以通过添加零初始化的校正FFN并进行微调转换为上下文就绪模型。我们在多种宽度、深度、块大小及两个数据集上进行评估，所有比较均针对标准变换器、变体及消融实验。一个D=5D{=}5模型击败了12层变换器，同时在A100上生成速度提升了1.7×1.7\times。在K=10K{=}10的情况下，单层模型（D=1D{=}1）击败了6层变换器，推理速度提升2.6×2.6\times，且序列推理与并行K=10K{=}10的困惑度差异在0.01以内。该架构在宽表示和长上下文中获益最多。在指针追踪任务中，使用BPTT训练的D=1D{=}1模型解决了所有10个组合层级，而标准变换器则表现出阶梯状的深度依赖性。 ## 1 引言 在自回归生成中，一个标准NN层变换器预测下一个token，为其分配一个上下文无关的基础嵌入，然后将其全部或部分NN层用于重新进行上下文化。这种从上下文到token ID再回到上下文的往返是架构的产物，而非问题本身。 上下文就绪推理。*上下文就绪变换器*缩短了这一往返过程。考虑从左到右的序列生成：当模型处理token tt时，块输出ztz_{t}编码了token 0,...,t0,\ldots,t的上下文。当token t+1t{+}1带着嵌入et+1e_{t+1}到来时，一个校正网络根据ztz_{t}和et+1e_{t+1}计算校正。token作为et+1e_{t+1}加上该校正进入块（携带来自所有之前token的上下文信息），而非原始嵌入。由于校正已完成了部分工作，因此需要更少的层数来完全上下文化该token。 训练挑战。这个序列过程在推理时是精确的，但本质上是串行的：每个位置的校正依赖于前一个位置已完全计算的块输出。经典RNN面临同样问题，并通过时间反向传播（BPTT）解决，该算法对整个序列的所有TT位置展开循环——使得串行训练的深度与序列长度成正比。我们采用不同方法。我们通过在整个序列上*展开*KK次来近似该串行过程，其中KK是与TT无关的小常数。所有TT位置仍然并行处理——就像标准变换器一样——但校正在KK次展开步骤（而非TT次）中得到精炼。从所有位置的原始嵌入开始，我们运行一个共享权重的块，并根据位置0,...,t−10,\ldots,t{-}1的块输出计算每个位置tt的校正。然后我们用这些校正更新所有嵌入并重复。每一步展开都使用越来越上下文化的输入来精炼校正。关键在于，KK控制计算图的深度——而非TT——因此经典RNN需要O(T)O(T)深度的展开，而上下文就绪变换器仅需O(K)O(K)。在实践中，对于测试的深度（D≥5D\geq 5；见表6 (https://arxiv.org/html/2606.27538#S5.T6)），K=5K=5足以收敛。 为何训练时的KK步展开能在推理时实现零迭代。两个设计选择使之成为可能。 - •**非累积校正**。每次迭代从头开始计算校正，而非基于前一次。迭代形式为x(k)=e+Fθ(x(k−1),e)x^{(k)}=e+F_{\theta}(x^{(k-1)},e)：基础嵌入加上一个依赖于前一次迭代块输出和token嵌入的校正。与残差网络（计算x(K)=e+f(x(0))+⋯+f(x(K−1))x^{(K)}=e+f(x^{(0)})+\cdots+f(x^{(K-1)})，即KK个校正之和）不同，我们的公式给出x(K)=e+Fθ(x(K−1),e)x^{(K)}=e+F_{\theta}(x^{(K-1)},e)：单个校正。之前的迭代仅用于使x(K−1)x^{(K-1)}接近不动点；一旦收敛，更多迭代产生相同输出。 - •**仅过去上下文化**。位置tt的校正依赖于两个量：块输出zt−1z_{t-1}（编码了token 0,...,t−10,\ldots,t{-}1的上下文）和token嵌入ete_{t}。由于zt−1z_{t-1}已在处理前一个token时被缓存，一旦token tt到达即可计算校正。 这两个属性共同意味着序列生成会自然地为每个新token产生收敛的校正，无需迭代（第4节 (https://arxiv.org/html/2606.27538#S4)）。 一种新型RNN。在序列推理时，位置tt的校正依赖于zt−1z_{t-1}，而zt−1z_{t-1}又依赖于zt−2z_{t-2}，依此类推——形成一个在所有TT位置上展开的循环计算。非累积 + 仅过去的结构使得高效训练成为可能：它将序列循环转换为一个不动点问题，因此我们可以通过在整个序列上并行展开KK步来训练（K≪TK\ll T），而不是精确的TT步BPTT。这不等同于BPTT——超过前KK个位置的位置会接收到近似梯度而非精确梯度——但在实践中，对于测试的深度，K=5K=5已足够。 实验证据。我们在宽度C∈{50,...,2048}C\in\{50,\ldots,2048\}、深度D∈{1,...,23}D\in\{1,\ldots,23\}、块大小{64,256,512,1024}\{64,256,512,1024\}、两个数据集（OpenWebText, Wikipedia）以及一个合成推理任务（第5节 (https://arxiv.org/html/2606.27538#S5)）上进行评估。最具影响力的发现：在C=1120C=1120时，D=5D{=}5击败了12层变换器（C=768C=768），将推理深度减半，并在A100上生成速度提升1.7×1.7\times。在K=10K{=}10时，C=2048C=2048的D=1D{=}1模型击败了6层变换器（C=1088C=1088），推理速度提升2.6×2.6\times。序列推理与K=10K{=}10训练的困惑度差异在0.01以内，证实流式精确性在实践中成立。在指针追踪任务中，使用BPTT训练的D=1D{=}1模型解决了所有10个组合层级，而标准变换器表现出阶梯状的深度依赖性。该架构在宽表示和长上下文中获益最多，并且需要一个专用的校正网络才能有效（第5.8节 (https://arxiv.org/html/2606.27538#S5.SS8)）。任何预训练变换器都可以通过添加一个零初始化的校正FFN并进行微调来转换。 ## 2 相关工作 权重共享与深度自适应架构。ALBERT (Lan et al., 2020 (https://arxiv.org/html/2606.27538#bib.bib14))、Universal Transformers (Dehghani et al., 2019 (https://arxiv.org/html/2606.27538#bib.bib6)) 结合 ACT (Graves, 2016 (https://arxiv.org/html/2606.27538#bib.bib11))、Deep Equilibrium Models (Bai et al., 2019 (https://arxiv.org/html/2606.27538#bib.bib1)) 以及 Huginn (Geiping et al., 2025 (https://arxiv.org/html/2606.27538#bib.bib9)) 在层或迭代之间共享权重并迭代精炼隐藏状态，但并未使用本文提出的那种专用的过去输出/当前token感知的块前校正。我们以标准变换器作为基线，代表token如何进入块。在上下文就绪变换器中，权重共享自然来自于将序列校正过程展开为并行处理所有位置的训练迭代，而非出于参数效率的设计选择。 提前退出与层跳过。LayerSkip (Elhoushi et al., 2024 (https://arxiv.org/html/2606.27538#bib.bib7))、ADEPT (Yoo et al., 2026 (https://arxiv.org/html/2606.27538#bib.bib21)) 和 PonderNet (Banino et al., 2021 (https://arxiv.org/html/2606.27538#bib.bib3)) 减少了平均层数，但需要学习停止机制。上下文就绪使用固定深度，无停止准则。 Lookahead Decoding (Fu et al., 2024 (https://arxiv.org/html/2606.27538#bib.bib8)) 和 CLLMs (Kou et al., 2024 (https://arxiv.org/html/2606.27538#bib.bib13)) 将Jacobi迭代应用于标准变换器作为*解码策略*。Bai等人 (2021 (https://arxiv.org/html/2606.27538#bib.bib2)) 通过Jacobi式更新并行化不动点求解来加速DEQ推理。上下文就绪是一种*架构变更*：KK步展开可视为对校正不动点方程进行Jacobi迭代，但非累积的仅过去结构保证了单次从左到右的流式传递即可恢复精确校正，无需任何迭代。 次二次与循环替代方案。Mamba (Gu and Dao, 2024 (https://arxiv.org/html/2606.27538#bib.bib12))、RWKV (Peng et al., 2023 (https://arxiv.org/html/2606.27538#bib.bib17))、Griffin (De et al., 2024 (https://arxiv.org/html/2606.27538#bib.bib5)) 和 xLSTM (Beck et al., 2024 (https://arxiv.org/html/2606.27538#bib.bib4)) 用压缩的循环状态替换因果自注意力，以实现线性时间推理。上下文就绪变换器解决的是不同问题：它在块内部*保留*完整的因果自注意力，而是改变token进入块的方式。这两种方法是互补而非竞争的——原则上可以将块前校正应用于这些架构中的任何一种——标准变换器仍然是评估校正机制的自然基线。 变换器的计算复杂度。对数精度固定深度变换器限制在TC0\mathrm{TC}^{0} (Merrill and Sabharwal, 2023 (https://arxiv.org/html/2606.27538#bib.bib16))：无论宽度如何，它们都无法解决需要无界序列组合的问题。具有任意精度的RNN可以逃脱此限制 (Siegelmann and Sontag, 1995a (https://arxiv.org/html/2606.27538#bib.bib18); Siegelmann and Sontag, 1995b (https://arxiv.org/html/2606.27538#bib.bib19))，但经典RNN需要梯度流经所有TT个时间步（BPTT），这使得它们难以在长序列上训练。上下文就绪变换器在D=1D=1时在推理时是循环的，但训练时通过展开校正KK次而非精确的TT步BPTT——在推理时继承了循环结构，同时保留了变换器风格的并行训练。 ## 3 方法 ### 3.1 架构 我们首先描述序列从左到右生成时的架构——这是上下文就绪推理精确的场景。令TT表示序列长度，CC表示嵌入维度，VV表示词汇表大小。核心组件是*DD块单元*：DD个变换器层（每层由因果自注意力和带残差连接的前馈网络组成），下文详述。处理token tt需要前序所有token的块输出z0,...,zt−1∈RCz_{0},\ldots,z_{t-1}\in\mathbb{R}^{C}。令et∈RCe_{t}\in\mathbb{R}^{C}表示位置tt的token嵌入，并定义z−1=0z_{-1}=\mathbf{0}。 校正。一个专用的*校正FFN*根据缓存的块输出zt−1z_{t-1}和当前token嵌入ete_{t}生成位置tt的校正： correctiont=corr_ffn(LN(zt−1+et))\texttt{correction}_{t}=\texttt{corr\_ffn}\bigl(\texttt{LN}(z_{t-1}+e_{t})\bigr) (1) 校正FFN是一个前馈网络（Linear(C→4CC\to 4C) →\to GELU →\to Linear(4C→C4C\to C)），拥有独立于块FFN的自身权重。校正是*token感知*的：它依赖于zt−1z_{t-1}（token 0,...,t−10,\ldots,t{-}1的上下文）和ete_{t}（当前token嵌入）。由于两个输入在token tt到达时均可用，校正是因果的——不依赖于任何未来token。 上下文化。新token带着添加到其原始嵌入的校正进入块： x~t=et+correctiont\tilde{x}_{t}=e_{t}+\texttt{correction}_{t} (2) 块处理。DD块单元应用DD个变换器块，具有*独立*权重和标准残差连接： h(0)\displaystyle h^{(0)}=x~t\displaystyle=\tilde{x}_{t}a(i)\displaystyle a^{(i)}=h(i−1)+Attni(LNia(h(i−1));Ki,Vi),i=1,...,D\displaystyle=h^{(i-1)}+\texttt{Attn}_{i}\bigl(\texttt{LN}^{a}_{i}(h^{(i-1)});\;\mathcal{K}_{i},\mathcal{V}_{i}\bigr),\quad i=1,\ldots,Dh(i)\displaystyle h^{(i)}=a(i)+FFNi(LNif(a(i)))\displaystyle=a^{(i)}+\texttt{FFN}_{i}\bigl(\texttt{LN}^{f}_{i}(a^{(i)})\bigr)zt\displaystyle z_{t}=h(D)\displaystyle=h^{(D)} (3) 每个Attni\texttt{Attn}_{i}是带有旋转位置嵌入（RoPE）(Su et al., 2024 (https://arxiv.org/html/2606.27538#bib.bib20))的因果自注意力。每个FFNi\texttt{FFN}_{i}：Linear(C→4CC\to 4C) →\to GELU →\to Linear(4C→C4C\to C)。参数DD控制推理深度。 预测。 logitst=Whead⋅LNf(zt),Whead∈RV×C\texttt{logits}_{t}=W_{\text{head}}\cdot\texttt{LN}_{f}(z_{t}),\quad W_{\text{head}}\in\mathbb{R}^{V\times C} (4) 预测后，ztz_{t}被缓存，KV缓存Ki,Vi\mathcal{K}_{i},\mathcal{V}_{i}为未来token更新。 ### 3.2 并行训练 序列推理是精确的，但本质上是串行的。对于训练，我们将校正过程展开KK次，每一步并行处理所有TT个位置。梯度流经KK个展开步骤而非TT个时间步，使得该架构尽管在推理时是循环的，但可以像变换器一样进行训练。 给定token嵌入e=(e1,...,eT)e=(e_{1},\ldots,e_{T})，对于所有kk，z0(k)=0z_{0}^{(k)}=\mathbf{0}（来自第3.1节 (https://arxiv.org/html/2606.27538#S3.SS1)的初始缓存），初始化correction(0)=0\texttt{correction}^{(0)}=\mathbf{0}。对于k=1,...,Kk=1,\ldots,K： x~t(k−1)\displaystyle\tilde{x}_{t}^{(k-1)}=et+correctiont(k−1)\displaystyle=e_{t}+\texttt{correction}_{t}^{(k-1)} （上下文化）h(0)\displaystyle h^{(0)}=x~t(k−1)\displaystyle=\tilde{x}_{t}^{(k-1)}a(i)\displaystyle a^{(i)}=h(i−1)+Attni(LNia(h(i−1));Ki,Vi)\displaystyle=h^{(i-1)}+\texttt{Attn}_{i}\bigl(\texttt{LN}^{a}_{i}(h^{(i-1)});\;\mathcal{K}_{i},\mathcal{V}_{i}\bigr)i=1,...,D\displaystyle i=1,\ldots,Dh(i)\displaystyle h^{(i)}=a(i)+FFNi(LNif(a(i)))\displaystyle=a^{(i)}+\texttt{FFN}_{i}\bigl(\texttt{LN}^{f}_{i}(a^{(i)})\bigr)zt(k)\displaystyle z_{t}^{(k)}=h(D)\displaystyle=h^{(D)} （块输出）correctiont(k)\displaystyle\texttt{correction}_{t}^{(k)}=corr_ffn(LN(zt−1(k)+et))\displaystyle=\texttt{corr\_ffn}\bigl(\texttt{LN}(z_{t-1}^{(k)}+e_{t})\bigr)t=1,...,T\displaystyle t=1,\ldots,T (5) 在迭代kk之间，DD个变换器块共享权重，但在层ii之间具有独立权重。 非累积校正。每次迭代完全替换前一次校正：x~(k)=e+correction(k)\tilde{x}^{(k)}=e+\texttt{correction}^{(k)}，而非x~(k)=x~(k−1)+f(x~(k−1))\tilde{x}^{(k)}=\tilde{x}^{(k-1)}+f(\tilde{x}^{(k-1)})。