每日基础课

1. 生活类比： 想象你在给一段说话人录音调总音量。 传统 LayerNorm 相当于：先把整段声音的直流偏移（均值）移到 0 刻度，再按声音波动幅度（方差）缩放。 RMSNorm 相当于：直接看当前这一段声音的整体能量（均方根 RMS），能量大了就整体拧小，能量小了就整体放大。它不费力气去把中心挪到 0，因为特征本身的方向和相对能量已经足够表达信息了。

2. 核心数学公式： 对输入特征向量 $x = [x_1, x_2, ..., x_d]$： $$\text{RMS}(x) = \sqrt{\frac{1}{d} \sum_{i=1}^d x_i^2 + \epsilon}$$ $$y = \frac{x}{\text{RMS}(x)} \odot \gamma$$

符号逐个拆解： - $d$：特征维度大小（例如音频特征维度 512）。 - $x_i$：该维度上的第 $i$ 个具体数值。 - $\epsilon$（Epsilon）：一个极小的常数（如 $10^{-6}$），防止分母为 0 导致计算崩溃。 - $\text{RMS}(x)$：均方根值，代表这组特征的整体平均强度（物理上的“有效值”）。 - $\gamma$（Gamma）：可学习的缩放权重，形状与最后一维相同，让模型保留自主放大或缩小某些维度的能力。 - $\odot$：逐元素相乘。 注意：相比 LayerNorm，它完全去掉了减均值 $\mu$ 的步骤，也去掉了偏置项 $\beta$。

• 它解决什么问题？没有它会怎样？
• 解决的问题：深度 Transformer 在层数加深时，激活值的尺度会逐层爆炸或衰减，导致梯度消失或梯度爆炸。
• 没有它会怎样：深层网络无法稳定训练；同时，传统的 LayerNorm 每次都要算均值、减均值、算方差，在 GPU 上会产生多次访存开销，拖慢训练和推理速度。

实现思路： 1. 在特征维度（`dim=-1`）上先求平方，再求均值，加上 $\epsilon$ 后开根号（或直接用 `torch.rsqrt` 求倒数开根号，速度更快更稳定）。 2. 将原输入乘以该倒数，最后乘以可学习参数 $\gamma$。 3. 对齐官方 API 行为。

完整可运行代码：

import torch
import torch.nn as nn

class CustomRMSNorm(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 forward(self, x: torch.Tensor) -> torch.Tensor:
        # 1. 计算均方根倒数: 1 / sqrt(mean(x^2) + eps)
        # 推荐在 float32 下计算方差，防止半精度下溢/溢出
        input_dtype = x.dtype
        x_fp32 = x.to(torch.float32)
        variance = x_fp32.pow(2).mean(dim=-1, keepdim=True)
        rsqrt = torch.rsqrt(variance + self.eps)
        
        # 2. 归一化并转回原数据类型，最后应用可学习权重 weight
        x_normed = (x_fp32 * rsqrt).to(input_dtype)
        return x_normed * self.weight

# ================= 验证环节 =================
if __name__ == "__main__":
    torch.manual_seed(42)
    
    # 模拟音频/视频特征: (Batch=2, Frames=4, Hidden_Dim=8)
    batch_size, time_frames, dim = 2, 4, 8
    dummy_input = torch.randn(batch_size, time_frames, dim)
    
    # 实例化自定义 RMSNorm
    my_rmsnorm = CustomRMSNorm(dim=dim, eps=1e-6)
    custom_output = my_rmsnorm(dummy_input)
    
    # 官方标准公式验证 (数学基准)
    rms = torch.sqrt(dummy_input.pow(2).mean(-1, keepdim=True) + 1e-6)
    expected_output = (dummy_input / rms) * torch.ones(dim)
    
    # 校验
    max_diff = (custom_output - expected_output).abs().max().item()
    print("输入形状:", dummy_input.shape)
    print("输出形状:", custom_output.shape)
    print(f"与标准公式的最大绝对误差: {max_diff:.8e}")
    assert max_diff < 1e-6, "验证失败：数值误差过大！"
    print(" 验收通过：手写 RMSNorm 数值与逻辑完全正确。")