层归一化、损失函数与 Universal Transformer 延展

小象实战讲义 · 自然语言处理

在掌握了 Transformer 的核心组件——自注意力与位置编码之后，本节我们将深入其内部的两个关键细节：层归一化（Layer Normalization）与损失函数的计算。这些技术是稳定深层网络训练、实现有效优化的基石。最后，我们将探讨一种基于自适应计算时间（Adaptive Computation Time, ACT）的 Transformer 变体——Universal Transformer，它为解决“模型深度如何自适应”这一经典问题提供了新颖的思路。本节内容将帮助你理解现代 Transformer 架构中不可或缺的工程技巧，并了解模型设计的前沿探索。

💡 核心导读

• 层归一化（Layer Norm）：理解其在 Transformer 中作为“残差学习”稳定器的作用，掌握其与批归一化（Batch Norm）的核心区别。
• Softmax 层与损失函数：从解码器的线性层出发，掌握如何将隐藏状态映射为词表上的概率分布，并使用交叉熵损失进行优化。
• 自适应计算时间（ACT）：学习一种让 RNN 在单个时间步内进行“深度思考”的机制，理解其如何通过“停机概率”动态控制计算量。
• Universal Transformer：了解如何将 ACT 机制融入 Transformer，构建一个编码器和解码器层数可动态调整的“通用”模型。

层归一化：稳定深层 Transformer 的训练

在 Transformer 的编码器和解码器内部，每个子层（如自注意力层、前馈网络）的输出都会紧接着一个层归一化（Layer Normalization, Layer Norm）操作。它的核心作用是稳定深层网络的训练过程。

回顾我们之前构建的编码器流程：输入词向量 $X$ 加上位置编码 $P$ 后，得到新的输入 $\tilde{X}$。$\tilde{X}$ 经过多头自注意力（Multi-Head Self-Attention）计算后，我们得到了一个中间输出，通常记为 $Z_{attn}$。

# 示例：模拟自注意力后的输出 Z_attn (假设 batch_size=2, seq_len=3, hidden_dim=4)
import numpy as np

# 假设自注意力层输出的 Z_attn
Z_attn = np.array([
    [[1.2, -0.5, 0.8, 2.1],  # 序列1，词1
     [0.3, 1.5, -1.2, 0.4],  # 序列1，词2
     [-0.9, 0.2, 1.8, -0.3]], # 序列1，词3
    [[2.0, 0.1, -0.7, 1.5],  # 序列2，词1
     [-0.5, 1.8, 0.3, -1.0], # 序列2，词2
     [1.1, -0.4, 0.9, 0.6]]  # 序列2，词3
])
print("自注意力输出 Z_attn 形状:", Z_attn.shape)

自注意力输出 Z_attn 形状: (2, 3, 4)

在标准的 Transformer 中，我们不会直接将 $Z_{attn}$ 送入前馈网络。相反，我们引入了一个残差连接（Residual Connection）和层归一化。

残差连接借鉴了 ResNet 的思想：将子层的输入 $\tilde{X}$ 直接加到其输出 $Z_{attn}$ 上，形成 $Z_{res} = \tilde{X} + Z_{attn}$。这样，网络只需要学习输入与输出之间的残差（即变化部分），这有助于缓解深层网络中的梯度消失问题。

层归一化则作用于这个残差和 $Z_{res}$ 上。与批归一化（Batch Norm）在批次维度上进行归一化不同，层归一化是在特征维度（hidden dimension）上，对单个样本的单个位置（token）的所有特征进行归一化。

其计算公式为： $$\text{LayerNorm}(x) = \gamma \odot \frac{x - \mu}{\sigma} + \beta$$ 其中：

• $x$ 是单个位置的特征向量（例如维度为 512）。
• $\mu$ 是该向量所有元素的均值，$\sigma$ 是其标准差（加上一个小常数 $\epsilon$ 防止除零）。
• $\gamma$（缩放参数）和 $\beta$（平移参数）是可学习的参数，维度与 $x$ 相同。

# —— 重新构建上文用到的微型数据（本节各代码块相互独立，可单独运行）
import numpy as np
Z_attn = np.array([
    [[1.2, -0.5, 0.8, 2.1],  # 序列1，词1
     [0.3, 1.5, -1.2, 0.4],  # 序列1，词2
     [-0.9, 0.2, 1.8, -0.3]], # 序列1，词3
    [[2.0, 0.1, -0.7, 1.5],  # 序列2，词1
     [-0.5, 1.8, 0.3, -1.0], # 序列2，词2
     [1.1, -0.4, 0.9, 0.6]]  # 序列2，词3
])

def layer_norm(x, gamma=1.0, beta=0.0, eps=1e-5):
    """
    简化的层归一化实现（单个样本，单个位置）。
    x: 输入向量 (hidden_dim,)
    gamma, beta: 可学习参数 (hidden_dim,)
    """
    mu = np.mean(x)
    sigma = np.std(x)
    x_norm = (x - mu) / np.sqrt(sigma**2 + eps)
    return gamma * x_norm + beta

# 取 Z_res 中第一个样本的第一个位置作为示例 (假设已加残差)
sample_vector = np.array([1.5, -0.5, 0.8, 2.2])
gamma = np.ones(4)  # 假设初始化为1
beta = np.zeros(4)  # 假设初始化为0
output = layer_norm(sample_vector, gamma, beta)
print("输入向量:", sample_vector)
print("层归一化后:", np.round(output, 4))

输入向量: [ 1.5 -0.5  0.8  2.2]
层归一化后: [ 0.5013 -1.5038 -0.2005  1.203 ]

经过层归一化后，我们才得到该子层真正的输出 $Z_1$，它将被送入前馈网络（Feed-Forward Network, FFN）。FFN 的输出同样会经过另一个残差连接和层归一化。因此，一个完整的 Transformer 编码器块可以概括为： Positional Encoding → Multi-Head Self-Attention → Add & LayerNorm → Feed-Forward Network → Add & LayerNorm。

Softmax 层与交叉熵损失：从解码到优化

解码器（Decoder）的最终目标是生成目标语言序列。在 Transformer 的解码器中，最后一个子层（FFN + LayerNorm）的输出是一个维度为 [batch_size, seq_len, hidden_dim] 的张量。为了将其转换为词表上的概率分布，我们需要经过一个 线性层（Linear Layer） 和一个 Softmax 层。

线性层：从隐藏空间到词表空间

线性层是一个全连接层，其权重矩阵 $W_{proj}$ 的形状为 [hidden_dim, vocab_size]，其中 vocab_size 是目标词表的大小（例如 10,000）。它将每个位置的隐藏向量映射到一个 vocab_size 维的 logits 向量。logits 可以理解为每个词未归一化的“得分”。

# —— 重新构建上文用到的微型数据（本节各代码块相互独立，可单独运行）
import numpy as np
Z_attn = np.array([
    [[1.2, -0.5, 0.8, 2.1],  # 序列1，词1
     [0.3, 1.5, -1.2, 0.4],  # 序列1，词2
     [-0.9, 0.2, 1.8, -0.3]], # 序列1，词3
    [[2.0, 0.1, -0.7, 1.5],  # 序列2，词1
     [-0.5, 1.8, 0.3, -1.0], # 序列2，词2
     [1.1, -0.4, 0.9, 0.6]]  # 序列2，词3
])
def layer_norm(x, gamma=1.0, beta=0.0, eps=1e-5):
    """
    简化的层归一化实现（单个样本，单个位置）。
    x: 输入向量 (hidden_dim,)
    gamma, beta: 可学习参数 (hidden_dim,)
    """
    mu = np.mean(x)
    sigma = np.std(x)
    x_norm = (x - mu) / np.sqrt(sigma**2 + eps)
    return gamma * x_norm + beta
sample_vector = np.array([1.5, -0.5, 0.8, 2.2])
gamma = np.ones(4)  # 假设初始化为1
beta = np.zeros(4)  # 假设初始化为0
output = layer_norm(sample_vector, gamma, beta)

# 模拟解码器最后一个隐藏状态 (batch_size=1, seq_len=2, hidden_dim=4)
hidden_state = np.array([[[0.8, -0.2, 1.5, 0.3], [0.1, 0.9, -0.4, 0.7]]])
vocab_size = 6  # 假设一个小词表: ['a', 'am', 'i', 'thanks', 'student', '<eos>']
hidden_dim = hidden_state.shape[-1]

# 随机初始化投影权重
np.random.seed(42)
W_proj = np.random.randn(hidden_dim, vocab_size) * 0.1

# 计算 logits: (1, 2, 4) @ (4, 6) -> (1, 2, 6)
logits = np.dot(hidden_state, W_proj)
print("隐藏状态形状:", hidden_state.shape)
print("投影权重形状:", W_proj.shape)
print("Logits 形状:", logits.shape)
print("第一个位置的 logits:", np.round(logits[0, 0], 4))

隐藏状态形状: (1, 2, 4)
投影权重形状: (4, 6)
Logits 形状: (1, 2, 6)
第一个位置的 logits: [ 0.0172 -0.3558 -0.1536  0.0199 -0.1594 -0.005 ]

Softmax：从得分到概率

Softmax 函数将 logits 向量转换成一个概率分布。对于向量 $z$，第 $i$ 个元素的概率为： $$p_i = \frac{e^{z_i}}{\sum_{j=1}^{V} e^{z_j}}$$ 这样，所有词的概率之和为 1，概率最大的词即被选为当前时间步的输出。

# —— 重新构建上文用到的微型数据（本节各代码块相互独立，可单独运行）
import numpy as np
np.random.seed(42)  # 固定随机种子，确保输出可复现
Z_attn = np.array([
    [[1.2, -0.5, 0.8, 2.1],  # 序列1，词1
     [0.3, 1.5, -1.2, 0.4],  # 序列1，词2
     [-0.9, 0.2, 1.8, -0.3]], # 序列1，词3
    [[2.0, 0.1, -0.7, 1.5],  # 序列2，词1
     [-0.5, 1.8, 0.3, -1.0], # 序列2，词2
     [1.1, -0.4, 0.9, 0.6]]  # 序列2，词3
])
def layer_norm(x, gamma=1.0, beta=0.0, eps=1e-5):
    """
    简化的层归一化实现（单个样本，单个位置）。
    x: 输入向量 (hidden_dim,)
    gamma, beta: 可学习参数 (hidden_dim,)
    """
    mu = np.mean(x)
    sigma = np.std(x)
    x_norm = (x - mu) / np.sqrt(sigma**2 + eps)
    return gamma * x_norm + beta
sample_vector = np.array([1.5, -0.5, 0.8, 2.2])
gamma = np.ones(4)  # 假设初始化为1
beta = np.zeros(4)  # 假设初始化为0
output = layer_norm(sample_vector, gamma, beta)
hidden_state = np.array([[[0.8, -0.2, 1.5, 0.3], [0.1, 0.9, -0.4, 0.7]]])
vocab_size = 6  # 假设一个小词表: ['a', 'am', 'i', 'thanks', 'student', '<eos>']
hidden_dim = hidden_state.shape[-1]
W_proj = np.random.randn(hidden_dim, vocab_size) * 0.1
logits = np.dot(hidden_state, W_proj)

def softmax(z):
    exp_z = np.exp(z - np.max(z))  # 减去最大值防止数值溢出
    return exp_z / np.sum(exp_z)

# 对第一个位置的 logits 应用 softmax
first_logits = logits[0, 0]
probs = softmax(first_logits)
print("Logits:", np.round(first_logits, 4))
print("Softmax 概率分布:", np.round(probs, 4))
print("概率和:", np.round(np.sum(probs), 4))

Logits: [ 0.0172 -0.3558 -0.1536  0.0199 -0.1594 -0.005 ]
Softmax 概率分布: [0.1869 0.1287 0.1576 0.1874 0.1566 0.1828]
概率和: 1.0

交叉熵损失：衡量预测与目标的差距

训练时，我们需要一个损失函数来指导模型参数的更新。在分类任务（如预测下一个词）中，最常用的是交叉熵损失（Cross-Entropy Loss）。

假设我们的目标词是“thanks”，在词表中的索引是 3（从0开始）。那么它的 One-Hot 表示 是一个只在索引 3 处为 1，其余为 0 的向量：$y = [0, 0, 0, 1, 0, 0]$。

交叉熵损失衡量的是模型预测的概率分布 $p$ 与真实分布 $y$ 之间的差异： $$L = -\sum_{i=1}^{V} y_i \log(p_i)$$ 本例为教学推导采用 One-Hot 硬标签，只有目标索引 $t$ 处 $y_t=1$，因此上式简化为： $$L = -\log(p_t)$$ 需要注意：One-Hot 只是未做标签平滑时的硬标签特例。原始 Transformer 训练采用 $\varepsilon_{ls}=0.1$ 的标签平滑，目标是平滑后的分布 $q$，此时应按完整式 $L=-\sum_i q_i \log(p_i)$ 计算，而不是严格的 One-Hot 形式。 即，损失是目标词预测概率的负对数。概率越接近1，损失越小。

# —— 重新构建上文用到的微型数据（本节各代码块相互独立，可单独运行）
import numpy as np
np.random.seed(42)  # 固定随机种子，确保输出可复现
Z_attn = np.array([
    [[1.2, -0.5, 0.8, 2.1],  # 序列1，词1
     [0.3, 1.5, -1.2, 0.4],  # 序列1，词2
     [-0.9, 0.2, 1.8, -0.3]], # 序列1，词3
    [[2.0, 0.1, -0.7, 1.5],  # 序列2，词1
     [-0.5, 1.8, 0.3, -1.0], # 序列2，词2
     [1.1, -0.4, 0.9, 0.6]]  # 序列2，词3
])
def layer_norm(x, gamma=1.0, beta=0.0, eps=1e-5):
    """
    简化的层归一化实现（单个样本，单个位置）。
    x: 输入向量 (hidden_dim,)
    gamma, beta: 可学习参数 (hidden_dim,)
    """
    mu = np.mean(x)
    sigma = np.std(x)
    x_norm = (x - mu) / np.sqrt(sigma**2 + eps)
    return gamma * x_norm + beta
sample_vector = np.array([1.5, -0.5, 0.8, 2.2])
gamma = np.ones(4)  # 假设初始化为1
beta = np.zeros(4)  # 假设初始化为0
output = layer_norm(sample_vector, gamma, beta)
hidden_state = np.array([[[0.8, -0.2, 1.5, 0.3], [0.1, 0.9, -0.4, 0.7]]])
vocab_size = 6  # 假设一个小词表: ['a', 'am', 'i', 'thanks', 'student', '<eos>']
hidden_dim = hidden_state.shape[-1]
W_proj = np.random.randn(hidden_dim, vocab_size) * 0.1
logits = np.dot(hidden_state, W_proj)
def softmax(z):
    exp_z = np.exp(z - np.max(z))  # 减去最大值防止数值溢出
    return exp_z / np.sum(exp_z)
first_logits = logits[0, 0]
probs = softmax(first_logits)

def cross_entropy_loss(probs, target_idx):
    return -np.log(probs[target_idx] + 1e-8)  # 加一个小数防止 log(0)

target_word = 'thanks'
vocab = ['a', 'am', 'i', 'thanks', 'student', '<eos>']
target_idx = vocab.index(target_word)

loss = cross_entropy_loss(probs, target_idx)
print(f"目标词: '{target_word}' (索引 {target_idx})")
print(f"目标词预测概率: {probs[target_idx]:.4f}")
print(f"交叉熵损失: {loss:.4f}")

目标词: 'thanks' (索引 3)
目标词预测概率: 0.1874
交叉熵损失: 1.6745

整个解码过程是自回归的：模型根据已生成的词（初始为开始符）和编码器的输出，预测下一个词的概率分布，选取概率最高的词作为输出，并将其作为下一时间步的输入，直到生成结束符 <eos>。

自适应计算时间（ACT）：让 RNN “深度思考”

标准的 RNN 在每个时间步进行固定且简单的计算。自适应计算时间（Adaptive Computation Time, ACT） 是一种机制，允许 RNN 在单个时间步内根据输入复杂度，动态地进行多次“微时间步”的计算，从而更深入地处理当前信息。

ACT 的核心思想

想象一下你的工作日：周一邮件繁多，你需要花更多时间（进行更多思考步骤）来处理；周三任务简单，很快就能完成。ACT 希望 RNN 具备类似的能力。在 RNN 的每个“宏观”时间步 $t$ 内，ACT 引入了一个内部循环，可以执行 $N_t$ 次“微观”计算步骤。$N_t$ 不是固定的，而是由模型根据当前输入动态决定的。

ACT 的运作机制

1. 初始状态：在时间步 $t$，我们有上一时刻的隐藏状态 $s_{t-1}$ 和当前输入 $x_t$。
2. 内部循环：
   • 从 $n=1$ 开始，ACT 单元计算一个中间隐藏状态 $s_{t}^{(n)}$（计算方式类似于普通 RNN）。
   • 同时，它计算一个 停机权重（halting weight） $h_t^{(n)}$，这是一个介于 0 和 1 之间的数（通过 Sigmoid 函数得到），可以理解为“在完成第 $n$ 步计算后，决定停机的概率”。
3. 累积与判断：
   • 维护一个剩余的“停机预算” $R$，初始为 1。
   • 每次计算后，从 $R$ 中减去当前步的停机权重 $h_t^{(n)}$。
   • 如果 $R$ 小于一个很小的阈值 $\epsilon$（如 0.01），则停止内部循环。
4. 输出：
   • 假设内部循环进行了 $N$ 步。我们需要对 $N$ 个中间状态 $s_t^{(1)}, ..., s_t^{(N)}$ 和对应的输出 $y_t^{(n)}$ 进行加权平均，以得到时间步 $t$ 的最终隐藏状态 $s_t$ 和输出 $y_t$。
   • 权重需要归一化。通常，前 $N-1$ 步的权重就是其停机权重 $h_t^{(n)}$，最后一步的权重是剩余的预算 $R_{t}^{(N-1)}$（确保所有权重之和为1）。
5. 训练中的计算量惩罚：训练 ACT 时，除主任务损失外，还需在目标函数中加入一项 ponder cost（思考成本），对模型在该样本上使用的总计算量（各微时间步停机权重之和）施加惩罚。该项的系数越大，模型越倾向于少用计算步；系数越小，则更偏向任务精度。它在预测精度与降低计算量之间作权衡，主任务损失本身并不会自动促使模型节省计算。

图示：ACT 在一个时间步内的计算流程 (以3步为例)

时间步 t-1 隐藏状态 s_{t-1}
         |
         v
     输入 x_t
         |
         v
+-----------------------+
|     内部循环开始       |
|  n=1:                 |
|  s_t^(1) = f(s_{t-1}, x_t) |
|  h_t^(1) = σ( W_h * s_t^(1) ) |
|  R = 1 - h_t^(1)      |
|                       |
|  n=2: (if R > ε)      |
|  s_t^(2) = f(s_t^(1), x_t) |
|  h_t^(2) = σ( W_h * s_t^(2) ) |
|  R = R - h_t^(2)      |
|                       |
|  n=3: (if R > ε)      |
|  s_t^(3) = f(s_t^(2), x_t) |
|  h_t^(3) = σ( W_h * s_t^(3) ) |
|  R = R - h_t^(3)      |
|  (假设此时 R < ε，停止) |
+-----------------------+
         |
         v
计算最终状态 s_t 和输出 y_t：
s_t = Σ_{n=1 to 3} (权重_n * s_t^(n))
y_t = Σ_{n=1 to 3} (权重_n * y_t^(n))
其中，权重_1 = h_t^(1), 权重_2 = h_t^(2), 权重_3 = R_before_step3

本质上，ACT 在单个 RNN 时间步内部引入了一个注意力机制，该机制关注的是不同计算深度（微观步）下的状态，并根据输入的复杂度动态分配计算资源。

Universal Transformer：动态深度的 Transformer

标准的 Transformer 需要预先固定编码器和解码器的层数（如 6 层）。Universal Transformer 将 ACT 的思想应用到了 Transformer 上，使其编码器和解码器的层数（即计算深度）能够动态自适应。

核心架构

Universal Transformer 的总体架构与 Transformer 相似，但有一个关键区别：它不再是固定数量的编码器/解码器块堆叠，而是同一个编码器块（或解码器块）在时间维度上循环执行 $T$ 次。这里的“时间”指的是深度方向上的循环，而非序列长度方向。

1. 输入：词嵌入 + 位置编码。
2. 循环块（Recurrent Block）：这个块包含 Transformer 编码器的核心操作：自注意力、层归一化、前馈网络等。
3. 自适应深度：通过集成 ACT 机制，模型可以动态决定这个循环块需要执行多少次（即深度 $T$ 是多少）。对于序列中不同的位置（token），模型甚至可以分配不同的计算步骤。
4. 输出：经过 $T$ 步循环后，得到最终的编码表示，传递给解码器。解码器部分也可以采用类似的自适应深度结构。

意义与局限

Universal Transformer 试图解决神经网络架构设计中的一个根本问题：对于不同的输入样本或同一序列中的不同部分，其所需的模型复杂度（深度）可能是不同的。通过 ACT，它提供了一种优雅的、可微分的方式来自适应计算量。此外，在内存充足等理论假设下，其深度方向上的循环机制使模型具备计算通用性（computationally universal）；这一“通用性”定位与 ACT 带来的自适应计算量是相关但不同的两件事：前者指理论上可计算的函数类，后者指按输入动态分配计算步。

然而，在实际应用中，尤其是大规模预训练场景下，Universal Transformer 并未成为主流。主要原因在于其计算效率：动态深度的循环计算难以像固定层数的 Transformer 那样进行高度并行化优化。此外，训练 ACT 机制本身也需要额外的技巧和成本。

尽管如此，Universal Transformer 代表了模型设计的一个重要研究方向——条件计算（Conditional Computation），即让模型的计算图根据输入动态变化。需要注意时间线：混合专家（Mixture of Experts, MoE）这一思想早在 1991 年的自适应局部专家混合模型中就已提出，稀疏门控 MoE 也在 2017 年发表，均早于 Universal Transformer；UT 的条件计算思路与它们相呼应，后续 Switch Transformer 等模型继续发展了稀疏条件计算，不能说 MoE 是 UT 之后的延续。

📝 动手练一练

1. 概念辨析：层归一化（LayerNorm）和批归一化（BatchNorm）在归一化的维度上有何本质区别？为什么在 Transformer 这类序列模型中，LayerNorm 通常比 BatchNorm 更适用？

2. 手算练习：假设一个解码器输出某个位置的 logits 向量为 [2.0, 1.0, 0.1, -1.0]，对应的词表是 ['猫', '狗', '鸟', '鱼']。如果真实的下一个词是“狗”，请计算此处的交叉熵损失。

本章小结

本节我们深入剖析了 Transformer 架构中的两个关键工程技术细节，并探索了一个前沿的变体模型：

• 层归一化（LayerNorm） 与残差连接协同工作，是稳定深层 Transformer 训练、加速收敛的关键组件。它通过对单个样本单个位置的特征进行归一化，避免了批归一化在序列任务中的弊端。
• Softmax 层与交叉熵损失 构成了解码器输出的最后一步，将隐藏状态映射为词表上的概率分布；本例按 One-Hot 硬标签推导交叉熵，原始 Transformer 则使用 $\varepsilon_{ls}=0.1$ 的标签平滑，以平滑目标分布计算损失。
• 自适应计算时间（ACT） 是一种创新的机制，使 RNN 能够在单个时间步内进行可变次数的内部计算，实现了计算量的动态分配。
• Universal Transformer 将 ACT 的思想引入 Transformer，通过深度方向上的循环和自适应停机机制，构建了一个编码器/解码器深度可动态调整的“通用”模型。它代表了条件计算这一重要研究方向，尽管因效率问题未成为主流，但其思想影响深远。

📋 行动清单

• 在纸上推导一遍层归一化的公式，并理解其与批归一化在计算维度上的差异。
• 编写一个简单的函数，输入 logits 向量和目标词索引，输出 softmax 概率和交叉熵损失，并用小例子验证。
• 思考：除了 Universal Transformer，还有哪些模型设计体现了“条件计算”或“动态结构”的思想？（提示：可以搜索“Mixture of Experts”, “Switch Transformer”）

—— 小象教研组