每日基础课

【生活类比】 想象你在调试一段录音的音量： - LayerNorm：先检测这行音频有没有整体静电杂音偏置（计算均值并减去），再看声音整体响度（计算方差），最后把音量拉到标准大小。 - RMSNorm：发现现在的神经网络内部特征其实很少有严重“偏置”，每次测均值、减均值完全是多余的计算开销。所以它直接用特征的整体能量（均方根 RMS）做统一缩放。动作少了一半，音量调得一样稳。

【必要公式与符号解释】 RMSNorm 的计算公式只有两步： $$y = \frac{x}{\text{RMS}(x)} \odot \gamma$$ 其中： $$\text{RMS}(x) = \sqrt{\frac{1}{d} \sum_{i=1}^d x_i^2 + \epsilon}$$

逐个符号拆解： - $x$：输入特征向量（例如语音帧或 Token 的隐藏特征向量，维度为 $d$）。 - $d$：特征维度大小（如 512、4096）。 - $\text{RMS}(x)$：均方根（Root Mean Square），衡量该向量的平均能量幅度。 - $\epsilon$：极小的常数（如 $1e-6$），防止分母除以 0。 - $\gamma$：可学习的缩放参数（权重向量，形状与 $d$ 相同）。 - $\odot$：逐元素相乘。 - 注意：**没有偏置项 $\beta$（bias），也没有减去均值 $\mu$ 的计算**。

【它解决什么问题？没有它会怎样？】 - 解决的问题：在深层神经网络（如 32 层、64 层的 Transformer）中，前向传播会让特征数值越来越大或越来越小，反向传播时导致梯度爆炸或梯度消失。RMSNorm 强行把每一层输出的能量尺度拉回可控范围。 - 没有它会怎样：网络层数一旦超过十几层，训练就会瞬间发散（Loss 变 NaN）或者几乎无法收敛。

【实现思路】 1. 继承 `nn.Module`，初始化可学习参数 `weight`（全 1 初始化）和微小值 `eps`。 2. 在 `forward` 中，沿最后一个维度计算均方值：`x.pow(2).mean(-1, keepdim=True)`。 3. 开根号并加上 `eps`，得到 `rms`。 4. 用 `x / rms` 乘以 `self.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
        # RMSNorm 只有 scale 参数 (gamma)，没有 bias
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 1. 沿特征维度 (最后一维) 计算均方值 (Mean Square)
        # 为防止低精度下平方和溢出，通常保证在 float32 运算
        orig_dtype = x.dtype
        x_fp32 = x.to(torch.float32)
        variance = x_fp32.pow(2).mean(dim=-1, keepdim=True)
        
        # 2. 计算 RMS 并做归一化
        x_normed = x_fp32 * torch.rsqrt(variance + self.eps)
        
        # 3. 转回原始精度并乘以可学习的 weight
        return (x_normed.to(orig_dtype)) * self.weight

# ================= 验收测试 =================
if __name__ == "__main__":
    torch.manual_seed(42)

    # 模拟说话人日志中的音频特征: Batch=2, 序列长度=100帧, 特征维度=256
    batch_size, seq_len, hidden_dim = 2, 100, 256
    audio_frames = torch.randn(batch_size, seq_len, hidden_dim) * 5.0 + 3.0  # 放大能量并加入偏移

    # 实例化模块
    rms_norm = RMSNorm(dim=hidden_dim)

    # 前向传播
    output = rms_norm(audio_frames)

    # 打印形状验证
    print("输入形状:", audio_frames.shape)
    print("输出形状:", output.shape)
    assert audio_frames.shape == output.shape, "形状校验失败！"

    # 验证输出能量 (未乘以 weight 前 RMS 应该接近 1.0)
    rms_val = torch.sqrt(output.pow(2).mean(dim=-1)).mean().item()
    print(f"归一化后的平均 RMS 值: {rms_val:.4f} (预期接近 1.0)")

    # 梯度反向传播验证
    loss = output.sum()
    loss.backward()
    print("权重 weight 的梯度形状:", rms_norm.weight.grad.shape)
    assert rms_norm.weight.grad is not None, "梯度反向传播失败！"
    print(">> RMSNorm 手写验证全部通过！")