Adaptive Bregman Proximal Stochastic Gradient with a Stabilized Barzilai--Borwein Step Size
Chenhan Jin, Shengze Xu, Binghui Xie, Kaiwen Zhou, Fan Jia, James Cheng, Tieyong Zeng
The Chinese University of Hong Kong · University of Utah · Beijing Normal-Hong Kong Baptist University · Guangzhou Nanfang College
math.OC, cs.LG
Submitted: 2026-08-12
Updated: 2026-08-13
License: http://creativecommons.org/licenses/by/4.0/
Importance score: 75/100
The gist: The paper addresses the finite-sum composite optimization problem on a finite-dimensional normed vector space E: min F(x):= f(x) + h(x), where f(x):= (1/n) Σ fi(x), with h being a proper convex
Terminology
Summary
The paper addresses the finite-sum composite optimization problem on a finite-dimensional normed vector space E:
min F(x):= f(x) + h(x), where f(x):= (1/n) Σ fi(x),
with h being a proper convex regularizer whose proximal subproblem is tractable. This model includes regularized empirical risk minimization, constrained learning, and structured estimation.
The authors note that "Bregman proximal stochastic gradient (BPSG) methods bring variance-reduced composite optimization to objectives whose geometry is poorly captured by Euclidean smoothness. Their performance, however, remains sensitive to the step size: raw stochastic curvature estimates can fluctuate sharply, whereas line searches add repeated proximal evaluations."
The central question is: Can a line-search-free BPSG method adapt to stochastic curvature while preserving convergence in non-Euclidean composite optimization?
The paper introduces Ada-BPSG, a line-search-free BPSG method with three main contributions:
-
A stable adaptive step from SAGA information: The method aggregates recent component-wise secant pairs by a mediant (denominator-weighted mean of local BB ratios), which
suppresses ratios associated with nearly singular curvature estimates.
Clipping and a monotone update turn the raw candidate into a bounded, line-search-free step-size sequence. -
A convergence analysis aligned with the algorithm: The analysis is carried out in a general finite-dimensional normed space, relying only on the dual norm, dual-norm Young inequalities, and the three-point Bregman inequality—never invoking Euclidean polarization identity or norm decomposability.
-
Experiments that test the design mechanism: Including a curvature diagnostic, logistic regression tests, a simplex-constrained Poisson inverse problem under the entropy kernel, real hyperspectral unmixing data, and sparse nonnegative matrix factorization.
The method builds on BPSG-SAGA by adding a curvature-adaptive scalar. The construction uses the classical BB secant ratio, identifies why arithmetic averages of stochastic ratios remain unstable, and replaces that average with a safeguarded mediant computed from information already stored by SAGA.
The raw adaptive candidate is:
η̂ k BB = (1/α) · ξ k / (δ k + θξ k) if ξ k > 0, else η k-1,
where ξ k = Σ ⟨y i,t, s i,t⟩ and δ k = Σ ∥y i,t∥2∗ are window sums of secant pairs from the SAGA table. The safeguard maps this to q k = clip[η min,η max](η̂ k BB), and the monotone update η k = max η k-1, q k preserves the interval and supplies monotonicity for the convex Bregman telescope.
For linear-prediction losses f i(x) = Ψ i(⟨a i, x⟩), the implementation reduces table memory from O(nd) to O(n) by storing only scalar predictions ⟨a i, φ i k⟩.
Convex analysis: Under Assumption 3.4 (including relative smoothness, component local smoothness, and a convex self-bounding condition), with the safeguard satisfying η max ≤ (L̄ + 32σ b M/b)-1, the paper proves an O(n/K) ergodic rate:
E[F(x̄ K)] - F(x*) ≤ (1/K)[(n/(4b) + 1/2)(F(x 0) - F(x*)) + (1/η min)D ψ(x*, x 0)].
Restarted linear rate: Under relative quadratic growth F(x) - F(x*) ≥ μD ψ(x*, x), restarting the averaged convex guarantee yields geometric convergence: after τ restart rounds, E[F(x(τ))] - F(x*) ≤ 2-τ(F(x(0)) - F(x*)).
Nonconvex analysis: For possibly nonconvex f with F bounded below, the paper proves an O(1/K) bound on the expected squared Bregman proximal residual:
E∥R R∥2 ≤ 2(F(x 0) - F inf)/(Kη max),
where R R is the normalized full-gradient proximal displacement.
The paper emphasizes that "the admissible cap η max ≤ (L̄ + 32σ b M/b)-1 encodes exactly the smoothness constants L̄ and M that a well-tuned fixed step would use. Consequently Theorem 3.6 does not improve on the O(n/K) rate of fixed-step BPSG-SAGA; it certifies that adapting within [η min, η max] does not degrade that rate."
The analysis is conducted entirely in the general normed space E: "it uses only the dual norm, dual-norm Young inequalities, and the three-point identity of Lemma A.1, and at no point invokes the Euclidean polarization identity... or the norm decomposability on which the Euclidean adaptive-BB analyses rely."
Euclidean logistic regression: On mushrooms, ijcnn1, w8a, and covtype datasets, SAGA-BB (the Euclidean specialization) reaches the smallest plotted objective gaps with the fewest effective passes on all four datasets.
In initial-step sensitivity tests, SAGA-BB maintains a low gradient norm across the entire grid on every dataset,
while vanilla VR methods vary by several orders of magnitude across the grid and can become unstable at its extremes.
Simplex-constrained Poisson inverse problem: Under the entropy kernel with l1/l∞ geometry, Ada-BPSG decays steadily to 1.3 × 10-5 along the O(n/K) reference slope, exhibiting the convex-rate behavior predicted by Theorem 3.6 in a bona fide non-Euclidean instance.
The fixed-step baselines remain near 10-2
because worst-case M is inflated by the log singularity. Ada-BPSG improves on the best a-priori baseline by more than two orders of magnitude.
Hyperspectral unmixing (Samson scene): With real measured data, Ada-BPSG drives the gap to 2.2 × 10-6—more than two orders of magnitude below the best a-priori baseline—tracking the O(n/K) reference slope.
The adaptive rule stays flat at ≈ 2 × 10-6 across five orders of magnitude of initial step.
Sparse NMF: On the ORL dataset with quartic kernel, Ada-BPSG descends faster than the stochastic BPSG variants and remains slightly below the two extrapolated BPSGE curves throughout the four sparsity regimes.
The paper notes: "The current theory analyzes the clipped sequence on a deterministic bounded region with a strongly convex kernel; the admissible safeguard depends on L̄ and the local component constant M. General finite-sum models also retain the O(nd) SAGA table, although linear-prediction losses admit the O(n) implementation." Extending the stabilized candidate to memory-light SVRG, SARAH, or SPIDER estimators, and deriving safeguards under weaker kernel and localization conditions, are identified as natural next steps.
Improvements for AI systems
Based on the paper, here are specific improvements to AI systems:
1. Adaptive step-size control for non-Euclidean optimization
-
Replace fixed or line-search-based step sizes in stochastic gradient methods with the safeguarded mediant-based adaptive rule (Ada-BPSG)
-
The improved system automatically adjusts step sizes using secant information from SAGA, suppressing unstable curvature estimates
-
This enables training on manifolds and non-Euclidean geometries (e.g., simplex constraints, entropy kernels) without manual tuning or expensive line searches
2. Robust optimization for ill-conditioned problems
-
Use the clipping and monotone update mechanism to prevent step-size explosion when curvature estimates are near-singular
-
The improved system maintains convergence guarantees (O(n/K) convex, O(1/K) nonconvex) even when raw stochastic curvature fluctuates sharply
-
This is particularly valuable for problems with log-singularities or extreme condition numbers where fixed-step methods fail
3. Memory-efficient implementation for linear-prediction models
-
For losses of the form f i(x) = Ψ i(⟨a i, x⟩), store only scalar predictions instead of full gradient tables
-
This reduces memory from O(nd) to O(n), enabling application to high-dimensional problems (e.g., large-scale logistic regression, sparse NMF) where memory is the bottleneck
4. Automatic initialization robustness
-
The adaptive rule makes performance nearly invariant to initial step-size choice (experiments show stability across five orders of magnitude)
-
The improved system eliminates the need for grid search over initial learning rates, which is a major practical bottleneck in deep learning and regularized ERM
5. Nonconvex optimization with guaranteed residual decay
-
Apply the method to nonconvex objectives (e.g., NMF, neural network training) with the O(1/K) bound on Bregman proximal residuals
-
The improved system provides a principled stopping criterion based on the proximal residual, enabling reliable convergence detection
6. Restartable convergence acceleration
-
Use the restarted linear rate under relative quadratic growth to achieve geometric convergence for strongly convex-like problems
-
The improved system automatically detects when to restart (e.g., every K iterations) to accelerate convergence without additional hyperparameters
7. Unified framework for constrained and structured problems
-
Handle composite objectives with tractable proximal operators (e.g., l1 regularization, simplex constraints, nonnegativity)
-
The improved system can solve constrained learning, structured estimation, and inverse problems in a single algorithmic framework without Euclidean assumptions
Abstract
Bregman proximal stochastic gradient (BPSG) methods bring variance-reduced composite optimization to objectives whose geometry is poorly captured by Euclidean smoothness. Their performance, however, remains sensitive to the step size: raw stochastic curvature estimates can fluctuate sharply, whereas line searches add repeated proximal evaluations. We introduce Ada-BPSG, a line-search-free BPSG method that couples the SAGA gradient table with a stabilized Barzilai--Borwein (BB) candidate. A mediant aggregates incremental secant information so that nearly singular local ratios receive little weight, and an explicit safeguard translates the resulting curvature estimate into the bounded step-size sequence required for convergence. This design yields a direct analytical chain from relative smoothness and component-wise variance control to convergence in finite-dimensional normed spaces. We prove an O(n/K) ergodic rate for convex objectives, a restarted linear rate under relative quadratic growth, and an O(1/K) bound for a Bregman proximal residual in the nonconvex setting. On logistic regression and sparse nonnegative matrix factorization, Ada-BPSG combines low objective values with substantially less sensitivity to the initial step size than standard variance-reduced baselines, while avoiding line search.
Sources
- Stabilized Barzilai-Borwein method
- SVRG Meets AdaGrad: Painless Variance Reduction
- Adam: A Method for Stochastic Optimization
- AI-SARAH: Adaptive and Implicit Stochastic Recursive Gradient Methods
- A Stochastic Bregman Primal-Dual Splitting Algorithm for Composite Optimization
- An Improved Optimal Proximal Gradient Algorithm for Non-Blind Image Deblurring
- AdaBB: Adaptive Barzilai-Borwein Method for Convex Optimization
- Hyperspectral Unmixing: Ground Truth Labeling, Datasets, Benchmark Performances and Survey
Related papers
- Lions and Muons: Optimization via Stochastic Frank-Wolfe under Heavy-Tailed Noise
- Adam-HNAG: A Convergent Reformulation of Adam with Accelerated Rate
- Incremental Learning in Mirror Flows
- Online Control via Counterfactual Tracking
- Asynchronous Replanning in Two Population Linear Quadratic Mean Field Games: Information Requirements and Stability
- Petrov-Galerkin operator inference with application to stability-encouraging identification