Layerwise LQR for Geometry-Aware Optimization of Deep Networks¶

会议: NeurIPS 2026
arXiv: 2605.04230
代码: https://github.com/SimonDufLab/LLQR
领域: 优化/理论
关键词: 二阶优化, 线性二次调节器, 几何感知优化, 自然梯度, 结构化预条件子

一句话总结¶

将深度网络下基于散度诱导的二阶几何最速下降步严格等价重写为层级有限时域线性二次调节器（LQR），并提出通过直接最小化 LQR 目标周期性学习结构化逆预条件子的松弛算法，在避免显式构建与求逆全局曲率矩阵的同时完整保留层间耦合信息。

研究背景与动机¶

深度学习优化算法的迭代更新通常可视为在特定几何空间中的最速下降法。标准梯度下降对应于欧几里得范数下的最速下降，而牛顿法、自然梯度法（NGD）以及近年提出的 K-FAC、Shampoo、Muon 等实用优化器，则试图通过引入二阶曲率或结构化预条件子来改善损失曲面的病态条件数。优化几何的选择不仅决定了算法的局部收敛速率，更深刻塑造了参数优化轨迹的隐式偏置（implicit bias），进而直接影响深度模型的泛化性质与表征形成。

然而，在现代深度神经网络中实现严谨的几何感知优化面临严重的计算瓶颈。真正的二阶或散度度量矩阵（如全局海森矩阵、高斯-牛顿矩阵或 Fisher 信息矩阵）因反向传播链式法则而在网络各层之间高度稠密耦合。在数千万至数十亿参数规模下，显式实例化并求解 $H \Delta \theta = -g$ 是不可能的。为实现工程可扩展性，现存的二阶近似方法（如 K-FAC）普遍采用“过早施加结构约束”的策略——在推导更新目标之前，就强行将全局曲率矩阵切分为逐层的分块对角阵，甚至进一步引入激活与反向导数的独立性假设来做 Kronecker 分解。这种先验结构割裂了网络各层在前向与反向传播中的协同动力学，直接丢弃了层间曲率耦合，导致最终更新方向严重失真。

本文的切入视角是逆向思考结构化近似的施加节点：不要在曲率模型上过早切断层间联系，而是先完整保留网络的顺序结构与跨层依赖，将稠密几何更新精确转化为离散时间最优控制问题。核心 idea：将广义散度诱导的二阶下降步严格等价表述为有限时域线性二次调节器（LQR），利用前向传播线性化表征状态动力学、散度展开定义二次代价，进而通过直接在 LQR 目标下优化结构化逆预条件子（如 E-KFAC/K-FAC），实现无求逆、保持层间耦合感知的高效二阶优化包装器。

方法详解¶

整体框架¶

LLQR 框架由理论精确等价与实用算法松弛两部分构成。输入为当前网络迭代步的参数 $\bm{\theta}^k$、微批次数据及前向激活值，输出为经过几何感知校准的更新步 $\Delta \bm{\theta}^k$。在精确层面，定理证明了只要度量源自层级可分解的散度（涵盖牛顿、高斯-牛顿、自然梯度及中间层度量），其最速下降二次模型与一个有限时域 LQR 轨迹控制问题严格等价，可通过经典 Riccati 后向递推与前向 rollout 精确求解。在工程实用层面，为免去 Riccati 矩阵递推的大规模雅可比存储与求逆开销，算法将更新步参数化为以结构化逆预条件子作用于梯度的形式 $\Delta \bm{\theta}_i = - \bm{U}_i \nabla_{\bm{\theta}_i} L(\bm{\theta}^k)$，周期性地在全局 LQR 目标函数上对 $\bm{U}$ 进行极少步内层优化，随后以指数移动平均（EMA）平滑更新，并在常规迭代步中以极低开销直接作用于梯度向量。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["前向计算与微批次采样<br/>获取状态激活与损失梯度"] --> B["阶段1：层级前向动力学线性化<br/>提取局部雅可比算子 A_i 与 B_i"]
    B --> C["阶段2：散度诱导 LQR 代价构造<br/>伴随变量对齐消去交叉项，确定 Q_i, R_i, M_i"]
    C --> D["阶段3：结构化逆预条件子松弛优化<br/>最小化全局 LQR 目标拟合 U_i（无显式求逆）"]
    D --> E["阶段4：跨迭代 EMA 平滑与外层包装<br/>结构化矩阵向量乘更新主优化器参数"]

关键设计¶

1. 散度诱导二次模型的层级 LQR 精确等价 现有二阶优化难以处理跨层耦合的核心在于全局曲率矩阵 $\bm{H}$ 维度过大。为解开其层级结构，论文考虑形式为 $D(\bm{\theta}^k, \bm{\theta}) = \psi_N(\bm{x}_N(\bm{\theta})) + \sum_{i=1}^{N-1} \psi_i(\bm{x}_i(\bm{\theta}), \bm{\theta}_i)$ 的广义层级可分解散度，利用拉格朗日乘子向量 $\bm{p} = [\bm{p}_1, \dots, \bm{p}_N]$ 将层间映射约束 $f_i(\bm{x}_i, \bm{\theta}_i) - \bm{x}_{i+1} = \mathbf{0}$ 引入增广哈密顿泛函 $\mathcal{H}$。通过选择满足伴随递推式 $\bm{p}_N^k = \nabla_{\bm{x}_N} \psi_N(\bm{x}_N^k)$ 及 $\bm{p}_i^k = \bm{A}_i^\top \bm{p}_{i+1} + \nabla_{\bm{x}_i} \psi_i(\bm{x}_i^k, \bm{\theta}_i^k)$ 的特定乘子，迫使 $\nabla_{\bm{x}} \mathcal{H} = \mathbf{0}$，从而使原本极其复杂的二次形式严格分解为局部状态摄动 $\delta \bm{x}_i$ 与控制摄动 $\delta \bm{\theta}_i$ 的局部块成本之和： $$\Delta \bm{\theta}^\top \bm{H} \Delta \bm{\theta} = \delta \bm{x}_N^\top \bm{Q}_N \delta \bm{x}_N + \sum_{i=0}^{N-1} \begin{bmatrix} \delta \bm{x}_i \\ \delta \bm{\theta}_i \end{bmatrix}^\top \begin{bmatrix} \bm{Q}_i & \bm{M}_i^\top \\ \bm{M}_i & \bm{R}_i \end{bmatrix} \begin{bmatrix} \delta \bm{x}_i \\ \delta \bm{\theta}_i \end{bmatrix}$$ 约束动力学为 $\delta \bm{x}_{i+1} = \bm{A}_i \delta \bm{x}_i + \bm{B}_i \delta \bm{\theta}_i$（其中 $\delta \bm{x}_0 = \mathbf{0}$）。最速下降问题被严格转化为标准有限时域 LQR 控制问题。当使用 Riccati 递推求解时，曲率求逆的复杂度从全网络总维度的立方降至逐层状态与参数维度的立方和，更重要的是确立了网络动力学与下降几何分离的理论基准。

2. 结构化逆预条件子直接优化松弛（无显式求逆） 虽然精确 LQR 提供了理论基准，但在现代深度网络中逐层运行 Riccati 递推仍需存储庞大雅可比矩阵并执行矩阵求逆。LLQR 的核心算法突破在于更新参数化：不求解未约束的控制变量 $\delta \bm{\theta}_i$，而是强制定式 $\delta \bm{\theta}_i = - \bm{U}_i \nabla_{\bm{\theta}_i} L(\bm{\theta}^k)$，直接在全局 LQR 目标函数上寻找最优的逆预条件算子 $\bm{U} = \mathrm{diag}(\bm{U}_0, \dots, \bm{U}_{N-1})$。将这一参数化代入 LQR 目标后，目标函数变为关于 $\bm{U}$ 的凸二次函数： $$\min_{\bm{U}} \; \nabla_{\bm{x}_N} \ell(\bm{x}_N)^\top \delta \bm{x}_N + \frac{1}{2} \delta \bm{x}_N^\top \bm{Q}_N \delta \bm{x}_N + \sum_{i=0}^{N-1} \left( \frac{1}{2} \delta \bm{x}_i^\top \bm{Q}_i \delta \bm{x}_i + \frac{1}{2} \delta \bm{\theta}_i^\top \bm{R}_i \delta \bm{\theta}_i + \delta \bm{\theta}_i^\top \bm{M}_i \delta \bm{x}_i \right)$$ 其中状态前向演化服从 $\delta \bm{x}_{i+1} = \bm{A}_i \delta \bm{x}_i - \bm{B}_i \bm{U}_i \nabla_{\bm{\theta}_i} L(\bm{\theta}^k)$。这种松弛机制带来了双重飞跃：首先，$\bm{U}$ 直接作为逆算子被优化拟合，整个流程中不需要对任何曲率矩阵或 Kronecker 因子求逆，彻底规避了二阶优化中最脆弱且计算沉重的矩阵求逆步；其次，不同于传统 K-FAC 先做分块对角截断再独立求逆，LLQR 中的 $\bm{U}_i$ 是在包含全局层间扰动演化链的统一 LQR 成本下拉动优化的，即便限制 $\bm{U}$ 自身为块对角结构，各层的参数也感知到了跨层动力学的协同约束。

3. 结构化参数化族与无缝外层优化器包装 为在内存与表达力之间灵活权衡，LLQR 支持多种逆预条件子参数化结构 $\mathcal{S}$，主要包括对角（Diagonal）、K-FAC 与 E-KFAC： - K-FAC 结构：将权重梯度矩阵 $\bm{X}_i \in \mathbb{R}^{m_i \times n_i}$ 的预条件动作参数化为 $\bm{Y}_i = \bm{C}_i \bm{X}_i \bm{D}_i^\top$，仅需维护较小的因子矩阵； - E-KFAC 结构：解耦坐标变换与尺度缩放，维护左右正交变换因子 $\bm{Q}_i^{\mathrm{L}}, \bm{Q}_i^{\mathrm{R}}$ 与可学习的对角逆缩放向量 $\bm{s}_i \in \mathbb{R}^{m_i n_i}$，即 $\bm{Y}_i = \bm{Q}_i^{\mathrm{L}} \operatorname{unvec}(\bm{s}_i \odot \operatorname{vec}((\bm{Q}_i^{\mathrm{L}})^\top \bm{X}_i \bm{Q}_i^{\mathrm{R}})) (\bm{Q}_i^{\mathrm{R}})^\top$。在 LLQR 中，E-KFAC 无需显式特征值分解，而是作为一种高表达力的结构化归纳偏置； - 周期拟合与外层包装：逆预条件子无需每步重构，算法每隔 $n$ 步（如 ImageNet 上每千余步，相当于每轮仅 2-4 次）执行 $T$ 步（通常仅 25 步）内层微小 SGD 优化，随后通过 EMA 平滑系数 $\beta \in [0.85, 0.95]$ 更新历史 $\bm{U}$。在绝大多数常规迭代中，预条件步仅是一次极快的矩阵向量乘法，产出的预条件梯度 $\Delta \bm{\theta}_k = \bm{U} g_k$ 直接作为无缝包装器输入到 SGDM 或 AdamW 中执行参数更新。

损失函数 / 训练策略¶

在工程实现中，矩阵 $\bm{Q}_i, \bm{R}_i, \bm{M}_i$ 及雅可比乘积从不显式物化，而是利用自动微分的 JVP（雅可比-向量积）、VJP（向量-雅可比积）和 HVP（海森-向量积）按需算子化评估。针对深层网络激活占用内存过高的问题，算法引入分块切片（Chunking）机制，将内层拟合的大微批次沿 batch 维度拆为小 chunk 串行累加梯度。阻尼项可形式化纳入控制惩罚 $\bar{\bm{R}}_i = \bm{R}_i + \lambda_i \bm{I}$，但实验表明由于直接优化逆算子且内层步数受限，即使阻尼设为 0 系统也具有极高的数值稳定性。

实验关键数据¶

主实验¶

论文在不同架构和数据集上系统验证了 LLQR 作为优化器包装器的加速效果与最终精度。在 CIFAR-10/100 上，LLQR 显著改善了 ResNet-18、PyramidNet-110、VGG-16-BN 及 WideResNet-28-10 的表现；在 ImageNet（ResNet-50）和机器翻译（Fairseq Transformer）上同样展现出优异的跨域泛化能力。

数据集 / 任务	模型 / 架构	主优化器	预条件结构与散度	LLQR 性能 (Top-1 / BLEU)	Baseline 性能	相对耗时倍率
CIFAR-10	ResNet-18	SGDM	E-KFAC (NGD)	96.37 ± 0.05%	96.19 ± 0.07%	1.31×
CIFAR-10	ResNet-18	AdamW	E-KFAC (NGD)	94.95 ± 0.05%	94.54 ± 0.11%	1.15×
CIFAR-100	ResNet-18	SGDM	E-KFAC (NGD)	79.96 ± 0.10%	78.97 ± 0.22%	1.32×
CIFAR-100	ResNet-18	AdamW	E-KFAC (NGD)	77.07 ± 0.16%	76.09 ± 0.27%	1.13×
CIFAR-100	PyramidNet-110	SGDM	E-KFAC (NGD)	84.40 ± 0.21%	83.96 ± 0.11%	1.41×
CIFAR-100	VGG-16-BN	SGDM	E-KFAC (NGD)	76.32 ± 0.16%	75.65 ± 0.12%	1.17×
CIFAR-100	WRN-28-10	SGDM	E-KFAC (NGD)	82.80 ± 0.17%	82.50 ± 0.08%	1.25×
ImageNet (100 ep)	ResNet-50	SGDM	E-KFAC (NGD)	78.05 ± 0.05%	77.42 ± 0.12%	1.032×
ImageNet (100 ep)	ResNet-50	SGDM	K-FAC (NGD)	77.98 ± 0.08%	77.42 ± 0.12%	1.027×
IWSLT14 De–En	Transformer	AdamW	E-KFAC (NGD)	34.51 ± 0.05 BLEU	34.24 ± 0.12 BLEU	1.16×

在与当前先进近似二阶优化器（严格对齐时钟预算的 AdaFisher 标准协议）的横向对比中，LLQR 表现出绝对的精度与计算开销优势：

数据集与协议	优化器配置	预条件形式	实际测试精度 (Top-1)	每轮训练耗时倍率
CIFAR-100 (BS=256)	SGDM Baseline	无预条件	77.52 ± 0.13%	1.00×
CIFAR-100 (BS=256)	LLQR (Ours)	E-KFAC (NGD)	78.34 ± 0.08%	1.25×
CIFAR-100 (BS=256)	AdaFisher†	动态对角 Fisher	77.28 ± 0.20%	1.33×
CIFAR-100 (BS=256)	K-FAC†	经典 Kronecker	76.03 ± 0.30%	2.32×
CIFAR-100 (BS=256)	Shampoo†	张量积预条件	76.78 ± 0.20%	10.82×
ImageNet (90 ep)	SGDM Baseline	无预条件	77.19 ± 0.02%	1.00×
ImageNet (90 ep)	LLQR (Ours)	E-KFAC (NGD)	77.59 ± 0.04%	1.034×
ImageNet (90 ep)	AdaFisher†	动态对角 Fisher	76.95%	基线等时预算
ImageNet (90 ep)	K-FAC†	经典 Kronecker	70.96%	基线等时预算

消融实验¶

针对预条件子结构表达力、散度类型以及内层优化设计的细粒度消融分析如下：

配置维度	实验设定 / 消融变体	核心表现与指标 (ResNet-18 CIFAR-100 SGDM)	机制说明与分析
预条件结构 $\mathcal{S}$	对角结构 (Diagonal)	79.20 ± 0.30% (耗时 1.08×)	仅能做逐坐标重缩放，无法捕捉跨输入输出维度的局部曲率
预条件结构 $\mathcal{S}$	K-FAC	79.22 ± 0.11% (耗时 1.23×)	引入 Kronecker 积因子，表征能力显著提升
预条件结构 $\mathcal{S}$	E-KFAC (完整设计)	79.96 ± 0.10% (耗时 1.32×)	解耦特征变换基底与各向异性缩放，提供最高精度的曲率校准
几何散度类型	牛顿散度 (Bregman Gap)	79.83 ± 0.08%	逼近真实损失 Hessian，但在随机小批次下噪声略高于自然梯度
几何散度类型	自然梯度 (KL 散度)	79.96 ± 0.10%	基于预测分布的 Fisher 几何更加鲁棒平滑，整体效果最佳
逆矩阵初值策略	历史热启动 (Warm-start)	早期略快但后期不稳定，ImageNet 掉点至 77.4%	极值特征积累导致轨迹偏离，不如单位阵重置配合 EMA 鲁棒
逆矩阵更新策略	单位阵重置 + EMA 平滑	稳定收敛，达到 78.05%	保证内层优化探索方向不过拟合单批次噪声，训练稳定
内层优化求解器	共轭梯度法 (CG + LineSearch)	带来显著计算拖慢且对迭代步数敏感	破坏了 GPU 向量化并增加调参负担，表现不如 SGD
阻尼系数 $\lambda$	非零阻尼 ($\lambda > 0$)	小阻尼无明显收益，大阻尼导致性能劣化	由于直接在逆空间优化且内层步数截断，无需额外正规化阻尼

关键发现¶

大模型长周期任务下开销被极大摊薄：在 ImageNet（ResNet-50）百轮训练中，由于单轮步数长达 5000 步，每 1500 步仅需 25 次前向反向步来重拟合预条件子，使得整体额外时间开销仅为 2.7%~3.2%（耗时倍率 1.03×）。而在等时钟对比中，LLQR 仅用 90 轮（用时为标准 100 轮的 0.93×）即可达到 77.59%，全面超越跑满 100 轮的标准 SGDM（77.42%）。
层间耦合感知击碎对角二阶瓶颈：在经典非凸 Rosenbrock 谷底与 CIFAR 实验中，直接利用 Hessian 对角元逆做预条件的方法（如局部对角牛顿法）会因为丢弃层间交叉项而在不定曲率区域停滞或偏离；而 LLQR 即使在限制预条件子结构自由度的前提下，由于是在全局动力学轨迹上优化逆算子，其更新方向与真实全牛顿步的余弦相似度始终保持在高位。
显著加速 Transformer 的 Grokking 突现泛化：在 5 个模运算算法任务（加、减、乘、除、指数运算，模 97）中，LLQR 显著缩短了 Transformer 在过拟合平台期的停滞时间，将达到 95% 测试集准确率所需的步数大幅提前，且在绝对壁钟时间上实现了真实加速。

亮点与洞察¶

将最速下降直接等价于离散时间控制（LQR）：摆脱了以往仅对标准多层感知机或纯牛顿步套用最优控制的狭隘视角，将控制视界拓展到由任意层级可分解散度（涵盖 Fisher、Gauss-Newton 及中间层表征度量）诱导的几何最速下降，建立了优化论与现代控制论的严密桥梁。
逆空间直接参数化与优化，彻底免除求逆操作：传统二阶优化器的计算死结在于“先近似曲率矩阵、再计算逆矩阵”，不仅开销高达立方阶，而且数值极易因病态特征值发散；LLQR 反其道而行之，直接对逆预条件子进行参数化，并在 LQR 凸二次目标下拉动更新，将繁琐的数值线性代数求逆转化为可自动微分的参数拟合问题。
跨层曲率信息向块内自由度的优雅投影：在结构化松弛中，LLQR 并不像传统方法那样在建模前粗暴假设层间独立，而是将完整的层间前向线性化扰动反馈进 LQR 目标中；因此各层即便只维护 Kronecker 因子，其学到的取值也浸润了跨层动力学的信息。

局限与展望¶

作者承认的局限：相比于纯一阶优化器，LLQR 依然需要额外的 GPU 显存用于缓存前向激活以执行 JVP/VJP/HVP 计算，并需在内存中维护预条件子的因子参数，对显存极度紧张的超大模型训练构成挑战。
实验与设定层面的局限：目前的 LLQR 实现主要针对标准的前馈层、卷积层及常规注意力投影层进行了结构化适配；对于具备复杂分支、动态稀疏激活（如 MoE）或长序列循环状态的模型，前向动力学矩阵的构建与分块切片复杂度显著增加。
未来改进方向：可进一步将 LLQR 推广至扩散模型（在去噪时间步与网络深度两个维度同时引入层级散度度量），或者与隐式神经系统（Deep Equilibrium Models、Neural ODEs）相结合，推导出无反向传播深度的连续时间几何优化器。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次建立广义散度二阶优化与层级 LQR 的精确数学等价，并独创直接优化逆算子的松弛机制。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖经典玩具函数、多模型 CIFAR、ResNet-50 ImageNet、Transformer 翻译及 Grokking 实验，数据严谨详实。
写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰优美，附录数学展开完整自洽，算法动机与工程取舍阐述透彻。
价值: ⭐⭐⭐⭐⭐ 突破了传统二阶优化在计算开销与求逆稳定性上的工程瓶颈，为工业界大模型二阶优化提供了全新范式。

配置维度	实验设定 / 消融变体	核心表现与指标 (ResNet-18 CIFAR-100 SGDM)	机制说明与分析
预条件结构 \(\mathcal{S}\)	对角结构 (Diagonal)	79.20 ± 0.30% (耗时 1.08×)	仅能做逐坐标重缩放，无法捕捉跨输入输出维度的局部曲率
预条件结构 \(\mathcal{S}\)	K-FAC	79.22 ± 0.11% (耗时 1.23×)	引入 Kronecker 积因子，表征能力显著提升
预条件结构 \(\mathcal{S}\)	E-KFAC (完整设计)	79.96 ± 0.10% (耗时 1.32×)	解耦特征变换基底与各向异性缩放，提供最高精度的曲率校准
几何散度类型	牛顿散度 (Bregman Gap)	79.83 ± 0.08%	逼近真实损失 Hessian，但在随机小批次下噪声略高于自然梯度
几何散度类型	自然梯度 (KL 散度)	79.96 ± 0.10%	基于预测分布的 Fisher 几何更加鲁棒平滑，整体效果最佳
逆矩阵初值策略	历史热启动 (Warm-start)	早期略快但后期不稳定，ImageNet 掉点至 77.4%	极值特征积累导致轨迹偏离，不如单位阵重置配合 EMA 鲁棒
逆矩阵更新策略	单位阵重置 + EMA 平滑	稳定收敛，达到 78.05%	保证内层优化探索方向不过拟合单批次噪声，训练稳定
内层优化求解器	共轭梯度法 (CG + LineSearch)	带来显著计算拖慢且对迭代步数敏感	破坏了 GPU 向量化并增加调参负担，表现不如 SGD
阻尼系数 \(\lambda\)	非零阻尼 (\(\lambda > 0\))	小阻尼无明显收益，大阻尼导致性能劣化	由于直接在逆空间优化且内层步数截断，无需额外正规化阻尼