跳转至

A Differentiable Interior-Point Method in Single Precision

会议: NeurIPS 2026
arXiv: 2605.17913
论文: https://arxiv.org/abs/2605.17913
代码: https://github.com/qpax-solver/qpax
领域: 优化/理论
关键词: 可微凸优化, 二次规划 (QP), 内点法 (IPM), 单精度浮点 (f32), 隐式互补松弛

一句话总结

针对传统内点法在单精度浮点运算下因互补松弛条件导致线性系统病态发散的瓶颈,提出基于 softplus 辅助变量的隐式互补紧致投影与谱有界 KKT 系统,首次在 f32 下实现了高鲁棒二次规划求解与高保真隐函数求导。

研究背景与动机

约束凸优化与二次规划(QP)是现代控制系统、动力学仿真、运筹学及机器学习决策模块的算法基石。随着端到端深度学习系统的扩展,将带约束优化层作为可微神经算子无缝嵌入神经网络并基于数据端到端优化超参数,已成为可微物理和神经控制领域的核心研究方向。在各类优化算法中,原始-对偶内点法(Primal-Dual IPM)凭借超线性收敛速度、理论完备性及高精度求解能力,历来是高精度凸优化求解器的黄金标准。

然而,传统内点法严重依赖双精度(float64)浮点数,这一软硬件断层阻碍了其在现代机器学习体系中的推广。现代深度学习计算架构与 GPU/TPU 硬件加速器主要针对单精度(float32)乃至更低精度的混合精度吞吐量进行了极致优化。在传统内点法中,为了满足 KKT 互补松弛条件 \(s \odot z \to 0\),牛顿方程紧致化后的对角缩放项 \(D(s \oslash z)\) 会出现极度分化:非激活约束的松弛变量保持正值而乘子趋近于 0 导致比例发散,激活约束的松弛变量趋于 0 导致比例湮灭。这种内在的谱病态在 float64 下尚能被直接分解法容忍,但在单精度 float32 下会迅速引发严重的数值溢出、下溢、步长破坏以及隐函数微分阶段的梯度失效(频繁出现 NaN)。

本文正是为了打破这一长期存在的基础性障碍,探讨如何构建能够在单精度浮点运算下可靠求解并稳定微分的内点法求解器。作者的核心视角是:若抛弃对互补条件的显式逼近,改用通过变换恒等满足互补流形的隐式参数化表达,便能从根本上重塑 KKT 牛顿系统的谱特征。本文的核心 idea 是通过引入 softplus 辅助变量映射恒等满足互补松弛条件,构建中间对角块谱严格有界在 \([0, 1]\) 且大部静态对称的紧致 KKT 线性系统,配合基于中心路径反向松弛的隐函数微分框架,实现首个完全原生支持单精度的高精度可微内点法求解器。

方法详解

整体框架

整个算法由三个连续解耦的计算阶段组成:前向高精度 QP 求解(Algorithm 1)、中心路径反向松弛(Algorithm 2)以及隐函数求导与梯度回传(Algorithm 3)。在前向求解阶段,求解器利用辅助变量 \(v\) 将不等式约束松弛变量 \(s\) 与拉格朗日乘子 \(z\) 投影到互补流形上,利用部分紧致化谱有界对称 KKT 系统迭代将残差推向设定容差;随后,算法无需重新初始化,直接从收敛解出发沿着中心路径反向追踪至预设的平滑松弛因子 \(\kappa_{\mathrm{relax}}\);最后,在松弛点处复用已分解缓存的 KKT 因子,通过隐函数定理计算目标损失关于问题数据的高精度向量-雅可比积(VJP)。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入 QP 参数<br/>Q, q, A, b, G, h"] --> B["前向高精度求解:隐式互补流形投影<br/>软加变换耦合 v 与中心参数 κ"]
    B --> C["KKT 线性求解:谱有界部分紧致对称系统<br/>LDLᵀ 分解并缓存静态对称块因子"]
    C --> D["反向平滑松弛:中心路径反向追踪<br/>从收敛解平滑外推至 κ_relax 消除非光滑次梯度"]
    D --> E["隐式微分阶段:复用分解缓存提取高保真梯度<br/>解析重构 ∇Q, ∇q, ∇A, ∇b, ∇G, ∇h 回传网络"]

关键设计

1. 隐式互补流形投影:通过 softplus 变换消除病态发散

传统内点法在迭代收敛时显式扰动互补条件 \(z \odot s = \kappa \to 0\),造成对角比值 \(s_i / z_i\) 出现极端尺度失衡(跨越 \(0\) 到 \(\infty\)),这是单精度数值崩溃的根源。为彻底规避该谱发散,该设计引入辅助变量 \(v \in \mathbb{R}^p\) 与参数化缩回映射(Retraction Map)\(b_\kappa(v): \mathbb{R} \to \mathbb{R}_+\),定义为加权 softplus 函数: $\(b_\kappa(v) = \frac{v + \sqrt{v^2 + 4\kappa}}{2}\)$ 该映射具有两项代数性质:首先,通过恒等式 \(b_\kappa(v) \cdot b_\kappa(-v) = \kappa\),只要令乘子与松弛变量满足 \(z = b_\kappa(v)\) 且 \(s = b_\kappa(-v)\),互补松弛条件在整个迭代空间内自然且精确地成立;其次,其导数算子严格满足导数和守恒与导数严格有界性: $\(\partial_v b_\kappa(v) + \partial_v b_\kappa(-v) = 1, \qquad 0 < \partial_v b_\kappa(v) \leq 1\)$ 在低精度浮点下,为避免当 \(v < 0\) 且 \(v^2 \gg \kappa\) 时出现灾难性相消(Catastrophic Cancellation),底层实现采用两段分支公式平滑过渡:当 \(v \ge 0\) 时使用直接解析式;当 \(v < 0\) 时通过有理化变形为 \(\frac{2\kappa}{\sqrt{v^2 + 4\kappa} - v}\)。这保证了变量始终严格驻留在对偶可行域内部,彻底切断了除零与下溢通道。

2. 谱有界部分紧致对称系统:保障 f32 下 LDLᵀ 分解数值稳定

未紧致的全量牛顿系统在单精度下难以有效传递辅助增量 \(\Delta v\) 的梯度信号。为了同时获得计算效率与数值条件数保证,该方法通过消元推导出一种特有的部分紧致对称 KKT 方程组: $\(\begin{bmatrix} Q - G^\top G & G^\top & A^\top \\ G & -B_\kappa(-v) & 0 \\ A & 0 & 0 \end{bmatrix} \begin{bmatrix} \Delta x \\ \Delta v \\ \Delta y \end{bmatrix} = -\begin{bmatrix} r_t - G^\top (r_i + r_z - r_s) \\ r_i - r_s - c r_\kappa \\ r_e \end{bmatrix}\)$ 其中对角矩阵 \(B_\kappa(-v) = \mathrm{Diag}(\partial_v b_\kappa(-v))\),标量敏感度 \(c = \partial_\kappa b_\kappa(v)\)。与传统紧致方程中的发散项 \(D(s \oslash z)\) 截然不同,该系统的唯一随迭代变化的模块为对角阵 \(B_\kappa(-v)\),且其对角元素被严格禁锢在开区间 \((0, 1]\) 之间。外部鞍点结构中的其余各个子块(\(Q - G^\top G\)、\(G\)、\(A\))在迭代过程中完全恒定。这带来了双重决定性优势:其一,线性系统的极值特征值被严格约束,避免了收敛极限处系统矩阵条件数的急剧恶化,使经典的对称不定 \(LDL^\top\) 分解在单精度下能够稳健执行;其二,由于变化部分仅仅是受限的对角项,极大地利于快速低秩更新或因子重用。

3. 中心路径反向追踪与平滑松弛隐式微分

直接在精确收敛点(\(\kappa^* \approx 0\))根据隐函数定理求导会因激活集边界处的非光滑接触几何导致次梯度信息退化或出现病态阶跃。该设计提出中心路径反向松弛机制(Reverse Continuation):在 Algorithm 1 求解达到终止精度后,Algorithm 2 直接将收敛解作为热启动初始点,以受控步长将障碍参数从极小值提升至目标松弛参数 \(\kappa_{\mathrm{relax}}\)。因为此阶段相当于在已经收敛的稳定凸流形上反向回溯,收敛极其迅速且其 KKT 矩阵结构与前向求解高度重合,前向缓存的分解因子 \(\tilde{K}\) 可被高复用率利用。在获得平滑点 \((x, y, z, s)\) 后,隐式微分(Algorithm 3)通过求解单次伴随 KKT 系统计算微分项 \((dx, dy, dz, ds)\),并直接利用紧凑张量积计算解析梯度(如 \(\nabla_Q \ell = \frac{(dx)x^\top + x(dx)^\top}{2}\),\(\nabla_q \ell = dx\) 等),完全避免了构造巨型全雅可比矩阵,且免除了传统内点法中复杂的拉格朗日乘子重缩放步骤。

实验关键数据

主实验

论文在多面体欧氏投影基准(Euclidean projection onto a polytope)、动态双层轨迹规划(Bilevel trajectory optimization)以及多智能体安全滤波器学习(Multi-agent safety filter learning)三个场景下,系统评测了单精度(f32)和双精度(f64)下隐式互补(Implicit)与传统显式互补(Explicit)内点法的表现。在欧氏投影中,以基于真实激活集构建的硬投影正交投影算子 \(J_{\mathrm{hard}} = I_n - G_A^\top (G_A G_A^\top)^{-1} G_A\) 计算的解析导数 \(g_{\mathrm{hard}} = J_{\mathrm{hard}} v\) 作为真实梯度 Ground Truth。

实验任务与维度配置 评测指标 本文 (Implicit f32) 传统基准 (Explicit f32) 对比增益 / 现象说明
欧氏多面体投影 (\(n=20, p=25, m=15, \kappa=10^{-6}\)) 梯度绝对误差 \(\mathbf{\approx 10^{-6}}\) \(\approx 10^{-1}\) (停滞) 精度提升 5 个数量级,紧贴松弛理论下界
欧氏多面体投影 (\(n>5\), 变约束比 \(m_r\)) 求解与求导失败率 (NaN) \(\mathbf{0\%}\) (完全稳定) \(\approx 50\% - 100\%\) 消除多重激活集导致的数值溢出崩溃
双层轨迹规划 (森林场景, \(\kappa_{\mathrm{relax}}=10^{-4}\)) 外层 L-BFGS 耗时与收敛性 \(\mathbf{100\%}\) 稳定收敛 (媲美 f64) 梯度失效 / 迭代崩溃 (NaN) 首次实现在紧松弛下低精度双层轨迹优化
多智能体安全滤波 (7 智能体, 70 约束, CBF-QP) 仿真测试碰撞率 (Pairwise/Obs) \(\mathbf{0\%}\) 碰撞违背 训练崩溃 (NaN, 无法运行) 成功解锁多智能体端到端低精度训练
多智能体安全滤波 (2~9 智能体 GPU 训练) 最大支持 Batch Size \(\mathbf{2\times}\) 提升 (相对于 f64) 无法运行 (\(n \ge 2\) 崩溃) 显存减半,单卡批大小直接翻倍

消融实验

消融实验深入追踪了传统显式内点法在单精度(f32)下出现浮点异常(NaN)的细粒度操作源头。在 \(n=20, p=25\) 的多面体投影中,固定 63 个测试实例,沿不同激活约束数 \(m \in \{10, 12, 15\}\) 与不同中心参数 \(\kappa_{\mathrm{relax}} \in [10^{-2}, 10^{-9}]\) 进行首现 NaN 故障归因分析。

激活约束数 \(m\) 松弛因子 \(\kappa_{\mathrm{relax}}\) 总测试样例数 仿射预测步 (Pred.) NaN 频次 反向松弛步 (Rel.) NaN 频次 本文隐式方法 (Implicit f32) 失败频次 说明
\(m=10\) \(10^{-2}\) 63 26 0 0 显式法前向预测步首现不稳定
\(m=10\) \(10^{-4}\) 63 26 7 0 紧松弛下显式法反向松弛开始失效
\(m=10\) \(10^{-6}\) 63 26 16 0 显式法累计失效比例超过 66%
\(m=10\) \(10^{-9}\) 63 26 33 0 显式法几近全部崩溃 (59/63)
\(m=15\) \(10^{-3}\) 63 25 1 0 约束数增加时显式法在较大 \(\kappa\) 即崩溃
\(m=15\) \(10^{-6}\) 63 25 14 0 隐式法在任意尺度与深度松弛下均 0 崩溃
\(m=15\) \(10^{-9}\) 63 25 36 0 显式法累计 61/63 报废,隐式法全流程 0 异常

关键发现

  • 显式内点法在 f32 下崩溃具有固定路径:NaN 追踪分析揭示,传统显式方法在单精度下的失效绝非随机噪声,而是百分之百集中在仿射预测方向求解(Affine Predictor Step)和后继的中心路径松弛迭代(Relaxation Continuation)中;其根源在于 \(z \oslash s\) 比例失衡引发除零与矩阵条件数爆炸。
  • 谱有界性对单精度求解是决定性的:本文隐式方法消除了对 \(B_\kappa(v)\) 的直接依赖,在所有测试维度和松弛参数下保持了 0% 的失败率,并且其梯度误差随 \(\kappa_{\mathrm{relax}}\) 的缩小严格呈现理论线性下降,直至触及 \(10^{-6}\) 的单精度机平极限。
  • 端到端训练中的硬件收益直接兑现:在 CBF-QP 多智能体神经安全滤波器学习中,Implicit f32 不仅消除了 Explicit f32 在 \(n \ge 2\) 时必然遭遇的训练发散,其占用的显存仅为 f64 的一半,使得同等 GPU 显存限制下的训练 Batch Size 严格翻倍。

亮点与洞察

  • 变换巧妙消除极端发散项:利用 softplus 映射天然的互补对称性与导数有界性(值域严格限定在 \((0, 1]\)),把传统牛顿方程中具有极端病态对角块 \(D(s \oslash z)\) 替换为谱有界的静态鞍点块,设计极其优雅且代数结构自洽。
  • 计算图与因子高度复用:从前向求解到反向路径松弛,再到向量-雅可比积回传,三个阶段均运行在拓扑等价的线性系统上,使得昂贵的矩阵因式分解在不同计算阶段之间得以高度缓存与复用,显著摊平了可微优化的开销。
  • 打通学术算法与生产硬件的断层:以往嵌入式凸优化与强化学习控制层多因单精度数值不稳定性而被迫退回一阶算法(如 ADMM/OSQP)或强制运行昂贵的 float64,该研究成功证明内点法通过精巧数学重构完全可在现代机器学习硬件原生精度下高效运行。

局限与展望

  • Hessian 块的局部稠密化:在部分紧致化推导中,矩阵左上角子块从原始的目标二次型 \(Q\) 转变为 \(Q - G^\top G\),当不等式约束矩阵 \(G\) 具有高度稀疏但跨度广泛的结构时,该变换可能破坏原有的大规模稀疏性,增加直接求解器的计算负担。
  • 目前仅局限于二次规划(QP):虽然论文框架展现了优异的普适性潜力,但正文仅针对二次规划(QP)进行了详尽公式推导与实验验证,向二阶锥规划(SOCP)及半定规划(SDP)等广义凸对称锥的拓展仍有待进一步深入开发。
  • 未来方向:探索针对大规模超稀疏场景下的无矩阵(Matrix-free)迭代 Krylov 求解器结合方案,并将该单精度可微层推广部署至人形机器人高频全身控制(Whole-body Control)与端到端强化学习大规模并行仿真环境中。

相关工作与启发

  • vs OptNet (Amos & Kolter, 2017) / CvxpyLayers (Agrawal et al., 2019): 传统可微凸优化层要么依赖算子分裂一阶算法,要么在内点法实现中默认使用 float64 避免病态;本文彻底从底层数值代数结构改造了内点法,首次赋予其在 float32 单精度下工业级求解与反向传播的稳定性。
  • vs 显式中心路径松弛法 (Tracy & Manchester, 2024): 后者揭示了中心路径平滑松弛在规避非光滑次梯度中的有效性,但沿用了经典显式残差体系,在低精度下极易因预测步病态崩溃;本文与其形成互补并从根源上加固了前向与松弛过程的谱稳定性。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 巧妙借助 softplus 映射代数性质重构出谱有界 KKT 系统,彻底突破了经典内点法数十年来对双精度的依赖。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖解析硬投影基准、细粒度 NaN 故障根因追踪、双层路径规划以及高维多智能体安全强化学习全流程。
  • 写作质量: ⭐⭐⭐⭐⭐ 数学推导严密精炼,对病态根源剖析透彻,图表数据丰富且极具说服力。
  • 价值: ⭐⭐⭐⭐⭐ 为大规模端到端可微物理、具身智能及硬件加速凸优化层提供了极为扎实且高性能的基础设施支持。