每日基础课

【生活类比】 想象你在调试一段说话人音频的音量： - LayerNorm 就像强迫症调音师：先测整段音频的直流偏移（均值），把波形强行平移到零点，再算音量方差，最后还要加一个可调节的偏移滑块。 - RMSNorm 就像实战派调音师：直接测音量的有效响度（均方根 RMS），响了就调小，轻了就调大，根本不管中心点在哪。实验证明，只要把幅度约束住，模型就能稳稳训练。

【核心公式与符号拆解】 输入一个特征向量 $x = [x_1, x_2, \dots, x_d]$：

1. 计算均方根（RMS）： $$\text{RMS}(x) = \sqrt{\frac{1}{d} \sum_{i=1}^d x_i^2 + \epsilon}$$

2. 归一化并缩放： $$y = \frac{x}{\text{RMS}(x)} \odot \gamma$$

• $x$：输入的单个 token 特征（比如某个音素帧或文本词向量），形状是 $[d]$。
• $d$：隐藏层维度（Hidden Dimension），例如 4096。
• $\epsilon$（epsilon）：一个极小的常数（如 $1e-5$ 或 $1e-6$），防止分母除以 0。
• $\gamma$（gamma）：可学习的缩放权重参数（Weight），形状也是 $[d]$，初始全为 1。
• $\odot$：逐元素相乘。
• 注意：没有可学习的偏置参数 $\beta$（bias），公式里也没有减均值 $\mu$。

【它解决什么问题？没有它会怎样？】 - 解决的问题：在深层 Transformer 中，信号逐层传递会导致数值越来越大（梯度爆炸）或越来越小（梯度消失）。归一化把激活值强行拉回稳定尺度。 - 没有它会怎样：模型稍微深一点就会在训练前几百步直接报 NaN（数值溢出），或者参数完全学不动。 - 为什么不直接用 LayerNorm：LayerNorm 每次都要算均值、减均值、算方差，多次遍历内存（Memory Bandwidth Bound）。RMSNorm 砍掉一半操作，显著减少了显存读写延迟。

【实现思路】 1. 继承 `nn.Module`，初始化可学习参数 `weight`（即 $\gamma$），形状为 `(dim,)`，初值填 1。 2. `forward` 阶段：对最后一维求平方均值 `x.pow(2).mean(-1, keepdim=True)`。 3. 开根号并加 `eps` 得到 RMS，最后计算 `(x / rms) * weight`。

【完整可运行代码】

import torch
import torch.nn as nn

class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        # 可学习参数 gamma，初始化为全 1
        self.weight = nn.Parameter(torch.ones(dim))

    def _norm(self, x: torch.Tensor) -> torch.Tensor:
        # 核心公式: x / sqrt(mean(x^2) + eps)
        # keepdim=True 保证广播机制能正确匹配维度 (B, T, 1)
        return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 先转为 float32 计算以防半精度溢出，算完转回原类型
        output = self._norm(x.float()).type_as(x)
        return output * self.weight

if __name__ == "__main__":
    torch.manual_seed(42)

    # 模拟输入: Batch_size=2, 音频帧序列长度 Seq_len=4, 特征维度 Dim=8
    B, T, D = 2, 4, 8
    dummy_audio_features = torch.randn(B, T, D)

    # 初始化 RMSNorm 层
    norm_layer = RMSNorm(dim=D, eps=1e-6)

    # 前向传播
    normalized_features = norm_layer(dummy_audio_features)

    print("=== 输入张量形状 ===")
    print(dummy_audio_features.shape)  # torch.Size([2, 4, 8])

    print("\n=== 输出张量形状 ===")
    print(normalized_features.shape)    # torch.Size([2, 4, 8])

    # 验证：检查未加权重前特征的 RMS 是否为 1
    raw_normed = norm_layer._norm(dummy_audio_features)
    calculated_rms = torch.sqrt(raw_normed.pow(2).mean(dim=-1))
    print("\n=== 归一化后的均方根（理论值应全为 1.0） ===")
    print(calculated_rms)