Keel：深度可扩展Transformer与Highway连接

1. 背景与动机

现代大语言模型（LLM）的进展主要依赖三条扩展路线：模型宽度、上下文长度和参数规模。然而，这三条路线都面临收益递减的问题：

扩展路线	问题
宽度扩展	快速饱和，计算复杂度 $O (d^{2})$
上下文扩展	成本急剧增加，外推能力有限
深度扩展	理论上表达力最优，但训练不稳定

深度扩展之所以困难，是因为传统的Post-LayerNorm架构在训练深层网络时会出现梯度消失问题。2019年”On Layer Normalization in the Pre-TayerNorm Transformer”指出Post-LN在训练初期不稳定，这导致了Pre-LayerNorm（Pre-LN）成为现代LLM的标准配置——将LayerNorm移到残差分支内部，使主分支成为恒等映射。

但Pre-LN牺牲了Post-LN的表达力。Keel论文（arXiv:2601.19895）重新审视Post-LN，发现其不稳定性的根本原因是ResNet风格的残差路径，并提出用Highway连接替代。

2. Pre-LN vs Post-LN 问题分析

2.1 Pre-LayerNorm架构

// Pre-LayerNorm (现代LLM标准)
class PreLNTransformerLayer {
    Tensor norm1(Tensor x) { return LayerNorm(x); }
    
    Tensor forward(Tensor x) {
        Tensor x_norm = norm1(x);
        Tensor attn_out = Attention(x_norm);
        x = x + attn_out;  // 残差分支
        
        x_norm = norm2(x);
        Tensor ffn_out = FFN(x_norm);
        x = x + ffn_out;   // 残差分支
        return x;
    }
}

Pre-LN的特点：

主分支始终是恒等映射 $x \leftarrow x + f (x)$
梯度可以无衰减地流向底层
训练稳定，但表达能力受限

2.2 Post-LayerNorm架构

// Post-LayerNorm (早期BERT等使用)
class PostLNTransformerLayer {
    Tensor forward(Tensor x) {
        Tensor attn_out = Attention(x);
        x = LayerNorm(x + attn_out);  // 归一化在加法之后
        
        Tensor ffn_out = FFN(x);
        x = LayerNorm(x + ffn_out);
        return x;
    }
}

Post-LN的特点：

归一化层作用于 $x + f (x)$ 的和
训练初期，残差分支输出尺度与主分支相近
LayerNorm的缩放效应导致底层梯度极小

2.3 梯度消失的数学分析

考虑第 $l$ 层的输出 $x^{l}$ ：

x^{l} = x^{l - 1} + f_{θ}^{l} (x^{l - 1})

反向传播时，损失 $L$ 对 $x^{l - 1}$ 的梯度为：

\frac{\partial L}{\partial x ^{l - 1}} = \frac{\partial L}{\partial x ^{l}} \cdot \frac{\partial x ^{l}}{\partial x ^{l - 1}} = \frac{\partial L}{\partial x ^{l}} \cdot (I + J_{f})

其中 $J_{f} = \frac{\partial f _{θ}^{l}}{\partial x ^{l - 1}}$ 是雅可比矩阵。

在训练初期， $f_{θ}^{l}$ 的输出方差与输入相近，导致 $J_{f}$ 的期望约为 $c \cdot I$ （ $c$ 是常数）。此时：

∥ \frac{\partial L}{\partial x ^{l - 1}} ∥ \approx ∣1 + c ∣ \cdot ∥ \frac{\partial L}{\partial x ^{l}} ∥

当 $c$ 为负且接近 $- 1$ 时，梯度会被衰减。LayerNorm的统计效应进一步放大了这个问题。

3. Keel架构设计

3.1 Highway连接机制

Highway网络的核心思想来自Hochreiter的早期工作：

y = T (x) \cdot H (x) + (1 - T (x)) \cdot x

其中：

$H (x)$ 是变换门（Transformation Gate）
$T (x) = σ (W_{T} x + b_{T})$ 是门控信号
$(1 - T (x))$ 是携带门（Carry Gate）

3.2 Keel的Highway残差设计

Keel将标准的加性残差替换为Highway连接：

// Keel Transformer层
class KeelTransformerLayer {
    Tensor highway_gate(Tensor x) {
        return sigmoid(W_gate * x + b_gate);  // [0,1] 范围的门控
    }
    
    Tensor forward(Tensor x) {
        // 计算变换
        Tensor h = Attention(x);
        h = x + h;  // 基础变换
        
        // Highway门控
        Tensor t = highway_gate(x);  // [0,1]
        
        // 门控组合
        Tensor y = t * h + (1 - t) * x;
        
        // 第二个块
        h = FFN(y);
        h = y + h;
        
        t = highway_gate(y);
        y = t * h + (1 - t) * y;
        
        return y;
    }
}

3.3 梯度流分析

使用Highway连接后，前向传播变为：

x^{l} = T (x^{l - 1}) \cdot f_{θ}^{l} (x^{l - 1}) + (1 - T (x^{l - 1})) \cdot x^{l - 1}

反向传播梯度为：

\frac{\partial L}{\partial x ^{l - 1}} = \frac{\partial L}{\partial x ^{l}} \cdot [(1 - T) \cdot I + T \cdot J_{f} + f \cdot \frac{\partial T}{\partial x ^{l - 1}}]

关键观察：

即使 $T \approx 1$ （变换主导），梯度仍可通过 $J_{f}$ 流动
即使 $T \approx 0$ （携带主导），梯度可通过 $(1 - T)$ 直接传递
Highway网络天然支持 自适应深度：网络自己决定每层应该变换多少、携带多少

4. 与标准Pre-LN的对比

特性	Pre-LN	Post-LN	Keel
训练稳定性	✓ 稳定	✗ 不稳定	✓ 稳定
表达能力	中等	高	高
梯度流	恒等保证	依赖初始化	自适应门控
深度可扩展	受限	不可能	支持1000+层
收敛速度	中等	快（初期）	快且稳定

4.1 深度缩放实验

Keel论文的实验显示：

深度     | Pre-LN Perplexity | Keel Perplexity
--------|-------------------|-----------------
64层    | 10.23            | 10.15
128层   | 9.87             | 9.71
256层   | 9.65             | 9.32
512层   | 9.52             | 8.89
1024层  | 不收敛           | 8.54

Keel在64-1024层的所有深度上都优于Pre-LN，且能够稳定训练超过1000层的模型。

4.2 任务性能对比

在Math & Code任务上，Keel比Pre-LN提升 16.5%：

任务类型	Pre-LN	Keel提升
Math	基准	+18.2%
Code	基准	+15.8%
常识推理	基准	+14.3%
阅读理解	基准	+12.7%

5. 门控信号的动态行为

训练过程中，Keel的门控值呈现有趣的模式：

# 门控值分析
def analyze_gate_values(model, dataloader):
    """分析训练过程中门控值的分布变化"""
    gate_values = {f'layer_{i}': [] for i in range(model.depth)}
    
    for batch in dataloader:
        for layer in range(model.depth):
            t = model.layers[layer].gate_value  # Highway门控
            gate_values[f'layer_{layer}'].append(t.mean().item())
    
    # 绘制门控值随层数变化的曲线
    # 观察到：浅层T值较高（需要变换），深层T值较低（携带信息）
    return gate_values

典型观察：

浅层（第1-32层）： $T \approx 0.7 - 0.9$ ，强变换
中层（第33-128层）： $T \approx 0.5 - 0.7$ ，平衡
深层（第129+层）： $T \approx 0.3 - 0.5$ ，信息携带

这表明Keel自动学会了自适应深度压缩：浅层进行主要变换，深层逐步压缩和精炼信息。

6. 与其他深度扩展技术的对比

6.1 残差累积 vs Highway

标准ResNet:     x_{l+1} = x_l + f(x_l)
Keel Highway:   x_{l+1} = T(x_l) * f(x_l) + (1-T(x_l)) * x_l

6.2 与LSTM门控的关系

LSTM的遗忘门和输入门与Highway连接有相似的形式：

c_{l} = f_{l} ⊙ c_{l - 1} + i_{l} ⊙ \tilde{c}_{l}

但LSTM门控作用于细胞状态，而Keel门控作用于隐藏状态。

7. 实现细节

7.1 门控初始化

为确保训练初期保持稳定，建议：

# Keel实现
class KeelLayer(nn.Module):
    def __init__(self, d_model, init_bias=2.0):
        super().__init__()
        self.gate = nn.Linear(d_model, 1)
        
        # 偏置初始化：使门控初期偏向携带
        nn.init.constant_(self.gate.bias, init_bias)  # sigmoid(2.0) ≈ 0.88
    
    def forward(self, x):
        h = self.attention(x)
        h = x + h
        
        t = self.gate(x).sigmoid()
        return t * h + (1 - t) * x

7.2 训练配置

Keel训练配置:
  学习率: 3e-4 (与标准Pre-LN相同)
  优化器: AdamW
  权重衰减: 0.1
  批大小: 4096
  梯度裁剪: 1.0
  训练步数: 100K+

8. 未来方向

无限深度架构：Keel打开了无限深度训练的可能性
动态计算：可根据输入复杂度自适应选择深度
与其他架构结合：如与Mamba、SSM结合
理论分析：更深入的梯度流和表达能力分析

Metaphor

探索

Keel：深度可扩展Transformer与Highway连接

Keel：深度可扩展Transformer与Highway连接

1. 背景与动机

2. Pre-LN vs Post-LN 问题分析

2.1 Pre-LayerNorm架构

2.2 Post-LayerNorm架构

2.3 梯度消失的数学分析

3. Keel架构设计

3.1 Highway连接机制

3.2 Keel的Highway残差设计

3.3 梯度流分析

4. 与标准Pre-LN的对比

4.1 深度缩放实验

4.2 任务性能对比

5. 门控信号的动态行为

6. 与其他深度扩展技术的对比

6.1 残差累积 vs Highway

6.2 与LSTM门控的关系

7. 实现细节

7.1 门控初始化

7.2 训练配置

8. 未来方向

9. 参考文献

关系图谱

目录