SGD in Multiclass Logistic Regression: Sequential Learning and Scaling Laws¶
会议: NeurIPS 2026
arXiv: 2609.07868
OpenReview: https://openreview.net/forum?id=lF5Y8IbCiq
领域: 优化/理论
关键词: 多分类逻辑回归, 神经标度律, 序列学习, 随机梯度下降, 算力最优分配
一句话总结¶
针对长尾幂律先验的大规模类别高维高斯混合模型,本文解析证明了多分类逻辑回归在交叉熵损失下的序列学习机制与三阶段标度律,并给出PCA降维下的算力最优缩放法则。
研究背景与动机¶
近年来以大语言模型为代表的经验标度律(Scaling Laws)展现出惊人的预测能力:随着算力、参数量与训练样本的增长,测试损失在多个数量级内遵循严格的幂律衰减。为了从理论上厘清标度律的数学成因,现有学术界主要聚焦于特征回归问题(如随机特征或带有幂律协方差谱的线性回归、加性浅层神经网络),推导出均方误差(MSE)关于模型容量与优化时间的幂律分解。然而,这一经典分析范式与现代自回归语言模型预训练之间存在根本性的结构脱节:真实语言模型的核心是离散词表上的 Next-Token 预测,优化目标是多分类交叉熵(Cross-Entropy)损失,且词表先验呈现典型的极端不平衡长尾分布(Zipf 定律)。
将回归分析直接套用到多分类交叉熵优化上面临三重严苛挑战:首先,类别数 \(K\) 趋向发散时,Softmax 函数引入了全局归一化耦合与强非线性饱和效应,破坏了均方误差下的解析解可分离性;其次,在回归问题中作为经典降维手段的各向同性高斯随机投影(Random Sketch),在多分类几何下由于 Johnson-Lindenstrauss 引理保持了类别均值间的两两距离与信噪比,无法在分类交叉熵中诱导出现实意义上的模型容量瓶颈;最后,非均匀类别分布导致不同类别的梯度信号强度跨越多个数量级,优化动态呈现出强烈的异步与跨类干涉特征。
面对上述理论空白,本文选择了一个极具代表性且解析可控的理论模型:类别数极大(\(1 \ll K \le d\))的高维高斯混合模型(GMM),其类别先验遵循幂律分布 \(\pi_i \asymp i^{-\alpha}\)(\(\alpha > 0\)),并通过梯度流与在线随机梯度下降(Online SGD)最小化交叉熵损失。核心 idea:多分类逻辑回归的梯度动力学本质上遵循由高频到低频的显式「逐类序列学习」机制,在幂律先验下交叉熵风险展现出特征三阶段标度律,且通过主成分分析(PCA)投影截断主导均值可严格诱导容量瓶颈,从而推导得出平衡模型容量与优化步数的算力最优(Compute-Optimal)标度律。
方法详解¶
整体框架¶
本文构建了类别数发散的高维正交均值高斯混合模型(GMM),输入样本满足 \(y \sim \text{Categorical}(\pi_1, \dots, \pi_K)\),在给定标签 \(y=i\) 条件下,\(\bm{x} \mid y=i \sim \mathcal{N}(\bm{\mu}_i, \sigma^2 \bm{I}_d)\),其中均值两两正交且满足单位范数 \(\langle \bm{\mu}_i, \bm{\mu}_j \rangle = \mathbb{I}\{i=j\}\)。模型采用无偏置项的多分类逻辑回归权重矩阵 \(\bm{W} = [\bm{w}_1, \dots, \bm{w}_K] \in \mathbb{R}^{d \times K}\),通过 Softmax 预测后验概率 \(p_i(\bm{x}) = \exp(\langle \bm{w}_i, \bm{x}\rangle) / \sum_{k=1}^K \exp(\langle \bm{w}_k, \bm{x}\rangle)\) 并以交叉熵损失驱动。
为了刻画复杂的非线性动力学,作者引入了连接权重与正交均值的充分统计量(Summary Statistics)矩阵 \(\theta_{ij}(t) := \langle \bm{w}_i(t), \bm{\mu}_j \rangle\),建立起连续时间动力学系统;随后引入数据驱动的经验 PCA 投影矩阵 \(\widehat{\bm{S}} \in \mathbb{R}^{m \times d}\) 实施容量瓶颈控制;最后通过凸性、光滑性分析与鞅差序列集中不等式,严格将总体梯度流连续轨迹离散化并转移至离散样本在线 SGD。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["高维长尾高斯混合输入<br/>y ~ Categorical(π), x|y=i ~ N(μ_i, σ²I)"] --> B["充分统计量动力学与逐类序列学习<br/>定义 θ_ij=⟨w_i, μ_j⟩ 跟踪解耦演化"]
B --> C["交叉熵风险的三阶段标度律<br/>初期平台 → 幂律衰减 → 隐式偏置收敛"]
C --> D["PCA容量控制与算力最优分配<br/>投影至Top-m主成分平衡容量误差与优化误差"]
D --> E["在线SGD离散化与鞅差收敛保证<br/>有效优化时间 t_n=η·n,继承总体标度律"]
关键设计¶
1. 充分统计量动力学与逐类序列学习:解耦高频到低频的类别捕获
在多分类逻辑回归中,由于 Softmax 具有移位不变性,零初始化 \(\bm{W}(0)=\bm{0}\) 保证了对任意类别 \(j\) 恒有守恒律 \(\sum_{i=1}^K \theta_{ij}(t) = 0\)。在理想无噪条件(\(\sigma^2 = 0\))下,跨类别统计量精确对称,满足 \(\theta_{ij}(t) = -\frac{1}{K-1} \theta_{jj}(t)\)(对所有 \(i \ne j\))。这一结构性对称将原高维动力学完全降维为 \(K\) 个关于真实类预测概率 \(p_{i|i}(t) := p_i(\bm{\mu}_i)\) 的一维自治微分方程: $\(\dot{p}_{i|i}(t) = \frac{K}{K-1} \pi_i p_{i|i}(t) (1 - p_{i|i}(t))^2, \quad p_{i|i}(0) = \frac{1}{K}.\)$ 由此解析求解得到类别 \(i\) 越过置信度阈值 \(c \in (1/K, 1)\) 的击中时间(Hitting Time): $\(t_i(c) = \frac{K-1}{K} \frac{1}{\pi_i} \left[ \log \frac{c}{1-c} + \frac{1}{1-c} + \log(K-1) - \frac{K}{K-1} \right] \asymp \frac{\log K}{\pi_i}.\)$ 当先验遵循幂律分布 \(\pi_i = i^{-\alpha}/Z_{K,\alpha}\) 时,各类别按照出现频次从大到小的顺序被梯度流逐一“激活”,学习时间 \(t_i \asymp Z_{K,\alpha} i^\alpha \log K\) 显式依赖于频率倒数。当引入类内方差 \(\sigma^2 > 0\) 时,高斯分布下的 Stein 引理在微分方程中引入了非零交互项 \(\sigma^2 \mathbb{E}_{\bm{x}} [p_i(\bm{x})(\mathbb{I}\{i=\ell\} - p_\ell(\bm{x}))] \theta_{\ell j}\);作者通过局部泰勒展开与梯度流权重范数上界(\(\max_i \|\bm{w}_i(t)\| \lesssim t\)),严谨证明了在噪声受限的时间视界 \(T \asymp 1/\sigma\)(即信噪比倒数)内,序列学习的时间尺度与拓扑轨迹依然受到精确控制。
2. 交叉熵风险的三阶段标度律:宏观幂律衰减的涌现
整体交叉熵损失(风险)由所有类别的瞬时误差加权叠加而成,即 \(\mathcal{R}(t) = -\sum_{i=1}^K \pi_i \log p_{i|i}(t)\)。作者根据类别是否被成功学习将类别集合动态划分为头部类集合(\(p_{i|i} \ge 1/2\))与尾部未学习类集合(\(p_{i|i} < 1/2\))。由于尾部类误差保持在 \(\log K\) 水平,而每个时刻被激活跨入头部的类别数量呈连续积分态,宏观风险随着时间演化清晰展现出三个阶段: - 初期平台阶段(Early Phase):在 \(0 \le t \lesssim Z_{K,\alpha} \log K\) 期间,即使最频繁的第 1 类也尚未收敛,\(p_{i|i} \approx 1/K\),损失保持在 \(\mathcal{R}(t) \asymp \log K\)。 - 中间幂律衰减阶段(Intermediate Phase):当 \(Z_{K,\alpha} \log K \lesssim t \lesssim Z_{K,\alpha} K^\alpha \log K\) 时,头部类别数逐渐累积。当 \(\alpha > 1\)(长尾分布的概率在 \(K \to \infty\) 时可和)时,尾部类别的积分主导了风险曲线,产生光滑的幂律衰减: $\(\mathcal{R}(t) \asymp (\log K)^{2 - 1/\alpha} t^{-(1 - 1/\alpha)}.\)$ 若 \(0 < \alpha < 1\),衰减形式变为 \((\log K)(1 - (t/(K\log K))^{1/\alpha - 1})\);若 \(\alpha = 1\),则呈对数缩放 \(\log(K (\log K)^2 / t)\)。 - 局部收敛阶段(Convergence Phase):当所有类别均被学习进入头部集合(\(t \gtrsim Z_{K,\alpha} K^\alpha \log K\))后,全局逻辑回归进入隐式偏差机制,损失呈现由类别数决定的局部收敛速率 \(\mathcal{R}(t) \asymp K/t\)。
3. PCA 容量控制与算力最优分配:破除高斯投影失效困境
在经典线性回归中,高斯随机投影能自然引入参数容量瓶颈。然而在多分类逻辑回归中,由于高维空间正交均值在经过随机投影矩阵 \(\bm{S} \in \mathbb{R}^{m \times d}\) 后,只要 \(m \gtrsim \log K\),根据 Johnson-Lindenstrauss 引理,均值间两两距离和各向同性噪声方差均保持不变,完全无法限制分类交叉熵的表达能力。为此,作者创新性地采用数据二阶矩矩阵 \(\bm{G} = \mathbb{E}[\bm{x}\bm{x}^\top] = \sum_{i=1}^K \pi_i \bm{\mu}_i \bm{\mu}_i^\top + \sigma^2 \bm{I}_d\) 的主成分分析(PCA)投影。保留前 \(m\) 个主特征方向将正交均值 \(\bm{\mu}_{m+1}, \dots, \bm{\mu}_K\) 坍缩至原点,迫使其输出恒为初始先验 \(1/K\)。在实际中,基于 \(N\) 个未标注样本构建经验二阶矩矩阵 \(\widehat{\bm{G}}\),利用 Davis-Kahan \(\sin\Theta\) 定理控制经验子空间距离 \(\varepsilon_{\text{PCA}} = \|\widehat{\bm{P}}_m - \bm{P}_m\|_{\text{op}}\)。在固定特征维度 \(m \le K/2\) 下,风险精确分解为逼近误差与优化误差: $\(\mathcal{R}_{\widehat{\bm{S}}}(t) \asymp \underbrace{m^{-(\alpha-1)} \log K}_{\mathcal{R}_{\text{approx}}} + \underbrace{(\log K)^{2 - 1/\alpha} t^{-(1 - 1/\alpha)}}_{\mathcal{R}_{\text{opt}}}.\)$ 在给定总计算量预算 \(C \asymp m K t\)(每步评估 \(K\) 个 Logit 需要 \(O(mK)\) FLOPs)的约束下,对 \(m\) 与 \(t\) 求解极值,推导出最优参数维度 \(m^*\) 与训练时间 \(t^*\) 的显式封闭解: $\(m^* \asymp \left(\frac{C}{K \log K}\right)^{\frac{1}{\alpha+1}}, \quad t^* \asymp (\log K)^{\frac{1}{\alpha+1}} \left(\frac{C}{K}\right)^{\frac{\alpha}{\alpha+1}}.\)$ 由此得到的全局算力最优标度律下限为: $\(\mathcal{R}_{\widehat{\bm{S}}}(C) \asymp (\log K)^{\frac{2\alpha}{\alpha+1}} \left(\frac{K}{C}\right)^{\frac{\alpha-1}{\alpha+1}}.\)$
4. 在线 SGD 离散化与鞅差收敛保证:从连续流到离散优化的严格过渡
为了将连续时间结论拓展到离散样本流优化,作者分析了步长为 \(\eta > 0\) 的在线 SGD。总体风险 \(\mathcal{R}(\bm{W})\) 具备凸性与全局 \(L\)-光滑性(\(L \le \sigma^2 + \pi_1 \asymp 1\))。首先,作者通过梯度的单调性建立总体梯度下降(GD)逼近连续梯度流(GF)的离散截断误差,证明在有效时间步 \(t_n = \eta n \le T\) 内有 \(\|\widetilde{\bm{W}}^{(n)} - \bm{W}(t_n)\|_F \le \frac{1}{2} L \eta T\);其次,定义随机梯度误差矩阵 \(\bm{D}_{n+1} = \nabla \ell(\bm{W}^{(n)}; \bm{u}^{(n+1)}, y^{(n+1)}) - \nabla \mathcal{R}(\bm{W}^{(n)})\),证明其条件方差受限于样本二阶矩 \(\mathbb{E}[\|\bm{D}_{n+1}\|_F^2 \mid \mathcal{F}_n] \le 2 M\),其中未降维模型 \(M = 1 + \sigma^2 d\),而降维后 \(M = 1 + \sigma^2 m\)。利用 Doob 停时理论与马尔可夫不等式,建立了误差在区间内的均匀高概率集中界: $\(\mathbb{P}\left(\max_{0 \le n \le N} \|\bm{W}^{(n)} - \widetilde{\bm{W}}^{(n)}\|_F \ge \eta \sqrt{\frac{2 M N}{\delta}}\right) \le \delta.\)$ 在学习率选取满足 \(\eta \lesssim \min\{1, \frac{\mathcal{R}(\bm{W}(T))}{T(1+T)}, \frac{\delta \mathcal{R}(\bm{W}(T))^2}{M T}\}\) 条件下,在线 SGD 的离散轨迹在高概率下严格继承总体梯度流的连续三阶段标度律。
损失函数 / 训练策略¶
模型针对多分类交叉熵损失 \(\mathcal{L}(\bm{W}; \bm{x}, y) = -\sum_{i=1}^K \mathbb{I}\{y=i\} \log p_i(\bm{x})\) 展开优化。在在线 SGD 算法中,初始化参数 \(\bm{W}^{(0)} = \bm{0}_{d \times K}\),在第 \(n\) 步独立采样 \((\bm{x}^{(n)}, y^{(n)})\),并执行单步更新: $\(\bm{W}^{(n+1)} \leftarrow \bm{W}^{(n)} - \eta \bm{x}^{(n)} (p(\bm{W}^{(n)}; \bm{x}^{(n)}) - \bm{e}_{y^{(n)}})^\top.\)$ 对于降维情形,输入 \(\bm{x}^{(n)}\) 替换为投影特征 \(\widehat{\bm{S}} \bm{x}^{(n)} \in \mathbb{R}^m\),参数矩阵相应变为 \(\bm{W} \in \mathbb{R}^{m \times K}\)。在小噪声条件 \(\sigma = o(1/\log K)\) 下,作者进一步刻画了最优贝叶斯分类器的不可约风险(Irreducible Risk)极限: $\(\mathcal{R}_{K,\bm{\pi}}^*(\sigma) := \inf_{\bm{W} \in \mathbb{R}^{d \times K}} \mathcal{R}(\bm{W}) \asymp \sigma e^{-\frac{1}{4\sigma^2}} \cdot \begin{cases} K, & 0 < \alpha < 1, \\ K/\log K, & \alpha = 1, \\ K^{2-\alpha}, & 1 < \alpha < 2, \\ (\log K)^2, & \alpha = 2, \\ 1, & \alpha > 2. \end{cases}\)$ 作者构造了模拟贝叶斯决策面的无偏置权重配置,证明在严格凸空间 \(\mathcal{W} = \{\bm{W}: \sum_i \bm{w}_i = \bm{0}\}\) 内,梯度流以 Frobenius 范数全局收敛至该唯一最小化子 \(\bm{W}^*\)。
实验关键数据¶
主实验¶
论文通过对总体常微分方程(ODE)进行自适应步长的欧拉数值离散(期望由蒙特卡洛积分逼近),系统验证了不同先验衰减指数 \(\alpha\)、特征保留维度 \(m\) 及噪声水平 \(\sigma\) 下的标度律理论预测。设置类别总数 \(K = 2048\),环境维度 \(d \ge K\)。
下表总结了在理论有效时间视界内,不同长尾先验指数 \(\alpha\) 对多分类逻辑回归经验衰减斜率(理论预测对比数值拟合)的影响:
| 先验指数 \(\alpha\) | 风险动态阶段 | 理论衰减速率 \(\mathcal{R}(t)\) | 观测经验斜率(Log-Log 拟合) | 吻合状态 |
|---|---|---|---|---|
| \(\alpha = 0.5\) (\(0 < \alpha < 1\)) | 中间学习阶段 | 缓慢次对数衰减 \((\log K)(1 - (t/t_K)^{-1})\) | 平台缓慢下倾,无幂律 | 严格吻合 |
| \(\alpha = 1.0\) (\(\alpha = 1\)) | 中间学习阶段 | 对数衰减 \(\log(K (\log K)^2 / t)\) | 曲线微凹对数衰减 | 严格吻合 |
| \(\alpha = 1.5\) (\(\alpha > 1\)) | 中间学习阶段 | \(t^{-(1 - 1/1.5)} = t^{-0.333}\) | \(t^{-0.33}\) | 严格吻合 |
| \(\alpha = 2.0\) (\(\alpha > 1\)) | 中间学习阶段 | \(t^{-(1 - 1/2.0)} = t^{-0.500}\) | \(t^{-0.50}\) | 严格吻合 |
| \(\alpha = 3.0\) (\(\alpha > 1\)) | 中间学习阶段 | \(t^{-(1 - 1/3.0)} = t^{-0.667}\) | \(t^{-0.67}\) | 严格吻合 |
| 全类别学习后(\(t \ge t_K\)) | 局部收敛阶段 | 隐式偏置主导 \(K / t\) | \(t^{-1.00}\) | 严格吻合 |
消融实验¶
针对 PCA 降维机制与类内高斯噪声 \(\sigma\) 对容量瓶颈和标度律边界的影响,论文在 \(K = 2048, \alpha = 2.0\) 条件下进行系统参数敏感性分析:
| 实验配置 | 噪声水平 \(\sigma\) | 降维维度 \(m\) | 容量瓶颈误差 \(\mathcal{R}_{\text{approx}}\) | 时间衰减指数 \(\beta\) | 有效标度视界 \(T\) |
|---|---|---|---|---|---|
| 无噪全维基线 | \(\sigma = 0\) | \(m = K\) (无降维) | 0 (无瓶颈) | 0.50 (\(t^{-0.5}\)) | \(T \to \infty\) |
| 弱噪全维基线 | \(\sigma = 0.1\) | \(m = K\) (无降维) | 0 (无瓶颈) | 0.50 (\(t^{-0.5}\)) | \(T \approx 10\) |
| 中噪全维配置 | \(\sigma = 0.16\) | \(m = K\) (无降维) | 0 (无瓶颈) | 0.50 (\(t^{-0.5}\)) | \(T \approx 6.25\) |
| PCA 截断 \(m=64\) | \(\sigma = 0.1\) | \(m = 64\) | \(\asymp 64^{-1.0} \log K\) | 0.50 (\(t^{-0.5}\)) | \(T \approx 10\) |
| PCA 截断 \(m=128\) | \(\sigma = 0.1\) | \(m = 128\) | \(\asymp 128^{-1.0} \log K\) | 0.50 (\(t^{-0.5}\)) | \(T \approx 10\) |
| PCA 截断 \(m=256\) | \(\sigma = 0.1\) | \(m = 256\) | \(\asymp 256^{-1.0} \log K\) | 0.50 (\(t^{-0.5}\)) | \(T \approx 10\) |
| 高斯随机投影 (对照) | \(\sigma = 0.1\) | \(m = 128\) | 无法形成瓶颈 (JL 保距) | 趋近全维行为 | 不符合容量幂律 |
关键发现¶
- 衰减指数与长尾厚度直接绑定:在 \(\alpha > 1\) 条件下,交叉熵优化误差衰减指数精确等于 \(1 - 1/\alpha\)。先验分布越陡峭(\(\alpha\) 越大),高频类别占比越高,前期幂律损失下降速度越快;而一旦所有类别被学完,损失曲线自动平滑换挡至与隐式偏置理论一致的 \(K/t\) 速率。
- 随机投影在多分类下容量失效:高斯随机特征投影在多分类交叉熵下无法建立与宽度成幂律反比的容量瓶颈,而基于输入协方差的 PCA 截断能够精确将低频类坍缩至原点,建立起严格的 \(m^{-(\alpha-1)} \log K\) 容量阶数。
- 噪声对时间视界的截断效应:噪声方差 \(\sigma^2 > 0\) 会在充分统计量动力学中引入不可忽视的类间交叉干涉项,使有效标度律视界受限于 \(T \asymp 1/\sigma\)。当训练时间超出该视界后,损失曲线由于高斯方差的不可约残差减速并脱离理论无噪曲线。
亮点与洞察¶
- 补齐分类交叉熵标度律的理论拼图:以往标度律数学理论几乎被均方误差(MSE)垄断,本文首次在具有离散发散类别空间和长尾先验的多分类交叉熵损失下,解析推导出具备现实对应性的幂律学习曲线。
- 揭示长尾分类下的「逐类激活」微观机制:宏观光滑连续的损失标度律并非来自均匀的整体压缩,而是由于类别频率的指数差异,驱动梯度流从高频向低频发生确定性的阶梯式序列学习,将加性模型的原子任务学习假说拓展至真实多分类分布。
- 明确多分类下 PCA 与随机投影的本质差异:指出在分类问题中仅靠随机投影无法诱导表征瓶颈,确立了基于特征谱截断实现多分类模型容量限制的正确数学构造。
局限与展望¶
- 摄动分析的时间界限限制:当前对于有噪高斯混合模型(\(\sigma^2 > 0\))的控制依赖局部泰勒展开与范数界,所得理论保证仅严格覆盖至 \(T \asymp 1/\sigma\) 之前,在此之后的晚期非线性噪声演化行为尚未闭合。
- 正交均值与各向同性方差假设:理论假定各个类别的先验均值两两严格正交且类内协方差矩阵为标准球形高斯分布,这与真实自然语言词嵌入中普遍存在的层次化语义聚类结构存在差距。
- 单遍在线 SGD 与固定优化器:当前随机优化证明严格局限于一次遍历(One-pass)的在线 SGD,未探讨多轮 Epoch 数据重用(Data Reuse)对算力最优分配比的影响,也未涵盖 Momentum 或 Adam 等自适应优化算法。
相关工作与启发¶
- vs 线性回归标度律 (Lin et al., 2024; Paquette et al., 2024):以往工作在具有幂律谱 \(\lambda_i \asymp i^{-\alpha}\) 的特征上分析均方误差回归,容量瓶颈由高斯随机特征引入;本文则将标度律扩展至具有长尾类别先验的多分类交叉熵,指出分类问题必须通过 PCA 截断方可诱导容量瓶颈,并揭示了独特的 \(\log K\) 因子依赖。
- vs 联想记忆模型 (Cabannes et al., 2024; Li et al., 2026):联想记忆将模式记忆视为纯无噪离散关联;本文不仅严格泛化了无噪情况下的序列学习时间,还处理了带有任意高维方差噪声的高斯混合分类,并在在线 SGD 采样随机性下给出了严格的鞅差集中控制。
- vs 逻辑回归隐式偏差理论 (Soudry et al., 2018; Ravi et al., 2024):经典隐式偏差文献聚焦于可分数据集上晚期向最大间隔方向收敛的渐近行为(\(K/t\) 速率);本文刻画了从初始化到收敛全生命周期中的早期平台期、中期序列学习幂律衰减期以及最终收敛期的完整相图。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ [首次为长尾多分类交叉熵逻辑回归建立了严谨的序列学习与算力最优标度律理论]
- 实验充分度: ⭐⭐⭐⭐☆ [高精度数值 ODE 模拟与蒙特卡洛验证详实,但在真实深度神经网络预训练上的验证仍待展开]
- 写作质量: ⭐⭐⭐⭐⭐ [数学推导严丝合缝,各阶段物理图像与直觉解释极为清晰明了]
- 价值: ⭐⭐⭐⭐⭐ [为理解大语言模型 Next-Token 预测中 Chinchilla 标度律的数学成因提供了重要的理论基石]