Select-then-differentiate: Solving Bilevel Optimization with Manifold Lower-level Solution Sets¶

会议: NeurIPS 2026
arXiv: 2605.09209
OpenReview: https://openreview.net/forum?id=NLXLf4kGZu
领域: 优化/理论
关键词: 双层优化, 流形解集, 乐观选择, 超梯度, 数据重加权

一句话总结¶

针对过参数化与非凸神经网络中下层解集呈现连通流形且非单值的挑战，本文提出“先选择后微分”（Select-then-differentiate）框架，证明了乐观选择唯一性是超目标可微的充要条件并推导出广义伪逆超梯度公式，进而构建了基于 Gibbs 采样与 Best-of-N 机制的 HG-MS 算法及其内在维度复杂度界。

研究背景与动机¶

现代机器学习中诸多前沿任务——从超参数优化、元学习、神经架构搜索，到大语言模型（LLM）的数据源重加权、安全微调和偏好对齐——在数学本质上均可建模为双层优化（Bilevel Optimization）问题。传统双层优化理论与算法建立在高度理想化的单值假设之上，即假定下层优化问题存在唯一的局部或全局极小值点，进而利用隐函数定理直接计算外层参数的超梯度。然而，在以深度神经网络为代表的过参数化非凸学习范式中，该假设被普遍打破：网络的损失曲面通常包含大量非孤立的全局极小值，这些解集在参数空间中构成连续连通的低维嵌入流形。在下层解映射呈现多值流形结构时，外层目标函数对不同下层最优解的评价存在巨大差异，使得外层超目标（Hyper-objective）严重依赖于下层的解选择机制。

在最符合实际需求的“乐观双层优化”（Optimistic Bilevel Optimization）设定下，系统要求在上层准则最好的前提下从多值下层解集中挑选最优解。然而，这种极小化内部的极小化运算会带来本质性的数学病态：当下层解集发生连续摄动时，最优下层解的选择可能会产生剧烈跳跃，导致外层超目标出现非光滑甚至不可微的奇点。以往拓宽单值假设的工作多依赖于下层目标的全局凸性或直接退化到次微分分析，难以指导实际的超梯度计算；而过参数化流形下的局部光滑性条件、超梯度闭式表达及其优化算法的理论收敛保证始终处于未解状态。

本文敏锐地洞察到：超目标不可微的根本根源不在于下层解集本身是否为单值点集，而在于上层最优选择（Optimistic Selection）是否唯一。只要在下层局部满足 Polyak–Łojasiewicz 流形几何条件（PŁ°）的前提下，上层目标在下层流形上的最优点是唯一的，超目标就依然具备严格的连续可微性。本文的核心 idea 是：确立“先选择后微分”（Select-then-differentiate）原则，证明乐观选择的唯一性与内在非退化性足以保证超目标的局部光滑可微，导出基于正规空间投影与摩尔-彭罗斯伪逆的显式超梯度解析式，并据此设计出首个复杂度仅受下层流形内在维度支配的 HG-MS 优化算法。

方法详解¶

整体框架¶

本文针对下层解集 $\mathcal{S}(\theta) := \arg\min_{x \in \mathbb{R}^d} g(\theta, x)$ 构成嵌入子流形的乐观双层优化问题 $\min_{\theta \in \Theta} F(\theta) := \min_{x \in \mathcal{S}(\theta)} f(\theta, x)$，建立了一套从理论正则性分析到实用算法实现的完整体系。整体算法逻辑遵循严格的两阶段流程：外层迭代首先在下层流形邻域内通过采样与上层评估挑选出最佳候选点（Minima Selection），随后在选定的极小值点处通过正则化共轭梯度求解法向隐式方程，计算出修正的伪逆超梯度以更新外层参数 $\theta$。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    In["当前外层参数 θ_t 与目标模型"] --> S1["Gibbs 测度近似与流形并行采样<br/>利用 ULA 沿低能级流形管道生成 N 个候选解"]
    S1 --> S2["乐观极小值选择 (BoN Selection)<br/>在上层目标 f(θ_t, ·) 下选取最优候选点 x_hat"]
    S2 --> S3["正则化伪逆超梯度求解<br/>岭正则化共轭梯度 (CG) 逼近法向敏感性逆算子"]
    S3 --> S4["超梯度投影更新<br/>沿 h_t 负方向更新外层参数并投影至参数空间 Θ"]
    S4 --> Out["输出最优外层混合权重 θ 与收敛模型"]

关键设计¶

1. 乐观选择唯一性与流形伪逆超梯度定理：突破单值假设的正则性理论

经典隐式微分要求下层海森矩阵 $\nabla_{xx}^2 g(\theta, x)$ 全秩可逆，但在流形解集设定下，流形切空间 $\mathcal{T}_x^\theta$ 恰好构成海森矩阵的零空间（$\ker \nabla_{xx}^2 g = \mathcal{T}_x^\theta$），海森矩阵在全空间严格奇异。作者证明，在下层目标满足局部 PŁ° 条件（保证解集为紧致光滑子流形且法向曲率一致正定）时，只要乐观选择集 $\mathcal{O}(\theta) := \arg\min_{x \in \mathcal{S}(\theta)} f(\theta, x)$ 为单元素集 $\{x^\star(\theta)\}$，外层超目标 $F(\theta)$ 即局部严格 $C^1$ 可微。在此基础上，利用流形正规坐标图与 Danskin 型包络定理，证明了切向梯度自动消除，最终超梯度解析式可显式表示为： $$ \nabla F(\theta) = \nabla_\theta f(\theta, x^\star(\theta)) - \nabla_{\theta x}^2 g(\theta, x^\star(\theta)) \big[\nabla_{xx}^2 g(\theta, x^\star(\theta))\big]^\dagger \nabla_x f(\theta, x^\star(\theta)) $$ 其中 $[\nabla_{xx}^2 g]^\dagger = \mathbf{N}_\theta (\mathbf{N}_\theta^\top \nabla_{xx}^2 g \mathbf{N}_\theta)^{-1} \mathbf{N}_\theta^\top$ 恰好是海森矩阵在法向空间 $\mathcal{N}_x^\theta$ 上的摩尔-彭罗斯伪逆。进一步地，论文揭示了若乐观极小值点在流形上的黎曼海森矩阵正定（内在非退化），则 $\nabla F(\theta)$ 局部利普希茨连续；若非退化性不满足，即使选择唯一，超梯度也可能失去任意正阶的 Hölder 连续性。

2. Gibbs 测度与 Best-of-N 乐观解筛选：摆脱环境维度的内在搜索机制

在算法层面，求解约束优化 $\min_{x \in \mathcal{S}(\theta)} f(\theta, x)$ 面临下层流形为隐式非凸曲面的困境。HG-MS 引入温度参数 $\lambda > 0$ 的 Gibbs 测度 $\mu_\theta^\lambda(\mathrm{d}x) \propto \exp(-g(\theta, x)/\lambda) \mathrm{d}x$。当 $\lambda \to 0$ 时，该测度以高概率严格浓缩在以 $\mathcal{S}(\theta)$ 为中心的细管区域（Tube Neighborhood）内。算法并行运行 $N$ 条不调整朗之万动力学（ULA）链抽取样本 $\{X_1, \dots, X_N\}$，并执行 Best-of-N (BoN) 操作挑选出上层损失最小的样本 $\hat{x}_N^\lambda(\theta) = \arg\min_{X_i} f(\theta, X_i)$。理论严格证明该选择器的均方误差满足： $$ \mathbb{E}\big[|x^\star(\theta) - \hat{x}_N^\lambda(\theta)|^2\big] = \mathcal{O}\big(N^{-1/\texttt{k}} + \lambda^{1/2}\big) $$ 极为关键的是，收敛阶数 $N^{-1/\texttt{k}}$ 中的 $\texttt{k}$ 是下层解流形的内在维度（Intrinsic Dimension），而非模型参数空间的环境维度 $d$。这使得该机制能够在过参数化深度学习中免受千亿维度的“维度灾难”拖累。

3. 岭正则化共轭梯度隐式步：抵御离轨扰动的数值稳定性保障

采样得到的候选点 $\hat{x}_N^\lambda(\theta)$ 往往位于流形邻域而非严格落在流形内部，导致海森矩阵的切向零特征值发生未知扰动变为近零小特征值，直接求逆极易引发严重的数值爆炸。HG-MS 引入带有自适应岭参数 $\gamma > 0$ 的线性系统 $(\nabla_{xx}^2 g + \gamma I) \tilde{v} = \nabla_x f$，利用仅需海森-向量积（HVP）的共轭梯度法（CG）迭代求解至容差 $\eta_t$。论文构建了第二矩误差分解框架，将超梯度总误差界定为候选点偏差、CG 残差、岭偏差以及流形管外逃逸概率四项之和，严格给出了保证外层优化达到 $\epsilon$-近似一阶稳定点所需的超参数调度与总神谕复杂度 $\widetilde{\mathcal{O}}(\epsilon^{-42-8\texttt{k}})$。

损失函数 / 训练策略¶

在外层参数更新中，HG-MS 采用投影梯度下降法： $$ \theta_{t+1} = \mathrm{Proj}_\Theta\big(\theta_t - \alpha_t \widehat{h}_t\big) $$ 其中超梯度估计量为 $\widehat{h}_t = \nabla_\theta f(\theta_t, \hat{x}_t) - \nabla_{\theta x}^2 g(\theta_t, \hat{x}_t) v_t$，$v_t$ 为经有界半径投影（截断保护）的共轭梯度求解向量。在面向实际大语言模型（LLM）多数据源配比优化时，为契合有限训练算力预算，算法派生出轻量化的多分支探索架构：用小批次 SGD 或 AdamW 训练多条独立的候选分支以模拟下层流形邻域分布，在验证集上硬筛选出损失最低的候选分支，随后在反向传播中计算超梯度以调整各数据源的 softmax 权重 $\omega = \mathrm{softmax}(\theta)$。

实验关键数据¶

主实验¶

针对真实场景下的 LLM 跨领域数据合成微调任务，论文在十源数学推理组合基准上对 Llama-3-8B、Qwen-2-7B、Gemma-2-9B 以及超大规模 Qwen2.5-32B 进行了对齐预算的 LoRA 微调评估。所有迭代类重加权方法均匹配相同的单卡 H100 训练时长预算（Qwen-2-7B 约 9 小时，Llama-3-8B 约 12 小时，Gemma-2-9B 约 15 小时，Qwen2.5-32B 固定 15 小时），评测指标为小学应用题 GSM8K 及竞赛级数学题 MATH 的最终测试准确率（%）。

模型基座	优化方法	外层重加权预算	GSM8K 准确率 (%)	MATH 准确率 (%)
Llama-3-8B	Uniform Weighting	0h (等权采样)	60.06	30.24
	RHO-LOSS	离线计算	65.43	25.89
	LESS	离线计算	74.45	29.93
	ScaleBiO	~12h	66.19	30.58
	HG-baseline (单分支超梯度)	~12h	73.77	30.49
	HG-MS (本文方法)	~12h	74.85	30.86
Qwen-2-7B	Uniform Weighting	0h (等权采样)	60.03	30.26
	RHO-LOSS	离线计算	63.15	33.40
	LESS	离线计算	65.28	29.24
	ScaleBiO	~9h	75.74	36.93
	HG-baseline (单分支超梯度)	~9h	68.46	34.27
	HG-MS (本文方法)	~9h	82.03	38.98
Gemma-2-9B	Uniform Weighting	0h (等权采样)	78.08	37.29
	RHO-LOSS	离线计算	78.47	37.35
	LESS	离线计算	79.28	38.00
	ScaleBiO	~15h	79.62	37.98
	HG-baseline (单分支超梯度)	~15h	80.14	38.23
	HG-MS (本文方法)	~15h	80.86	38.48
Qwen2.5-32B	Uniform Weighting	0h (等权采样)	78.19	43.30
	ScaleBiO	15h	88.40	51.29
	HG-baseline (单分支超梯度)	15h	86.81	46.82
	HG-MS (本文方法)	15h	92.58	53.32

(注：在 32B 大模型规模下，RHO-LOSS 与 LESS 均由于显存爆炸出现 OOM 崩溃；HG-MS 相比强基线 ScaleBiO 在 GSM8K 上大幅超出 4.18 个百分点，MATH 超出 2.03 个百分点。)

消融实验¶

在严格受控的双数据源 GPT-2（124M 全参数可训练）重加权实验中，作者设置了去噪微调（干净 Alpaca 与噪声注入数据，目标权重 $\omega^\star=(1,0)$）和多语言配比（Alpaca 中文与英文，目标配比 $\omega^\star=(0.6, 0.4)$）两项任务，验证显式极小值选择对超参数学习轨迹与验证性能的直接影响。

任务场景	优化算法	习得权重 $\omega$	权重误差 $\\|\omega - \omega^\star\\|_1$	验证集损失	测试集损失
数据去噪任务	Uniform	(0.500, 0.500)	1.000	2.571	2.610
($\omega^\star = (1, 0)$)	ScaleBiO	(0.922, 0.078)	0.156	2.336	2.443
	HG-baseline	(0.930, 0.070)	0.140	2.427	2.454
	HG-MS (本文)	(0.962, 0.038)	0.077	2.403	2.437
多语言配比任务	Uniform	(0.500, 0.500)	0.200	2.430	2.407
($\omega^\star = (0.6, 0.4)$)	ScaleBiO	(0.623, 0.377)	0.047	2.339	2.371
	HG-baseline	(0.540, 0.460)	0.119	2.327	2.307
	HG-MS (本文)	(0.514, 0.486)	0.172	2.323	2.298

关键发现¶

极小值筛选的泛化优势：在受控多语言任务中，ScaleBiO 学得的混合权重虽然在数值上更接近人为设定的目标配比 $\omega^\star$，但 HG-MS 却在验证集和测试集上取得了更低的交叉熵损失。这揭示出双层优化的本质在于寻优上层泛化性能，而非单纯复原预设标签比例；显式筛选能够有效驱使模型避开泛化较差的局部流形点。
单分支与多分支的决定性差距：在所有尺寸的 LLM 数学推理评测中，HG-MS 均全面压制单分支基准 HG-baseline（如在 Qwen-2-7B 上 GSM8K 从 68.46% 跃升至 82.03%，净增 13.57 个百分点）。这有力证实：在高度复杂的损失曲面中，随机选择的单个下层解通常距离真正的乐观最优点相去甚远，从而沿途产生带有严重偏差的伪超梯度。
大模型尺度下的可行性与鲁棒性：当参数量扩展至 32B 时，基于影响函数或参考模型的传统过滤算法（LESS、RHO-LOSS）均受限于昂贵的逆海森存储与前向推理而陷入 OOM 显存崩溃；HG-MS 采用轻量化分支探索与基于 HVP 的共轭梯度隐式步，具备极高的显存适配度与伸缩性。

亮点与洞察¶

充要的正则性分水岭：论文不仅给出了保证超目标 $C^1$ 可微与局部利普希茨光滑的充分条件，更构造了精巧的反例（Example 3.1 与 Example 3.5）证明：非唯一选择会在任意小邻域诱发可数无穷个不可微折点；而缺乏流形内在黎曼非退化性则会彻底摧毁超梯度的任何正阶 Hölder 连续性。
内在维度的理论突破：在非凸双层优化领域，理论界长久以来受制于环境参数维度 $d$ 的指数爆炸困扰。HG-MS 借助 Gibbs 测度的流形局部测度浓缩特性，证明了 Best-of-N 的逼近误差指数仅由流形内在维度 $\texttt{k}$ 决定，首次为过参数化网络中的双层优化建立了维度解耦的复杂度保障。
可复用的通用设计模式：对于涉及深度模型内循环的多任务学习、安全对齐以及 Prompt 自动化调优，均可复用“多随机种子分支前瞻 + 上层泛化准则筛选 + 隐式共轭梯度回传”的 Select-then-differentiate 范式，无需显式保存庞大的全轨迹计算图。

局限与展望¶

理论采样与实用分支的差距：理论收敛性证明基于可严格调控的 ULA 连续马尔可夫链采样，但受制于工业界实际算力，大模型实验中不得不采用少量前向优化分支（如 4 到 8 条单步或短步 SGD/AdamW 分支）作为代用品，理论的渐近覆盖与实际分支经验分布之间尚缺乏紧密的非渐近理论衔接。
内在维度依赖的最坏情况上界：神谕复杂度中包含了维度无关的高阶常数项 $\mathcal{O}(\epsilon^{-42-8\texttt{k}})$，虽然内在维度依赖是全局非凸极小化不可避免的本质代价，但常数项指数偏大，主要源于各阶段松弛界的串联放大，仍有待未来通过更紧密的耦合分析加以压缩。
多最优解冲突的平局处理：当上层目标在下层流形上存在天然对称性或平局时，唯一性条件被破坏。探索集值 Clarke 次微分或基于近端平滑化（Proximal Smoothing）的松弛化超梯度机制是未来重要的理论拓展方向。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次建立流形解集下乐观双层优化的可微性理论，给出伪逆超梯度公式与内在维度复杂度。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖受控双源任务、MNIST基准，以及涵盖 7B 至 32B 规模的 LLM 顶尖数学推理与指令微调实证。
写作质量: ⭐⭐⭐⭐⭐ 数学推导严丝合缝，几何直觉与图解清晰透彻，理论成果与算法实践衔接自然。
价值: ⭐⭐⭐⭐⭐ 为过参数化深度学习中日益普及的双层优化算法提供了坚实可靠的微积分理论支柱与工程范式。