🔢 mathematics

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

This paper introduces BROS, a memory-efficient single-loop stochastic bilevel optimization method that utilizes randomized subspaces and a Rademacher bi-probe correction to achieve the same $\mathcal O(\varepsilon^{-2})$ convergence rate as exact methods while significantly reducing peak memory usage.

Original authors: Hengrui Zhang, Boao Kong, Engao Zhang, Kun Yuan

Published 2026-05-12

📖 4 min read🧠 Deep dive

CC BY 4.0

Original authors: Hengrui Zhang, Boao Kong, Engao Zhang, Kun Yuan

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to tune a very complex machine, like a giant orchestra, to play a perfect song. In the world of AI, this machine is a neural network (the "lower-level" problem), and the "tuning knobs" you are adjusting are the hyperparameters (the "upper-level" problem).

The challenge is that to know which way to turn the knobs, you first have to listen to the orchestra, figure out exactly how every single musician is playing, and then calculate how changing a knob would change the entire song. This is called Bilevel Optimization.

The Problem: The "Memory" Bottleneck

The paper explains that for massive modern AI models (which have billions of parameters), trying to calculate these tuning directions all at once is like trying to carry a library in your backpack.

The Old Way: Existing methods try to calculate the perfect direction by keeping track of every single note and instrument simultaneously. This requires so much computer memory (RAM) that it crashes on large models.
The "Surrogate" Way: Other methods try to cheat by simplifying the math to save memory, but they end up giving you a slightly wrong direction, leading to a worse song.

The Solution: BROS (Bias-Corrected Randomized Subspaces)

The authors propose a new method called BROS. Here is how it works, using a simple analogy:

1. The "Spotlight" Strategy (Randomized Subspaces)

Instead of trying to listen to the entire orchestra at once (which is too heavy), BROS uses a spotlight.

It randomly picks a small group of musicians (a "subspace") to focus on for a moment.
It calculates how to tune the knobs based only on this small group.
Because it's only looking at a small group, it uses much less memory (up to 45% less in their tests).

2. The "Magic Mirror" Correction (Rademacher Bi-Probe)

Here is the tricky part: If you only listen to a small group, your calculation of the whole orchestra will be biased (wrong). It's like judging a whole symphony based only on the violins; you might miss the drums.

Most methods that use this "spotlight" approach just accept this error, which ruins the final result.
BROS's Secret Sauce: It uses a clever mathematical trick called a Rademacher bi-probe. Think of this as a "magic mirror" or a "correction lens."
After looking at the small group, BROS asks a few specific, random questions (using random +1 and -1 signals) to figure out exactly how the spotlight distorted the view.
It then mathematically "undoes" that distortion.

The Result: Best of Both Worlds

Because of this correction, BROS gets the best of both worlds:

Low Memory: It runs on smaller computers because it only processes small chunks of the model at a time.
High Accuracy: Because it corrects the bias, it finds the exact same perfect tuning direction as the heavy, memory-hungry methods. It doesn't compromise on quality.

What They Tested

The authors tested BROS on four real-world AI tasks:

Cleaning Messy Data: Fixing AI training data that has wrong labels (like correcting a student's homework that was graded incorrectly).
Mixing Data: Figuring out the perfect recipe of different data sources to train a language model.
Learning Representations: Teaching an AI how to "see" images better.
Reweighting Samples: Deciding which specific images are most important for an AI to learn from.

In all these tests, BROS used significantly less memory (reducing peak memory by up to 45%) while achieving nearly identical performance to the heavy, memory-intensive methods.

In a Nutshell

BROS is like a smart conductor who doesn't need to hear every single instrument in the orchestra at once to know how to tune the music. Instead, they listen to a small section, use a special mathematical trick to correct for what they missed, and end up conducting the whole orchestra perfectly—without needing a massive, expensive sound system.

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

Problem Statement
Stochastic bilevel optimization (SBO) is a fundamental framework for deep learning applications such as hyperparameter optimization, data reweighting, representation learning, and hyper-data cleaning. The standard formulation involves an upper-level variable $x$ and a lower-level variable $Y$ (typically the parameters of a neural network), where the upper objective depends on the solution of the lower optimization problem.

Existing exact single-loop SBO methods (e.g., MA-SOBA) provide strong convergence guarantees ( $O(\epsilon^{-2})$ sample complexity) but suffer from severe memory bottlenecks when the lower-level network is large. These methods rely on full-space Hessian-vector products (HVPs) and Jacobian-vector products (JVPs) to maintain auxiliary variables, requiring memory proportional to the full parameter count. Conversely, memory-efficient surrogate methods (e.g., penalty methods, zeroth-order approaches) reduce memory usage but often sacrifice exact convergence guarantees, require stronger assumptions, or optimize a surrogate problem rather than the original bilevel objective.

Methodology: BROS
The authors propose BROS (Bilevel optimization via RandOmized Subspaces), a memory-efficient single-loop SBO method designed for multilayer matrix-parameterized problems. BROS addresses the memory bottleneck by performing lower-level and auxiliary-level updates within randomized low-dimensional subspaces while preserving the exact single-loop hypergradient structure.

The core components of BROS are:

Randomized Subspace Decomposition:
Instead of updating the full lower-level parameters $Y \in \mathbb{R}^{m \times n}$ , BROS samples a scaled Haar projector $P \in \mathbb{R}^{m \times r}$ (where $r \ll m$ ) at each iteration. The lower-level update is computed in the reduced subspace $B \in \mathbb{R}^{r \times n}$ via the projected objective $\tilde{g}(x, B) = g(x, Y + PB)$ . The full-space variable is updated as $Y^{k+1} = Y^k + P B^k$ . This reduces the memory footprint of gradients, activations, and auxiliary directions from $O(mn)$ to $O(rn)$.
Bias-Corrected Auxiliary Updates:
A naive "sketch-then-lift" approach for the auxiliary variable $Z$ (which solves the linear system involving the Hessian) introduces bias because the lifted Hessian action $P \tilde{H} [P^\top Z]$ is not an unbiased estimator of the full-space Hessian action $H[Z]$ . This bias distorts the hypergradient and can lead to non-stationary limits.
To resolve this, BROS introduces a Rademacher bi-probe correction. The algorithm samples independent Rademacher vectors (sign vectors) $u$ and $\xi$ to estimate trace-like and transpose-like nuisance terms induced by the projection. These probes allow the construction of a corrected HVP estimator, denoted as $\widehat{H}[Z]$ , which is conditionally unbiased for the full-space Hessian action without requiring full-space matrix formation.
Single-Loop Recursion:
BROS maintains a coupled single-loop recursion over the upper variable $x$ , lower variable $Y$ , auxiliary variable $Z$ , and a moving-average hypergradient estimator $h$ . The updates utilize the projected lower gradient and the bias-corrected auxiliary direction.

Key Contributions

Memory-Efficient Recursion: BROS decomposes the lower-level and auxiliary recursions into randomized subspaces, significantly reducing the dimensionality of stored gradients, activations, and HVP/JVP computations.
Unbiased Hypergradient Recovery: The paper identifies the specific projection-induced bias in lifted auxiliary Hessian actions and proposes a layerwise Rademacher bi-probe correction. This ensures the estimator remains unbiased for the full multilayer Hessian action without imposing additional structural assumptions (e.g., block-diagonal Hessians).
Theoretical Guarantees: Under standard assumptions (smoothness, strong convexity of the lower level), the authors prove that BROS preserves the $O(\epsilon^{-2})$ sample complexity of exact single-loop methods (specifically MA-SOBA) for finding an $\epsilon$ -stationary point. The convergence rate depends on the projection ratio $\omega_P = \max(m_\ell/r_\ell)$ , recovering the full-space rate when $r=m$ .

Experimental Results
The authors evaluate BROS on four bilevel learning tasks:

MNIST Hyper-Data Cleaning: BROS achieves 84.30% test accuracy, closely matching the full-space MA-SOBA baseline (84.46%) and outperforming penalty-based and zeroth-order methods.
Data-Mixture Learning (280M GPT-style models): BROS matches the performance of MA-SOBA (final log-perplexity 2.8109 vs. 2.8098) while reducing peak proxy-training memory by approximately 44.9% compared to MA-SOBA and 27% compared to Penalty methods.
CIFAR-10 Hyper-Representation Learning: BROS achieves 78.03% final test accuracy, comparable to MA-SOBA (78.57%).
CIFAR-100 ViT Sample Reweighting: BROS achieves the best validation (49.88%) and test (49.05%) accuracy among compared methods, using only 6.42 GB of peak memory compared to over 23 GB for full-space baselines.

Significance and Claims
The paper claims that BROS successfully bridges the gap between memory efficiency and theoretical convergence in stochastic bilevel optimization. Unlike surrogate methods that trade off exactness for efficiency, BROS maintains the exact single-loop hypergradient update structure and its associated $O(\epsilon^{-2})$ convergence order. By utilizing randomized subspaces and a novel bias-correction mechanism, BROS enables bilevel optimization tasks (such as data reweighting and hyperparameter tuning) on large neural networks that would otherwise be infeasible due to memory constraints, reducing peak memory usage by up to 44.9% while closely matching full-space performance.