基于对数对齐比的训练时泛化诊断

# 训练时泛化诊断：基于对数对齐比
来源：https://arxiv.org/html/2605.28975
Ali Shehker ali\.shehper1@gmail\.com & Ashish Vaswani¹ ashish\.vaswani@gmail\.com

###### 摘要

我们研究了对数对齐比（LAR），一种衡量参数-激活对齐的指标，源自参数化理论。我们将其重新表述为权重谱 p（矩阵归一化平方奇异值的分布）与激活谱 q（输入在其奇异方向上归一化平方投影的分布）之间的重叠。我们发现，在两种不同的设定下，解嵌入层的 LAR 通过捕捉训练过程中 p 和 q 的扩散程度来追踪记忆与泛化之间的转变。在 grokking 现象中，LAR 预测了学习函数的有效维度：k ≈ n^{2(1 - LAR)}，其中 n 是矩阵的输入维度。在 3B 参数语言模型预训练中，其与无过拟合基线的偏差追踪了泛化差距，且其下降速率随过拟合临近而加快。LAR 可通过前向传播中已有的量进行计算，计算开销可忽略不计，且无需保留验证数据。

## 1 引言

大规模预训练推动了神经网络能力中最显著的近期进展（Kaplan et al., 2020；Hoffmann et al., 2022），随着模型和数据集规模的持续增长，每次训练运行的成本也随之增长。这使得对训练动态的研究以及设计廉价的诊断工具来监控它们变得异常有价值：一个能够仅从训练时量（无需昂贵的保留评估）就标记出过拟合、不稳定或饱和等问题的指标，可以为实践者节省计算和时间。

另一条关于神经网络*参数化*的互补工作路线，规定了初始化、参数乘数以及学习率应如何随宽度和深度缩放，为从小模型向大模型转移超参数提供了原则性方法。Yang et al.（2022）提出了最大更新参数化（μP），该参数化假设读出层的更新与传入激活完全对齐，从而实现超参数转移。Everett et al.（2024）证明了这一对齐假设比必要的更强，并通过放宽这一假设，获得了一个标准参数化配方，在多达 26.8B 参数的规模上、成千上万次实验中始终优于 μP。作为该分析的一部分，他们引入了*对数对齐比*（LAR）作为训练过程中该层参数矩阵与其输入激活对齐程度的经验度量。

虽然 Everett et al.（2024）将 LAR 用作参数-激活对齐的代理指标，但我们在本文中发现该指标具有第二个先前未被探索的含义：即它衡量了该层的激活和权重视谱是分散还是集中的程度。更精确地说，我们证明了 LAR 衡量了权重视谱 p 与激活谱 q 之间的重叠：当两个谱都集中在同一组小方向上时，LAR 很大；当其中任何一个分布分散或两者不对齐时，LAR 很小。

这一见解使 LAR 成为衡量*泛化*的自然候选指标——即模型泛化到训练数据之外而非过拟合的程度。经验上，我们观察到良好泛化的网络倾向于将其计算集中在少数方向上，而过拟合的网络则倾向于将其分散在多个方向上。大多数先前研究这一行为的指标要么只覆盖权重视谱，要么只覆盖激活谱，并且需要完整的权重矩阵 SVD 或激活的 PCA，计算成本高昂。相比之下，LAR 同时关注权重和激活，并且可以从模型前向传播中已有的量计算，计算开销可忽略不计。

我们研究了两种不同设定下模型解嵌入矩阵的 LAR 与泛化之间的关系。首先，在表现出 grokking 现象的小型算法任务中（Power et al., 2022），我们发现优化过程将 p 和 q 集中到一组共享的方向上，这导致 LAR 增长。当方向数 k 变小时，可能会发生 grokking，具体取决于任务难度和训练数据集大小等因素。LAR 预测了学习函数的维数：k ≈ n^{2(1 - LAR)}，我们将其与 k_95（解释解嵌入输入激活方差 95% 的主成分数量）进行比较。我们在所有实验中找到了最终检查点的紧密一致性，包括 12 个 grokking 任务以及对训练集比例的扫描。

其次，在 3B 参数语言模型预训练中，我们发现当模型泛化时，解嵌入层的 LAR 缓慢下降并趋于稳定，但当过拟合接近时则急剧下降。LAR 与无过拟合基线的偏差也紧密追踪了泛化差距。在过拟合之前 LAR 加速下降的特性可能使得无需评估的方法能够训练泛化模型——我们将其留给未来研究探讨。

在 grokking 和大规模设定中，我们发现了一致的图景：较低的 LAR 对应更分散的权重和激活分布，这与记忆重合；而较高的 LAR 对应更集中的分布和泛化。在每种情况下，我们观察到训练过程中 p 与 q 之间的相关性始终保持较高，因此 LAR 有效追踪了权重和激活分布偏离均匀分布的程度。

#### 贡献。

我们的主要贡献如下。

- **LAR 的分布重述**。我们证明了 LAR 可以写为权重视谱 p 与激活谱 q 之间的对数重叠 ∑ᵢ pᵢ qᵢ，同时给出了等价的协方差和相关性形式，使得该指标的行为可以用基本统计量来解释。
- **通过 LAR 获得有效维度**。我们证明当 W 将其能量集中在 k 个奇异方向上且激活沿着这些方向时，LAR = 1 - ½ logₙ k。我们通过在 12 个算法任务上将 LAR 预测的维度与 k_95 进行比较，在 grokking 设定中经验验证了这一点。
- **大规模下的过拟合诊断 LAR**。在使用 Adam 和 Muon 训练的 3B 参数 Gemma-2 模型上，我们证明了解嵌入层的 LAR 在非过拟合状态下趋于稳定，在过拟合下急剧下降，并且其与无过拟合基线的偏差紧密追踪了泛化差距。
- **LAR 作为 p 和 q 扩散程度的代理**。我们证明在实践中，优化使 p 和 q 高度相关，从而使 LAR 成为 p 和 q 方差的代理。在预训练实验中，解嵌入矩阵的最大奇异值远大于其余部分，并且是 p 扩散的主要驱动因素。这使得来自最大奇异方向的贡献 L₁ = 1 + ½ logₙ(p₁ q₁) 成为过拟合前过拟合迹象的替代指标。

## 2 相关工作

#### 泛化的谱度量。

多条工作线已将权重矩阵的谱特性与泛化联系起来。Bartlett et al.（2017）和 Neyshabur et al.（2018）推导了用谱范数乘积表示的基于边界的泛化界，而 Arora et al.（2018）利用训练网络的噪声稳定性给出了更紧的基于压缩的界。Jiang et al.（2020）后来在大规模实证研究中比较了数十种此类度量。经验上，Martin and Mahoney（2021）提出权重视谱密度在训练过程中发展出重尾，并将其用作无数据的泛化和早停信号。Yunis et al.（2024）表明权重的归一化有效秩在训练过程中广泛下降，并且低秩权重的出现与 grokking 中从记忆到泛化的转变一致。

相比之下，LAR 依赖于权重和激活两者，并且可以从矩阵范数本身计算。它允许一个显式的有效维度公式，在 grokking 实验中得到了验证，并在 3B 规模下追踪了泛化差距。此外，在 3B 预训练实验中（附录 D.1），有效秩在 Adam 下只能微弱地区分过拟合与泛化，而解嵌入层的 LAR 在 Adam 和 Muon 下都能清晰地区分二者。

#### 神经网络参数化。

另一条互补的工作线研究神经网络*参数化*，作为从小模型向大模型转移超参数的原则性方法。最大更新参数化（μP）（Yang and Hu, 2021；Yang et al., 2022）实现了学习率及相关超参数的宽度转移。后续工作将这一框架扩展到新的缩放轴：Dey et al.（2025）引入了 CompleteP，实现了跨深度的超参数转移，同时确保每一层的非懒惰学习；Mlodozeniec et al.（2025）进一步将框架扩展到覆盖每模块超参数、批次大小和 Token 视野。Everett et al.（2024）重新审视了最大更新参数化背后的对齐假设，并证明了放宽这一假设可以得到在高达 26.8B 参数规模下优于 μP 的逐层学习率处方。LAR 是在 Everett et al.（2024）中作为该分析的一部分被引入的，作为参数-激活对齐的经验度量。我们的工作将 LAR 重新用于泛化诊断：这个在参数化理论中用于探测训练期间对齐假设是否成立的同一量，独立于其在参数化理论中的用途，也追踪了过拟合动态。

#### Grokking。

Grokking 由 Power et al.（2022）首次发现，指的是在训练精度饱和后很久才观察到的从记忆到泛化的延迟转变。多条工作线提供了解释。Barak et al.（2022）研究了隐藏进度度量（即预测收敛时间的训练状态标量函数），并使用稀疏奇偶学习作为测试平台，其中 grokking 式的相变发生并且与其隐藏进度框架一致。Nanda et al.（2023）以机械方式逆向工程了模加法的泛化解，并定义了其出现的进度度量；Liu et al.（2023）将 grokking 归因于训练和测试损失景观几何之间的不匹配；Davies et al.（2022）将 grokking 和双下降统一为同一动态的实例；Varma et al.（2023）将 grokking 解释为记忆电路与更参数高效的泛化电路之间的竞争。在谱动力学文献中，Yunis et al.（2024）观察到 grokking 过程中验证损失下降与权重视谱有效秩的急剧下降同时发生。

我们并非提出新理论，而是将 grokking 作为测试平台，研究解嵌入层 LAR 在记忆到泛化转变过程中的行为。我们还发现 LAR 预测了学习函数的有效维度 k ≈ n^{2(1 - LAR)}，并在 12 个任务上进行了经验验证。

## 3 对数对齐比

### 3.1 定义

设 W ∈ ℝ^{m×n} 为矩阵，X ∈ ℝ^{b×n} 为一批输入向量，WX ∈ ℝ^{b×m} 为相应的输出批次。该矩阵的对数对齐比（LAR）由 Everett et al.（2024）定义为：

LAR = logₙ (||WX||_{RMS} / (||W||_{RMS} · ||X||_{RMS})) (1)

其中 ||·||_{RMS} 是矩阵的 RMS 范数。

设 W = ∑_{i=1}^{r} u_i s_i v_i^⊤ 为 W 的 SVD，其中 r = rank(W)。每个输入向量 x ∈ ℝⁿ 都可以在右奇异向量基中展开为 x = ∑_{i=1}^{n} x_i v_i，其中 x_i = v_i^⊤ x，而 v_{r+1}, ..., v_n 在 n > r 时完成 ℝⁿ 的一组标准正交基。

我们将*权重分布* p 和*激活分布* q 分别定义为 W 的归一化平方奇异值以及 x 在其奇异方向上的归一化平方投影：

p_i = s_i² / ∑_{j=1}^{r} s_j²，  q_i = ∑_{x ∈ B} (v_i^⊤ x)² / ∑_{x ∈ B} ||x||² (2)

其中 ||·|| 表示向量的 L₂ 范数。p 是 r 个元素上的分布，捕捉 W 如何将其能量分布在奇异方向之间；q 是 n 个元素上的分布，捕捉该批激活沿着这些方向和零空间的方差如何分布。¹

直接计算可得：

LAR = 1 + ½ logₙ (∑_{i=1}^{r} p_i q_i) (3)

为了看出这一点，写出：

||W||_{RMS}² = (1/(nm)) ∑_{j=1}^{r} s_j²，
||X||_{RMS}² = (1/(|B| n)) ∑_{x∈B} ||x||²，
||Wx||_{RMS}² = (1/(|B| m)) ∑_{x∈B} ∑_{i=1}^{r} s_i² x_i²。

平方比变为：

||Wx||_{RMS}² / (||W||_{RMS}² ||x||_{RMS}²) = [ (1/(|B| m)) ∑_{x∈B} ∑_{i} s_i² x_i² ] / [ (1/(nm)) ∑_{j} s_j² · (1/(|B| n)) ∑_{x∈B} ||x||² ] = n² · [ ∑_{i} s_i² ∑_{x∈B} x_i² ] / [ ∑_{j} s_j² · ∑_{x∈B} ||x||² ]