# Hyperparameter tuning via trajectory predictions: Stochastic prox-linear methods in matrix sensing

arXiv: https://arxiv.org/abs/2402.01599 · HTML: https://mlsanity.com/paper/2402.01599/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 4 of its 14 references that are in the map.
- First posted on arXiv 2 Feb 2024.

## Builds on (most cited first, 4 of 4)

- [High-dimensional limit theorems for SGD: Effective dynamics and critical scaling](https://mlsanity.com/paper/2206.04030/) (2022)
- [Efficiency of minimizing compositions of convex functions and smooth maps](https://mlsanity.com/paper/1605.00125/) (2016)
- [The importance of better models in stochastic optimization](https://mlsanity.com/paper/1903.08619/) (2019)
- [Stochastic algorithms with geometric step decay converge linearly on sharp functions](https://mlsanity.com/paper/1907.09547/) (2019)

## Built on by (most cited first, 4 of 4)

- [Accurate, provable and fast polychromatic tomographic reconstruction: A variational inequality approach](https://mlsanity.com/paper/2503.19925/) (2025)
- [Precise asymptotics of reweighted least-squares algorithms for linear diagonal networks](https://mlsanity.com/paper/2406.02769/) (2024)
- [Computationally efficient reductions between some statistical models](https://mlsanity.com/paper/2402.07717/) (2024)
- [Estimating Generalization Performance Along the Trajectory of Proximal SGD in Robust Regression](https://mlsanity.com/paper/2410.02629/) (2024)

## Abstract (arXiv)

Motivated by the desire to understand stochastic algorithms for nonconvex optimization that are robust to their hyperparameter choices, we analyze a mini-batched prox-linear iterative algorithm for the problem of recovering an unknown rank-1 matrix from rank-1 Gaussian measurements corrupted by noise. We derive a deterministic recursion that predicts the error of this method and show, using a non-asymptotic framework, that this prediction is accurate for any batch-size and a large range of step-sizes. In particular, our analysis reveals that this method, though stochastic, converges linearly from a local initialization with a fixed step-size to a statistical error floor. Our analysis also exposes how the batch-size, step-size, and noise level affect the (linear) convergence rate and the eventual statistical estimation error, and we demonstrate how to use our deterministic predictions to perform hyperparameter tuning (e.g. step-size and batch-size selection) without ever running the method. On a technical level, our analysis is enabled in part by showing that the fluctuations of the empirical iterates around our deterministic predictions scale with the error of the previous iterate.