输入凸神经网络训练的一种提升方法

# 一种用于输入凸神经网络训练的"提升"方法 来源：https://arxiv.org/html/2605.24274

###### 摘要

输入凸神经网络（ICNN）广泛应用于一系列学习任务——对数凹密度估计、凸势归一化流、最优传输以及高维贝叶斯后验的传输映射反演。所有这些任务共享一个结构性约束：ICNN的层间权重必须保持非负。强制执行该约束的标准方法是投影梯度下降（PGD），即投影到非负锥上，它应用一种硬性的、非光滑的投影——这是ADMM风格约束分裂的刚性罚极限——而其在经典意义上的收敛保证无法迁移到非光滑的ICNN训练场景中；可微分的替代方案，即softplus重参数化，则会随权重幅度指数级地衰减梯度，导致层间权重"死亡"和损失停滞，从而阻碍训练。为解决这一局限性，并受PDE约束反问题中参数扩展"提升"方法的启发，我们提出了**提升（lift）**方法：并非直接约束层间权重，而是训练一个无约束的超网络，该超网络根据输入批次的一个置换不变摘要来发射这些权重。这为训练动力学增加了一种随机性来源，使损失景观变得平缓，让迭代能够逃离直接softplus方法陷入停滞的梯度衰减区域。我们将这种平缓效应追溯至三个结构性要素——一个作为**松弛量**的可学习偏置、一个以目标批次为条件的超网络**主体**，以及通过批次随机性将两者耦合的**交叉协方差**——并且证明每个要素都是必要的：删除任何一个要素都会破坏承载平缓效应的交叉协方差。通过在从一维玩具目标到图像风格潜变量的多尺度对数凹能量基建模，以及在21维表格基准上的凸势归一化流实验中，我们证明**提升方法达到了比PGD和直接softplus都更低的测试损失，并将受平台限制的训练轨迹转变为山谷下降轨迹**。

请参阅图注 (a) 参数空间视图。请参阅图注 (b) 损失空间视图。

图 1: 三种正性重参数化在对数凹EBM训练中的表现，三种不同的结果（21维表格目标；报告每种方法在最低验证损失检查点处的测试负对数似然）。(a) 在一个二维PGD锚定切片上的损失景观（第5.3节），收敛后的超网络位于原点（金色星标）；(b) 同一运行中的留出验证损失随迭代次数的变化。**超网络**沿着盆地下降到最深损失；**PGD**停留在锥边界；**直接softplus**被困在读出肩部（第2.2节），其损失比另外两者高出一个数量级。提升方法相对于PGD的优势是引理1带来的景观平滑收益；相对于直接方法的差距是推论1的Kramers逃逸。部分复现结果的代码可在GitHub上获取。

## 1 引言

输入凸神经网络（ICNN）（Amos等人，2017）参数化了一类凸标量场，这些标量场驱动着概率建模和贝叶斯推理中的多个现代学习任务。它们是对数凹（Prékopa，1971；Saumard和Wellner，2014）能量基模型的负对数密度基础；其梯度定义了凸势归一化流（Huang等人，2021）；是用于传输映射后验采样的PCP-Map（Wang等人，2025；Bunne等人，2022）的凸势；以及ICNN参数化最优传输中的Brenier势（Makkuva等人，2020；Korotin等人，2021）。这些传输映射构建属于高维贝叶斯反问题生成式后验采样的更广泛趋势的一部分，在该领域中，基于分数和扩散模型的采样器也已被应用于地震成像（Baldassari等人，2024；Siahkoohi等人，2026）。在整个生态系统中，环境维度从玩具目标上的单个坐标变化到图像密度和PCP-Map应用上的数千像素，而所有这四个应用共享一个结构特征：输入凸性要求层间权重的正性，即 θ ⪰ 0。强制执行该约束的主要实用方法是投影梯度下降（PGD）（Amos等人，2017）：在 θ 上执行无约束步骤，然后进行投影 θ ← max(θ, 0)。该投影在活动集（迭代聚集处）上是不可微的，而PGD的经典收敛保证（对于具有闭凸约束的光滑凸问题为 O(1/k)，对于光滑非凸问题达到稳定点的率为 O(1/k)（Nesterov, 2018; Beck, 2017））假设了一个Lipschitz光滑的目标函数，而ICNN训练景观在其中并不满足，因此这些保证无法迁移。一种可微分的替代方案是将层间权重重参数化为 θ = ψ(θ̃)，其中 θ̃ ∈ ℝᵈ 无约束，ψ 为单调非负映射，这有两个真正的家族——光滑的softplus家族，以及基于非光滑max的ReLU型读出 ψ(θ̃) = max(θ̃, ε)，即投影到非负锥上（锥投影是 ε=0 的情况）。每一种都通过构造强制 θ ⪰ 0，并且每一种都引入了一个链式法则前置因子 ψ′(θ̃)，该因子乘以每一个下游梯度，并在参数空间的一个扩展区域上坍缩：softplus前置因子 ψ′(θ̃) 在 θ̃ → -∞ 时平滑消失（Hoedt和Klambauer，2023），而ReLU型前置因子在门控集上恒为零，并以可微性为代价匹配PGD。随机梯度下降仅在时间尺度上逃逸该衰减区域，该尺度随逆噪声水平呈指数增长（推论1）——即Kramers–Arrhenius（Kramers，1940；Hänggi等人，1990；Xie等人，2021）机制，从根本上慢于多项式级的经典PGD速率。现有的补救措施只治标不治本：专门的初始化方案（Hoedt和Klambauer，2023）各自驯服单个失效模式；带有正性约束的普通交替方向乘子法（ADMM）（Boyd等人，2011）通过松弛量上的闭式近端算子强制执行约束，但将原始块留在了没有数据条件重参数化的状态，其刚性罚极限 ρ → ∞ 在结构上会简化为锥上的PGD。受全波形反演（FWI）（Symes等人，2020；van Leeuwen和Herrmann，2013）中参数扩展“提升”方法的启发——其中将一个刚性变量重新投射到更大的空间中，以将困难的非凸问题转化为更平滑的问题——我们提出了**提升（lift）方法**。我们不直接约束层间权重，而是训练一个无约束的超网络（Ha等人，2017），该超网络根据输入批次发射这些权重。由于批次在每一步都会被重新采样，发射出的权重会随着训练进行而波动——这是训练动力学中的额外随机性来源。与普通的mini-batch梯度噪声不同，这种波动不会被使直接softplus停滞的梯度衰减所抑制，并且它会在训练原本会停滞的区域周围使损失景观变得平缓。我们将该效应追溯至该构造的三个结构性要素——一个可学习的**松弛量**、一个**批次条件化的主体**以及它们之间的随机耦合——并证明没有一个是可以省略的（定理1）。在一维对数凹能量基模型（EBM）和21维表格基准上的凸势归一化流中，提升方法达到了比PGD和直接softplus都更低的测试损失，将受平台限制的训练轨迹转变为山谷下降轨迹（图1）。

### 1.1 贡献

1.  **(1) 提升方法。** 我们提出了公式(1)的超网络发射ICNN参数化方法，该方法为训练动力学引入了一种额外的随机性来源，有效地平滑了读出肩部周围的损失景观。然后我们将该参数化解构为一种分裂变量分解（第3节，第4.2节），这一解释提供了提升方法的条件化优势，而无需援引ADMM收敛分析所需的速率理论假设。
2.  **(2) 三个联合必要的结构性要素。** 我们识别了提升方法条件化优势的三个要素——一个具有单位雅可比矩阵的**松弛量**、一个**数据条件化的主体**以及它们之间非零的**交叉协方差**——并证明每个要素对于一个操作可测的交叉协方差估计量都是必要的（定理1）。一个隐含的强凸化结果将交叉协方差与损失景观上的附加曲率模量联系起来，其范围限定在随机读出机制内。
3.  **(3) 跨两个ICNN范式的实验证据。** 在一维对数凹EBM目标上，一个四架构消融实验隔离了每个要素，一个30次随机种子的配对扫描限定了提升方法在对数凹族上的分布改进范围（第5.1节）。在21维表格目标以及用于凸势流的两维合成目标上，提升方法相比于直接softplus改善了测试损失，并产生了可测量的更平滑训练轨迹和更好条件化的损失景观几何（第5.2节，第5.3节）。

论文的其余部分安排如下。我们首先命名了激发提升方法的链式法则衰减病理（第2节），然后介绍了解决该问题的松弛量加批次摘要重参数化方法（第3节），并将条件化优势分解为三个结构性要素（第4节）。随后，我们将该构造与普通ADMM加正性约束进行对比（第4.2节），并在对数凹EBM训练和凸势流估计上展示实验证据（第5节）。PGD基线的消融实验（第6节）以及范围和相关工作的讨论（第7节）在结论（第8节）之前结束论证。

## 2 正性重参数化的ICNN在读出肩部衰减梯度

在介绍提升方法之前，我们命名了它所要解决的结构性病理：一个正性读出 ψ，其链式法则前置因子 ψ′ 在参数空间的一个扩展区域上坍缩，从而困住随机梯度下降（SGD）。第2.1节确定了ICNN训练设置；第2.2节定义了读出肩部；第2.3节证实两个读出族——光滑的softplus和非光滑的ReLU型——在那里统一失效。

### 2.1 ICNN训练与正性约束

我们通过数据驱动损失 L 上的SGD来训练ICNN参数化模型（Amos等人，2017）。运行示例是单组分ICNN-EBM，p_θ(x) ∝ exp(-E_θ(x))，其中E_θ是一个输入凸神经网络，其凸非递减激活函数和非负的层间权重 θ_l ⪰ 0 共同强制E_θ的输入凸性；凸势流、PCP-Map和ICNN参数化的最优传输共享相同的结构。正性通过一个单调非负读出 ψ: ℝ → ℝ_≥₀ 来强制执行，θ = ψ(θ̃)，其中 θ̃ ∈ ℝᵈ 无约束。典型实例是 ψ = softplus，其导数 ψ′(θ̃) ∈ (0,1) 通过链式法则乘以每一个下游梯度。在整篇文章中，θ ∈ ℝᵈ 表示正性约束所应用的扁平化层间权重向量；ICNN的其他参数（每层偏置和任何无约束权重）被吸收进函数 E_θ 中，但不单独分析，因为提升机制仅作用于约束权重。

### 2.2 读出肩部

轨迹 ψ′(θ̃) = σ_s ≪ 1 是宽度为 σ_s 的**softplus肩部**（图2）。该肩部是参数空间的一个扩展区域，而不是一个薄集：在 σ_s=0.05 时，它包含整个半线 θ̃ ≲ -3，因此进入它的迭代有空间徘徊。一旦进入肩部，链式法则前置因子 ψ′(θ̃) 在权重幅度上呈指数级小，而SGD仅在时间尺度上逃逸，该尺度随逆噪声水平呈指数增长——即第4.1节推论1所量化的Kramers–Arrhenius机制（Kramers，1940；Hänggi等人，1990）。相同的指数逃逸时间也出现在其他标准正性重参数化中（注1）。

请参阅图注

图 2: Softplus肩部：参数空间的一个扩展区域。