Research map / Optimization / Optimizers
Memory-Efficient Gradient Unrolling for Large-Scale Bi-level Optimization
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes (FG)^2U, an unbiased forward-gradient-based unrolling method for large-scale bi-level optimization that keeps memory constant in meta-parameter dimension and unroll depth, evaluated on data condensation, meta-learning LM adaptation, and PDE discovery.
Problem
Gradient-based bi-level optimization is hard to scale: reverse gradient unrolling has non-constant memory in unroll depth and parameter dimension, while implicit-function and truncated/Hessian-free approximations introduce bias.
Approach
Combine forward-mode gradient unrolling (tracking Z_t v instead of the full Jacobian) with forward gradient estimation over random directions v, yielding an unbiased stochastic meta-gradient estimate with constant memory, which also extends to zeroth-order finite-difference variants for black-box inner solvers.
Key results
- On data condensation (CIFAR-100, IPC=1), (FG)^2U reaches 8.22% test accuracy versus 3.96% for TRGU and 3.14% for Hessian-Free, and 4.52% for Neumann.
- On CaMeLS online adaptation with GPT2-XL (1.5B), (FG)^2U gets 8.89 EM / 14.42 F1 on StreamingQA versus 6.55 / 11.67 for CaMeLS with RGU.
- For PDE discovery, (FG)^2U-ZO with a spectral numerical solver significantly outperforms (FG)^2U with a PINN solver on both parameter and solution prediction.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use (FG)^2U instead of RGU for data condensation | 3-layer ConvNet data condensation · MNIST, IPC=1/10/50 | testing accuracy (%): 82.44 vs 92.32; 96.12 vs 96.79; 97.01 vs 97.72 | Hurt |
| Use (FG)^2U instead of RGU for data condensation | 3-layer ConvNet data condensation · CIFAR-10, IPC=1/10/50 | testing accuracy (%): 29.37 vs 34.08; 50.10 vs 53.15; 51.98 vs 56.37 | Hurt |
| Use (FG)^2U instead of RGU for data condensation | 3-layer ConvNet data condensation · CIFAR-100, IPC=1/10/50 | testing accuracy (%): 8.22 vs 15.61; 23.38 vs 25.42; 25.84 vs 28.52 | Hurt |
| Replace RGU with (FG)^2U in CaMeLS | CaMeLS with GPT2-XL · StreamingQA | EM: 7.93 -> 8.89 | Helped |
| Replace RGU with (FG)^2U in CaMeLS | CaMeLS with GPT2-XL · StreamingQA | F1: 12.94 -> 14.42 | Helped |
| Replace RGU with (FG)^2U in CaMeLS | CaMeLS with GPT2-XL · SQuAD-Seq | EM: 6.71 -> 7.37 | Helped |
| Replace RGU with (FG)^2U in CaMeLS | CaMeLS with GPT2-XL · SQuAD-Seq | F1: 9.65 -> 10.37 | Helped |
| Increase unrolled depth (ablation) | CaMeLS online adaptation · StreamingQA | F1 | Helped |
| Increase base-model size from DistilGPT2 to GPT2 (ablation) | CaMeLS online adaptation · StreamingQA and SQuAD-Seq | F1 | Helped |
| Use a numerical solver with (FG)^2U-ZO instead of a PINN solver | PDE parameter discovery · Burgers, Allen-Cahn, and KdV | relative L2 error | Helped |
Surprising findings
- The proposed unbiased estimator beats RGU-style reference on some settings only when unroll depth is increased, since RGU is memory-limited; the authors note RGU is an upper bound reference but impractical at scale.
Limitations
- Experiments are small-scale; the largest inner model is GPT-2 (1.5B).
- Convergence rate scales linearly with meta-parameter dimension N divided by sample size b, and forward gradient inherits efficiency issues.
- Black-box/zeroth-order variant (FG)^2U-ZO has limited exploration and introduces approximation bias.
Experimental setup
Datasets: MNIST CIFAR-10 CIFAR-100 StreamingQA SQuAD-Seq Burgers Allen-Cahn KdV
Benchmarks: MNIST CIFAR-10 CIFAR-100 StreamingQA SQuAD-Seq Burgers Allen-Cahn KdV
Metrics: testing accuracy exact match F1 relative L2 error GPU memory usage
Baselines: TRGU Hessian-Free Neumann series RGU CaMeLS + RGU PINN solver
Builds on: forward gradient learning reverse gradient unrolling CaMeLS truncated reverse gradient unrolling performance matching data condensation
Closest papers by meaning
- A Stochastic Approach to Bi-Level Optimization for Hyperparameter Optimization and Meta Learning (AAAI 2025 · similarity 0.89)
- An Accelerated Algorithm for Stochastic Bilevel Optimization under Unbounded Smoothness (NeurIPS 2024 · similarity 0.87)
- Bilevel Optimization under Unbounded Smoothness: A New Algorithm and Convergence Analysis (ICLR 2024 · similarity 0.87)
- BROS: Bias-Corrected Randomized Subspaces for Memory-Efficient Single-Loop Bilevel Optimization (NeurIPS 2026 · similarity 0.87)
- A Nearly Optimal Single Loop Algorithm for Stochastic Bilevel Optimization under Unbounded Smoothness (ICML 2024 · similarity 0.86)
- Moreau Envelope for Nonconvex Bi-Level Optimization: A Single-Loop and Hessian-Free Solution Strategy (ICML 2024 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.