🔢 mathematics

BROS: Bias-Corrected Randomized Subspaces for Memory-Efficient Single-Loop Bilevel Optimization

本文介绍了 BROS，这是一种内存高效的单循环随机双层优化方法，它利用随机子空间和 Rademacher 双探针校正，在显著降低峰值内存占用的同时，实现了与精确方法相同的 $\mathcal O(\varepsilon^{-2})$ 收敛速率。

原作者： Hengrui Zhang, Boao Kong, Engao Zhang, Kun Yuan

发布于 2026-05-12

📖 1 分钟阅读🧠 深度阅读

CC BY 4.0

原作者： Hengrui Zhang, Boao Kong, Engao Zhang, Kun Yuan

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在调校一台极其复杂的机器，就像一支庞大的管弦乐团，要让它演奏出一首完美的乐曲。在人工智能领域，这台机器就是神经网络（“下层”问题），而你正在调节的“调音旋钮”则是超参数（“上层”问题）。

挑战在于，为了知道该朝哪个方向转动旋钮，你首先必须聆听整个乐团，精确分析每一位乐手的演奏方式，然后计算调整某个旋钮会如何改变整首乐曲。这被称为双层优化。

问题：“内存”瓶颈

论文指出，对于拥有数十亿参数的现代大型 AI 模型而言，试图一次性计算出这些调音方向，就像试图把整个图书馆塞进背包里。

旧方法：现有方法试图通过同时追踪每一个音符和每一种乐器来计算完美的方向。这需要海量的计算机内存（RAM），以至于在大型模型上会导致崩溃。
“代理”方法：其他方法试图通过简化数学运算来节省内存，但这最终会给出略微错误的方向，导致乐曲效果变差。

解决方案：BROS（偏差校正随机子空间）

作者提出了一种名为BROS的新方法。其工作原理如下，借助一个简单的类比：

1. “聚光灯”策略（随机子空间）

BROS 不使用试图一次性聆听整个乐团的方法（因为这太沉重），而是使用聚光灯。

它随机挑选一小群乐手（一个“子空间”）进行短暂聚焦。
它仅基于这一小群乐手来计算如何调节旋钮。
由于只关注一小群乐手，它使用的内存要少得多（在其测试中减少了高达 45%）。

2. “魔镜”校正（拉达马赫双探针）

这里是棘手之处：如果你只聆听一小群乐手，你对整个乐团的计算就会产生偏差（错误）。这就像仅凭小提琴部分来评判整部交响乐；你可能会忽略鼓声。

大多数使用这种“聚光灯”方法的技术只是接受这种误差，从而破坏了最终结果。
BROS 的独门秘籍：它使用一种巧妙的数学技巧，称为拉达马赫双探针。你可以将其想象为一面“魔镜”或一个“校正透镜”。
在观察了一小群乐手之后，BROS 会提出几个特定的、随机的问题（利用随机的 +1 和 -1 信号），以精确弄清楚“聚光灯”是如何扭曲视角的。
随后，它在数学上“消除”了这种扭曲。

结果：鱼与熊掌兼得

得益于这种校正，BROS 实现了两全其美：

低内存：因为它每次只处理模型的一小部分，所以可以在较小的计算机上运行。
高准确度：由于它校正了偏差，它能找到与那些庞大且耗内存的方法完全相同的完美调音方向。它没有在质量上做出妥协。

测试内容

作者在四个现实世界的 AI 任务上测试了 BROS：

清理杂乱数据：修复带有错误标签的 AI 训练数据（例如纠正被错误评分的学生作业）。
混合数据：找出训练语言模型所需的完美数据源配方。
学习表征：教导 AI 如何更好地“看”图像。
样本重加权：决定哪些特定图像对 AI 学习最为重要。

在所有这些测试中，BROS 使用的内存显著更少（峰值内存减少了高达 45%），同时实现了与那些庞大且耗内存的方法几乎相同的性能。

nutshell

BROS 就像一位聪明的指挥家，无需一次性听到乐团中的每一件乐器，就能知道如何调校音乐。相反，他聆听一小部分，利用一种特殊的数学技巧来校正他所遗漏的内容，最终完美地指挥整个乐团——而无需庞大且昂贵的音响系统。

技术摘要：BROS：用于内存高效单循环双层优化的偏差校正随机子空间

问题陈述
随机双层优化（SBO）是深度学习应用（如超参数优化、数据重加权、表示学习和超数据清洗）的基础框架。其标准形式涉及上层变量 $x$ 和下层变量 $Y$ （通常是神经网络的参数），其中上层目标函数依赖于下层优化问题的解。

现有的精确单循环 SBO 方法（例如 MA-SOBA）提供了强大的收敛保证（ $O(\epsilon^{-2})$ 样本复杂度），但在下层网络规模较大时，会遭受严重的内存瓶颈。这些方法依赖全空间海森 - 向量积（HVPs）和雅可比 - 向量积（JVPs）来维护辅助变量，其内存需求与全量参数成正比。相反，内存高效的代理方法（例如惩罚方法、零阶方法）虽然降低了内存使用，但往往牺牲了精确的收敛保证，需要更强的假设，或者优化的是代理问题而非原始双层目标。

方法论：BROS
作者提出了BROS（通过随机子空间进行双层优化），这是一种专为多层矩阵参数化问题设计的内存高效单循环 SBO 方法。BROS 通过在随机低维子空间内执行下层和辅助层更新，同时保持精确的单循环超梯度结构，从而解决了内存瓶颈问题。

BROS 的核心组件包括：

随机子空间分解：
BROS 不再更新完整的下层参数 $Y \in \mathbb{R}^{m \times n}$ ，而是在每次迭代中采样一个缩放的海尔投影算子 $P \in \mathbb{R}^{m \times r}$ （其中 $r \ll m$ ）。下层更新通过投影目标函数 $\tilde{g}(x, B) = g(x, Y + PB)$ 在降维子空间 $B \in \mathbb{R}^{r \times n}$ 中计算。全空间变量更新为 $Y^{k+1} = Y^k + P B^k$ 。这将梯度、激活值和辅助方向的内存占用从 $O(mn) $降低到了$ O(rn)$。
偏差校正辅助更新：
对辅助变量 $Z$ （用于求解涉及海森矩阵的线性系统）采用朴素的“先草图后提升”方法会引入偏差，因为提升后的海森作用 $P \tilde{H} [P^\top Z]$ 并非全空间海森作用 $H[Z]$ 的无偏估计量。这种偏差会扭曲超梯度并可能导致非平稳极限。
为解决此问题，BROS 引入了Rademacher 双探针校正。算法采样独立的 Rademacher 向量（符号向量） $u$ 和 $\xi$ ，以估计由投影引起的类迹和类转置的干扰项。这些探针允许构建一个校正后的 HVP 估计量，记为 $\widehat{H}[Z]$ ，该估计量在无需构建全空间矩阵的情况下，对全空间海森作用具有条件无偏性。
单循环递归：
BROS 维持一个关于上层变量 $x$ 、下层变量 $Y$ 、辅助变量 $Z$ 以及移动平均超梯度估计量 $h$ 的耦合单循环递归。更新过程利用投影后的下层梯度和偏差校正后的辅助方向。

主要贡献

内存高效递归： BROS 将下层和辅助递归分解到随机子空间中，显著降低了存储的梯度、激活值以及 HVP/JVP 计算的维度。
无偏超梯度恢复： 论文识别了提升辅助海森作用中特定的投影诱导偏差，并提出了一种逐层 Rademacher 双探针校正。这确保了估计量对全多层海森作用保持无偏，而无需施加额外的结构假设（例如块对角海森矩阵）。
理论保证： 在标准假设下（平滑性、下层强凸性），作者证明了 BROS 保留了精确单循环方法（特别是 MA-SOBA）的 $O(\epsilon^{-2})$ 样本复杂度，用于寻找 $\epsilon$ -平稳点。收敛速率取决于投影比率 $\omega_P = \max(m_\ell/r_\ell)$ ，当 $r=m$ 时恢复全空间速率。

实验结果
作者在四个双层学习任务上评估了 BROS：

MNIST 超数据清洗： BROS 实现了 84.30% 的测试准确率，与全空间 MA-SOBA 基线（84.46%）高度吻合，并优于基于惩罚的方法和零阶方法。
数据混合学习（2.8 亿参数的 GPT 风格模型）： BROS 匹配了 MA-SOBA 的性能（最终对数困惑度 2.8109 对比 2.8098），同时将峰值代理训练内存减少了约44.9%（相较于 MA-SOBA）和 27%（相较于惩罚方法）。
CIFAR-10 超表示学习： BROS 实现了 78.03% 的最终测试准确率，与 MA-SOBA（78.57%）相当。
CIFAR-100 ViT 样本重加权： BROS 在比较的方法中取得了最佳的验证集（49.88%）和测试集（49.05%）准确率，其峰值内存使用量仅为 6.42 GB，而全空间基线则超过 23 GB。

意义与主张
论文主张，BROS 成功弥合了随机双层优化中内存效率与理论收敛之间的差距。与为了效率而牺牲精确性的代理方法不同，BROS 保持了精确的单循环超梯度更新结构及其相关的 $O(\epsilon^{-2})$ 收敛阶。通过利用随机子空间和一种新颖的偏差校正机制，BROS 使得在大型神经网络上进行双层优化任务（如数据重加权和超参数调整）成为可能，而这些任务此前因内存限制而不可行；它在将峰值内存使用量降低高达 44.9% 的同时，紧密匹配了全空间性能。