小思框架研究概览:跨层因果注意力(Cross-Layer Causal Attention)
关于自研序列架构 Beta12Transformer 的技术文章。参考实现tnl_torch/torch_models.py中的Beta12Transformer/Beta12Layer测试配置t19_flash、beta_1.2_test及其消融臂族。1. 一句话定位beta_1.2 在标准 decoder-only Transformer 的骨架完全不变残差流 pre-norm SwiGLU 因果掩码 next-token 预测的前提下唯一改动一处每层注意力的 KV 源不再只取前文的当前层表示而是取前文 token 已完工的更深层的表示。这个改动把模型深度从串行堆叠的层数变成一个可以沿上下文长度免费生长的维度。2. 动机标准 Transformer 的深度天花板标准 Transformer 有两个深度相关的天花板串行深度恒等于层数 L。每跳回一个 token 必须降一层层 l 的注意力读层 l 的 KV信息在 token 维度的传播与深度轴是绑死的。要传播 T 步串行依赖就需要 T 层——这在长上下文推理状态跟踪、多步组合、代码数据流中是不现实的。循环加深要花真金白银。层间循环/迭代方案Universal Transformer、looped Transformer用同一组权重重复应用 R 次来扩深度但每 token 计算量变成 R×L——深度是用 FLOPs 买的且权重循环带来训练稳定性问题。beta_1.2 的出发点token 反正要逐个生成让同层权重沿 token 链复用RNN 式的水平链是免费的串行深度。问题只在于如何让一个 token 在一次前向里读到前文已经算得更深的信息答案是直接把 KV 的来源扩展到深度轴。3. 核心机制3.1 施工态与成品施工态token 当前正在第 l 层被加工时的输入残差流h^(l-1)还没算完。成品token 已完成第 m 层计算后的表示h^(m)对 m 层来说已完工。标准 Transformer 中层 l 的注意力里query 与 KV 都取自施工态本层输入。beta_1.2 中query 来自本 token 的施工态h_i^(l-1)KV 来自前文 token 的成品h_j^(m)且允许 m ≥ l同层或更深。两侧来自不同 token、不同层天然是 cross-attention。3.2 KV 源集合形式化定义S(i,l){(j,m):ji, l≤m≤top(l)}∪{σ},top(l)min(lk, L−1) S(i,l) \{(j,m) : j i,\; l \le m \le \mathrm{top}(l)\} \cup \{\sigma\}, \qquad \mathrm{top}(l) \min(lk,\, L-1)S(i,l){(j,m):ji,l≤m≤top(l)}∪{σ},top(l)min(lk,L−1)即层 l 的注意力可以读取前文所有 token 在层 l…lk 的成品 KV外加一个自身槽位 σ。其中 k 是深度窗口depth_window是本架构的核心超参数。3.3 一次 softmax 里的三组 key组来源掩码说明① 跨层窗口组更早块中前文 token 在层 l…top(l) 的成品KV仅 j i每组一个源层 mKV 经过该层自己的 RMSNorm 与 K/V 投影② 块内组本块内 j ≤ i 的施工态KV源 本层输入严格下三角cT 时①为空②即完整因果注意力③ 自身槽位 σtoken 自己对角线深度偏置取 0另加每层可学习标量 b_σ3.4 深度轴的相对位置编码深度窗口让读第几层成为新的自由度必须告诉模型它在读多深。两个机制都只走注意力通路深度偏置每头一个标量 λ_h对跨层组分数减λ_h·(m−l)组序号 g。λ_max 上界约 2.0depth_bias开关控制。层编码每层一对可学习向量 e_m^K / e_m^V 直接加进 KV 投影输出f_l 加进 Qlayer_enc开关。3.5 参数量与标准层同构每个 Beta12Layer 的参数 4d²QKVO 3·d·hiddenSwiGLU 2d两个 RMSNorm与标准层完全一致。新增仅每层 e_k/e_v/f_q 三个 d 维向量O(Ld)、每头一个 λ_h、每层一个 b_σ。同参数公平对比成立——消融实验用resolve_fair_config协议同参数预算下缩放宽度/层数保证这一点。4. 依赖图与波前调度m l 的成品要等层 m 算完才存在因此前向计算必须沿反对角波前推进token 每减 1、层升 k。实现上按块切分block_size c这是理解该架构的关键开关c 1严格按规范逐单元串行块内无任何近似墙钟最贵约 T×L 倍标准前向的关键路径。c T整段一块历史集合为空只剩块内严格下三角注意力——精确退化为标准 RoPE Transformer。这既是理论上的特例也是实现正确性的自检臂。中间值如 c 32块内各层的跨层源 之前块的成品 KV增量 KV 在整块完成后才整体重绑块内读到的始终是块开始时的状态。块内读不到本块更早 token 的深层成品是对规范的小偏差但换来块内大矩阵并行。消融结论c 对质量的影响在噪声级见 §7对速度影响巨大——c 越大越快。5. 深度分析两种口径口径一单跳聚合深度 L k一个 token 一次前向内最深读到的成品经过了 (top(0)1) 次层计算再爬完自身栈。对 t19_flashL4, k26 次层应用。这是表达深度的保守口径注意它是同权重栈的复用不是异权重堆叠。口径二链式串行深度 ≈ (T/c) × L随 seqlen 线性增长层 l 的块内组构成水平链(i,l) ← (i-1,l) ← ...RNN 性质。串行关键路径约 (T/c)×L架构串行深度深度的成本标准 Transformer恒为 L—层间循环 TFloopedR × L恒定每token R×L 次层应用真金白银的 FLOPsbeta_1.2(T/c) × L随上下文增长链是因果结构免费给的每 token 只花 L×(1α) 次层应用α≈跨层组开销这是 beta_1.2 与循环 TF 的本质区别循环 TF 是每 token 自适应加深可控但要买beta_1.2 是随上下文自动加深免费但不可控。对串行依赖超过物理层数的任务长状态跟踪、长链组合推理beta_1.2 原理上没有固定层数上限。6. 推理与工程实现6.1 两种推理模式增量解码state 接口与标准 Transformer 的逐层 (K,V) 缓存格式一致可直接增量解码。有状态 RNN 模式inference_modeauto首轮对前缀做一次 prefill之后每步只算最新 token 缓存状态O(1) 每 token——链深度带来的 RNN 性质在推理端兑现为常数成本。6.2 生产路径use_sdpa把跨层窗口组 块内组 自身槽位合并为单一 KV 序列深度偏置、因果掩码全部编码进加性 float bias一次F.scaled_dot_product_attention完成flash kernel不物化 attn 矩阵。三组 key 的语义完全保留只是调度方式合并。6.3 确定性与显存KV 增量按块累积、块完成后整体重绑块内读到的始终是块开始时的状态梯度检查点重算读同一引用数值确定性成立。按波前块做梯度检查点把跨块 KV 拼接链的中间激活峰值压到一个块的量级长上下文/大模型训练不爆显存。可选qk_normQ/K 投影后 RMS 归一与deep_token_window深层源组只看最近 w 个 token长上下文下控制跨层组检索成本。7. 实验证据消融7.1 cblock_size消融——质量噪声级速度敏感~10M 参数、20k batch、B16、seed 42summary_20260925_021817模型test lossppl时间beta_1.2_test (c32)2.692414.770.876 s/batchbeta_1.2_c256 (c256)2.695114.810.124 s/batcht16_std_tf_rope标准 TF 对照2.694714.770.094 s/batchc32 与 c256 的质量差 ~0.003噪声级速度差 ~7×prefill 基准8192 上下文同样显示 c 大幅加速0.24 vs 1.45 ms/tok。c 是纯效率旋钮k 才是效果杠杆。7.2 深度效率消融——10 层跨层模型打平 20 层标准 TF同参数公平协议物理层数减半、加宽 SwiGLU 补齐总参数、600 batch内部实验记录summary_20260925_105014模型配置test loss时间t16_std_tf_ropeL20 标准 TF2.81390.094 sbeta_1.2_de10k6L10, k62.81250.398 sbeta_1.2_de10k2L10, k22.81280.341 sbeta_1.2_stride2L10, 深度采样步长22.81691.289 s半物理层的 beta_1.2 打平/略优于全深度标准 TF等效倍率约 2×且与 k 关系不大——深度主要来自链深度窗口的边际收益递减快。代价是墙钟慢 3.6–4.2×训练时块间串行的账目前用 c32 部分偿还。7.3 早期串行能力证据beta_1 的串行测试台变量搬运通道n8 0.859→0.968n32 0.374→0.432——链深度确实在做标准 TF 结构上做不到的事但长链利用率不完全这是当前明确的短板与改进方向。8. t19_flash当前主力配置t19_flash:{vocab_size:-1,d_model:1600,hidden_dim:3200,num_heads:8,num_layers:4,dropout:0.0,max_seq_len:16384,depth_window:2,block_size:32,lambda_max:2.0,depth_bias:True,layer_enc:True,self_slot:True,qk_norm:False,deep_token_window:0,tie_embedding:True,use_sdpa:True,},约 0.19B 参数含 tied embedding。物理上只有 4 层单跳聚合深度 6串行深度在 16K 上下文下关键路径约 (16384/32)×4 ≈ 2048 次层应用——以 4 层的钱买循环 500 轮的依赖深度。9. 相关工作定位标准 TransformercT 特例是本架构的自检臂。RNN / 状态空间模型共享串行深度随序列增长的性质但 beta_1.2 保留了注意力一次检索全上下文的宽度优势且 KV 范式与现有推理栈兼容。Universal / looped Transformer同为权重复用扩深度但方向不同——循环 TF 在深度轴循环每 token 付 R×L FLOPsbeta_1.2 在序列轴循环链深度免费。相同 FLOPs 预算下looped 的 R 被压缩beta_1.2 的深度/算力比更高looped 则保留 per-token 自适应计算ACT的优势。跨层 KV 复用CLA 类那些工作是复用 KV省算力beta_1.2 是读取更深层 KV扩能力方向相反。10. 局限与进行中的工作训练墙钟块间串行使同参数训练比标准 TF 慢数倍c 调大可偿还大部分但块内施工态注意力的显存随 c 增长存在权衡。深度窗口边际递减k2→k6 几乎无增益深度来源已经从窗口转移到链窗口更多是给链提供重组接口。长链利用率串行深度纸面无限实际利用受训练梯度穿过链的程度限制beta_1 搬运通道证据。评测口径现有结论均为 LM loss 口径、单 seed、早期 batch 数600–20k推理能力编程、状态跟踪的同台对比尚未完成。已规划的方向参数公平 sweep 标准 TF 层数、串行状态跟踪测试台、与 looped TF 臂的直接对比。11. 总结beta_1.2 用一个最小侵入的改动——把注意力的 KV 源沿深度轴打开——同时拿到了三样东西免费生长的串行深度RNN 性质突破固定层数的推理天花板与标准 Transformer 完全同构的参数量与骨架公平可比、即插即用cT 退化为标准 TF 的理论闭环使它成为标准架构的严格超集而非另一个模型。它的性格可以概括为一句话深度不靠堆层靠上下文。实验数据来源于小思框架内部所有消融为单 seed 结果解读请参考 §10 的局限性说明。