A Markovian Model for Learning-to-Optimize

arXiv:2408.11629 · cs.LG, math.PR · Submitted 2026-08-14 · Read on arXiv

Listen

Radio episode about this paper

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Next we'll be talking about the paper "A Markovian Model for Learning-to-Optimize".

Jane: The paper was written by Michael Sucker and Peter Ochs from University of Tübingen and Saarland University.

Tom: Stay tuned as we take you through the paper and discuss its implications.

Title: Tom: Welcome back to the show, everybody. Today we're diving into a fresh arXiv paper that's got a mouthful of a title: "A Markovian Model for Learning-to-Optimize." Jane, what's your first read on that?

Jane: Tom, I love this one. The title tells you exactly what they're doing—they're taking the idea of learning to optimize, which is this hot field where you train an algorithm to solve other optimization problems, and they're putting it on solid mathematical ground using Markov chains.

Tom: And for our listeners who might not be deep in the math weeds, a Markov chain is basically a process where the next step only depends on where you are right now, not the whole history. That's exactly how these iterative optimization algorithms work—you're at a point, you take a step, you're at a new point.

Jane: Right. And the authors, Michael Sucker and Peter Ochs from Tübingen and Saarland, they're saying, look, when we train an optimization algorithm, we're really training a stochastic process. The trajectory it produces—the whole sequence of points—that's the thing we care about, not just the final answer.

Tom: So they're not just saying "here's a cool new optimizer." They're building a whole probabilistic framework for understanding what happens when you train one. And the key word there is "stochastic"—there's randomness everywhere.

Jane: Exactly. And that randomness comes from four different sources. You've got the random choice of problem, the random starting point, the random hyperparameters, and the internal randomness of the algorithm itself. The paper calls this a "superposition" of randomness, which I think is a great way to put it.

Tom: That's a lot of moving parts. So the big claim here is that they can actually model the distribution of the entire trajectory—the whole infinite sequence of points—in a way that's unique and well-defined. That's the Markovian part.

Jane: And once you have that distribution, you can start asking real questions. Like, how long does this algorithm actually take to converge? What's its actual convergence rate? Not just on training data, but on new problems it's never seen.

Tom: That's the generalization question, and it's huge. Because right now, a lot of learned optimizers are just black boxes that work great in practice but have no guarantees. This paper is trying to change that.

Jane: And they're doing it with PAC-Bayesian bounds, which is a fancy way of saying they can give you high-probability guarantees. We'll get into the weeds on that in a bit, but for now, the takeaway is that they're not just hoping the algorithm works—they're proving it will, with high confidence.

Tom: I love it when theory and practice actually talk to each other. So we've got the framework, we've got the guarantees—what do they actually do with it? That's coming up next.

Jane: Stay tuned, because they run real experiments, and the results are pretty striking.

Summary: Tom: So we've established that "A Markovian Model for Learning-to-Optimize" builds this rigorous probabilistic foundation. Jane, walk me through what they actually prove.

Jane: Okay, so the centerpiece is Theorem thirty-three in the paper. It's a PAC-Bayesian generalization bound. In plain English, it says: if you train your optimizer on a bunch of problems, you can bound its performance on new problems from the same distribution. And the bound depends on how much data you have—it shrinks as you get more training problems.

Tom: And what do they apply this bound to? Because I remember you said convergence time and rate.

Jane: Exactly. They define two key quantities. The first is convergence time—how many iterations until the algorithm hits a stopping criterion, like the loss dropping below a threshold. The second is convergence rate—how much the loss contracts each iteration. Both of these are functions of the whole trajectory, which is why their Markovian model is so important.

Tom: So they're not just bounding the final loss. They're bounding the actual behavior of the algorithm over time. That's a much stronger statement.

Jane: Much stronger. And they do it with a clever trick. They define a "rate function" that looks at the loss at the stopping time versus the initial loss, raised to the power of one over the number of iterations. It's like asking, on average, by what factor does the algorithm shrink the loss each step?

Tom: And they also handle the case where the algorithm just doesn't converge within a budget, right? They cap it at some maximum number of iterations.

Jane: Right, they use t max as a cutoff. So the convergence time is always finite, which makes the math work. And they have a separate result for the probability of observing a trajectory with a certain property—like converging with a rate better than some threshold.

Tom: So they've got bounds for the expected time, the expected rate, and the probability of success. That's a pretty complete toolkit.

Jane: And here's the kicker—they don't just prove these bounds exist. They actually use them to train algorithms. The training procedure minimizes the PAC-bound itself, which means you're directly optimizing for a guarantee, not just for empirical performance.

Tom: That's a big philosophical shift. You're not just trying to make the loss small on training data; you're trying to make the bound small, which gives you a certificate for new data.

Jane: Precisely. And they validate this with five experiments. Quadratics, image processing, LASSO, training a neural network, and stochastic risk minimization. In most cases, the learned algorithm crushes the classical baseline.

Tom: I saw that in the quadratic experiment, the learned algorithm outperforms heavy-ball by orders of magnitude. That's not a small win.

Jane: Not at all. And in the image processing case, it beats Nesterov's accelerated gradient. But—and this is important—the paper is honest about the failures too. In some cases, the PAC-bound is vacuous, meaning it's too loose to be useful. And on a few problem instances, the learned algorithm just doesn't converge.

Tom: So it's not a magic bullet, but it's a real step forward. The question is, how do they actually build these algorithms? What's the architecture? That's where it gets really interesting.

Jane: That's our next segment. The engineering side of this paper is genuinely clever.

Improvements: Tom: So Jane, we've talked about the theory and the results. Now I want to get into the weeds of how they actually implement this. Because "learning-to-optimize" can mean a lot of things.

Jane: Right. And this is where the paper gets really practical. The key improvement over previous work is that they don't just learn a step size or a preconditioner. They learn the entire update rule, parameterized by a neural network, and they train it by minimizing the PAC-Bayesian bound directly.

Tom: So the algorithm itself is a neural network that takes in the current state, the gradient, maybe some history, and outputs the next point. That's the "learned optimizer" part.

Jane: Exactly. And they have different architectures for different problem classes. For the quadratic case, it's a coordinate-wise operation—the network processes each dimension independently, which makes sense because quadratics are separable.

Tom: And for the image processing problem, they use convolutional layers. That's clever because images have spatial structure, and a one times one convolution acts on each pixel independently but shares weights across the image.

Jane: Right. And for the LASSO problem, they do something really smart. They split the state into zero and non-zero entries, because the support of the solution—which coordinates are active—is crucial. The network learns to treat those differently.

Tom: So the architecture is problem-specific. But the training procedure is the same across all of them. Walk me through that.

Jane: So they start with a "warm start"—they train the network with a simple loss function that measures how much the loss decreases each iteration. That gives them a decent starting point. Then they sample a bunch of hyperparameter candidates around that starting point to build a discrete prior distribution.

Tom: And then they do the PAC-Bayesian optimization step, which is in closed form. That's the beauty of using a discrete prior—the optimization over the posterior distribution becomes tractable.

Jane: Exactly. They compute the posterior that minimizes the bound, and then they just pick the hyperparameter with the highest posterior probability. It's simple, it's fast, and it comes with a guarantee.

Tom: Now, I noticed something interesting in the experiments. The mean and median performance often diverge significantly. What's going on there?

Jane: That's the robustness issue. On most problems, the learned algorithm is fantastic. But there are a few problem instances where it just fails—it doesn't converge, or it converges very slowly. Those outliers drag the mean up, even though the median is great.

Tom: So the PAC-bound is giving you an average guarantee, but it's not protecting you against the worst case.

Jane: Right. And the paper acknowledges this. In the image processing experiment, about four percent of test problems don't converge within the budget. The bound still holds—it's an average over the distribution—but it's a reminder that these methods aren't universally robust.

Tom: But here's what I find exciting: the bound itself is a training signal. By minimizing it, you're explicitly trading off empirical performance against the complexity of the posterior distribution. That's a principled way to avoid overfitting.

Jane: And that's the real improvement. Previous work in learning-to-optimize often just trained on empirical loss and hoped for generalization. This paper makes generalization part of the objective. That's a fundamental shift.

Tom: So what does this mean for the field? Where does this leave us? I think that's the question for our next segment.

Jane: And I think Lu and Meng will have strong opinions on that.

Conclusion: Tom: Alright, we've covered the theory, the experiments, and the architecture of "A Markovian Model for Learning-to-Optimize." Let's bring in Lu and Meng to get their takes.

Lu: Thanks, Tom. I think the most exciting implication here is that this framework gives us a language to talk about learned optimizers with mathematical precision. We're not just saying "it works"—we're saying "here's the distribution of trajectories, here's the bound on convergence time, and here's the confidence level."

Meng: And from an engineering standpoint, that's huge. When I'm deploying a learned optimizer in production, I need to know it's not going to blow up on some edge case. This paper gives me a way to certify that, at least in expectation.

Jane: But Meng, you saw the experiments—the bound is sometimes vacuous. In the stochastic risk minimization case, the bound for the convergence rate was useless.

Meng: True, but that's a known limitation of PAC-Bayes bounds. They're conservative by design. And in the convergence time case, the bound was actually quite tight. So it's not all doom and gloom.

Lu: I'd push back on that slightly. The vacuous bounds in some experiments tell me that the theory is still ahead of the practice. The framework is sound, but we need better ways to tighten these bounds, maybe by using more data or better priors.

Tom: So what's the path forward? What would you two want to see next?

Lu: I'd love to see this framework extended to non-asymptotic results. Right now, the convergence time is capped at t max, which is fine for practical purposes, but the paper itself acknowledges that truly asymptotic guarantees are out of reach with this approach.

Meng: And I'd want to see more work on robustness. The fact that four percent of image processing problems fail is concerning. Can we detect those cases ahead of time? Can we fall back to a classical algorithm when the learned one is uncertain?

Jane: Those are both great directions. And I think the broader impact is cultural. This paper shows that we can have both—performance and guarantees. We don't have to choose between a fast learned optimizer and a provably convergent one.

Tom: And that's the message I want listeners to take away. "A Markovian Model for Learning-to-Optimize" isn't just a paper about optimization. It's a blueprint for how to do rigorous machine learning—where every claim is backed by a bound, and every bound is backed by a proof.

Lu: And it opens the door for other fields to adopt this approach. Anywhere you have an iterative stochastic process—reinforcement learning, MCMC sampling, even financial modeling—this framework could apply.

Meng: I'd add that the code and architectures are described in enough detail that a competent engineer could reproduce them. That's not always the case in theory-heavy papers.

Jane: So we're saying this paper is both theoretically deep and practically useful. That's a rare combination.

Tom: It is. And with that, we're going to wrap up our discussion of "A Markovian Model for Learning-to-Optimize." Thanks to Lu and Meng for joining us, and thanks to all our listeners for tuning in.

Jane: Next up, we've got a paper on diffusion models for protein design. That should be fun.

Tom: Can't wait. Until then, keep optimizing, everyone.

Michael Sucker, Peter Ochs

University of Tübingen · Saarland University

cs.LG, math.PR

Submitted: 2026-08-14

Updated: 2026-08-17

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 72/100

Key concepts

Markov Chain
A process where the next state only depends on the current state, not the entire history. Optimization algorithms are viewed this way because taking a step only requires knowing your current point.
Learning-to-Optimize
The field of training an algorithm (often a neural network) to solve other optimization problems. Instead of designing a fixed optimizer, the system learns the update rule itself.
PAC-Bayesian Bounds
A mathematical tool used to provide high-probability guarantees about an algorithm's performance on new, unseen data. It allows researchers to prove that the learned optimizer will perform well with high confidence.
Stochastic Process
A process involving randomness, which is inherent in optimization. The paper models the entire sequence of points (trajectory) generated by the algorithm as a stochastic process.

Terminology

Summary

Summary

This paper presents a probabilistic model for stochastic iterative algorithms, with a specific focus on optimization algorithms, and uses this model to derive PAC-Bayesian generalization bounds for functions defined on the trajectory of a learned algorithm. The authors state their primary contribution as follows: "In this work, we consider parametric stochastic iterative (optimization) algorithms to minimize parametric (loss) functions, and how to learn such algorithms with theoretical guarantees on their non-asymptotic convergence rate and convergence time."

The starting point is a parametric loss function l(s, θ) to be minimized in s for every realization of θ. A stochastic algorithm A is applied iteratively, yielding a sequence ξ = (ξ(t)) t∈N0 defined by the update rule: ξ(t+1) = A(α, θ, ξ(t), η(t+1)). Here, α represents hyperparameters, θ specifies the loss function, and η(t+1) models internal randomness. The authors emphasize that both convergence rate and stopping time are properties of the entire trajectory, not single iterates, so they model the distribution of the whole stochastic process.

The paper is structured as follows: after related work and preliminaries, Section 4 derives the probabilistic model by building it from ground up, basically starting from Equation (1). Section 5 uses this model to derive generalization bounds for learning such an algorithm based on data, including bounds for its non-asymptotic convergence rate and convergence time. Section 6 presents five practically relevant experiments.

Probabilistic Model: The model relies on two mild assumptions: Assumption 9 provides four Polish probability spaces: the state space (S, B(S), P I), the parameter space (P, B(P), P P), the hyperparameter space (H, B(H), P H), and the randomization space (R, B(R), P R). Assumption 11 provides a measurable loss function l: S × P → [0, ∞] and a measurable algorithmic update A: H × P × S × R → S.

The core of the model is the transition kernel defined in Definition 13: γ: H × P × S → S, given by γ((α, θ, x), B) = P R A(α, θ, x, ·) ∈ B. The joint transition kernel Γ for N problem instances is the product of individual kernels. The transition semi-group (γ t) t∈N0 is defined recursively, and Theorem 19 establishes that there exists a unique probability kernel ψ from H × P to S N0 such that the finite-dimensional distributions of the trajectory match the iterates generated by A. Similarly, a unique kernel Ψ from H × P N to S N N0 exists for the joint distribution of N trajectories. Lemma 21 shows that the joint distribution factorizes: Ψ(α, θ[N]) is (up to reordering) the product of the individual distributions ψ(α, θ n), meaning the processes are conditionally independent given the hyperparameters and parameters.

The probability space (Ω, A, P) is then defined with Ω:= H × P N × S N N0, and P:= P H ⊗ P P⊗N ⊗ Ψ. Lemma 24 establishes regular versions of conditional distributions, and Corollary 26 confirms conditional independence of the processes.

Stopping Times: The convergence set C ⊂ P × S is defined in Definition 28 as the set of points satisfying the convergence criterion for l(·, θ). Assumption 29 requires C to be measurable. The random time τ n:= τ max ∧ τ conv,n is introduced, where τ max is a maximal computational budget and τ conv,n is the first time the process enters the convergence set. Proposition 31 proves that τ n is a stopping time with respect to the natural filtration.

Generalization Results: The main theoretical result is Theorem 33, which provides a PAC-Bayesian generalization bound for any measurable function f: P × S N0 → [0, ∞) bounded by f max. The theorem states that for every λ ∈ (0, ∞) and ε > 0, with probability at least 1 - ε over the data set P[N], for all posterior distributions ρ ∈ P(P H):

ρ[E(P,ξ)H f] ≤ (1/N) Σ n=1 N ρ[E(P n,ξ n)H,P n f] + (D KL(ρ ∥ P H) + (λ/(2N)) f max2 - log(ε)) / λ.

This bound is then specialized in two corollaries:

  • Corollary 35 (Convergence Time): For the bounded function T (the stopping time, bounded by t max), the bound provides a guarantee on the average expected convergence time τ̄.

  • Corollary 38 (Convergence Rate): For the bounded rate function r b (defined via a contraction function c(θ, x, y) = l(x, θ)/l(y, θ) · 1 l(y, θ) > 0 and rate function r(θ, (z(t))) = (c(θ, z(T), z(0)))(1/T) · 1 T ≥ 1), the bound provides a guarantee on the expected bounded convergence rate r̄ b.

Additionally, Theorem 42 (attributed to Catoni) provides a tighter bound for the probability of observing a trajectory with a certain property (encoded in a measurable set A), using the function Φ a(p) = -(1/a) log(1 - [1 - exp(-a)]p). Remark 43 notes that by a union bound, all three results (convergence time, convergence rate, and property probability) can be applied simultaneously with confidence level 1 - ε.

Numerical Results: The paper presents five experiments:

  1. Quadratics: Strongly convex and smooth quadratic problems (d=200). The learned algorithm outperforms heavy-ball with friction (HBF) by orders of magnitude, reaching the convergence criterion faster. The PAC-bound for convergence time is reasonably tight, while the bound for convergence rate is not vacuous, but also not really tight.

  2. Image Processing: Convex and smooth image denoising/deblurring (d=30000). The learned algorithm outperforms Nesterov accelerated gradient (NAG) on typical problems, but fails to converge on about 4% of test problems. The PAC-bound for convergence rate is vacuous, but the bound for convergence time is quite tight.

  3. LASSO: Convex and non-smooth problem (d=70, p=35). The learned algorithm outperforms FISTA by several orders of magnitude, reaching the stopping criterion in about 200 iterations. Both PAC-bounds are reasonable tight.

  4. Training a Neural Network: Non-convex and non-smooth problem of training a two-layer ReLU network (p=351 parameters). The learned algorithm clearly outperforms Adam, reaching the ground-truth loss after about 25 iterations. The PAC-bound for convergence time is reasonably tight, while the rate bound varies strongly due to plateaus.

  5. Stochastic Empirical Risk Minimization: Same neural network training but with minibatches (size m=5). The learned algorithm (a preconditioned version of Adam) outperforms Adam, with a median convergence time of less than 500 iterations versus about 2000 for Adam. The PAC-bound for convergence time is reasonably tight, while the rate bound is vacuous.

Conclusion: The authors conclude that their framework allows for giving generalization guarantees for (nearly) any kind of statistics that one wants to have for such an algorithm, particularly convergence rates and stopping times. They note that both results are non-asymptotic, and that asymptotic events are inherently non-observable and thus may not be within reach for practical generalization results.

Improvements for AI systems

Based on the paper, here are specific improvements to AI systems and what the improved system can do:

1. Guaranteed-Convergence Learned Optimizers

  • Improvement: Replace black-box learned optimizers (e.g., learned step-size or update rules) with a system that, during training, explicitly minimizes a PAC-Bayesian upper bound on the expected convergence time and rate (Corollaries 35 and 38).

  • What it can do: For a new problem instance (e.g., a new image or dataset), the system outputs a hyperparameter distribution (or a single best hyperparameter) along with a provable high-probability upper bound on: (a) the number of iterations needed to reach a user-defined stopping criterion (e.g., gradient norm < ε), and (b) the per-iteration contraction factor of the loss. This means you can deploy the optimizer in safety-critical applications (medical imaging, autonomous systems) with a certificate that it will finish within a known time budget, rather than hoping it converges.

2. Early-Stopping with Statistical Certificates

  • Improvement: Use the stopping-time framework (Section 4.3, Proposition 31) to train a system that learns when to stop an iterative algorithm (e.g., during neural network training) to avoid overfitting.

  • What it can do: The system monitors the trajectory (loss, gradient norm, validation loss) and, at each iteration, computes a PAC-Bayesian bound on the probability that the current iterate is good enough (Theorem 42). It stops automatically when this probability exceeds a threshold (e.g., 95%). This gives a principled, data-driven alternative to fixed iteration counts or manual early-stopping, with a guarantee that you are not stopping too early or wasting compute.

3. Robust Hyperparameter Selection with Risk Certificates

  • Improvement: Replace standard hyperparameter tuning (grid search, Bayesian optimization) with the proposed PAC-Bayesian framework, which explicitly accounts for the distribution over hyperparameters (posterior ρ) and its divergence from a prior.

  • What it can do: Given a small dataset of problem instances, the system outputs a posterior distribution over hyperparameters (e.g., step-size, momentum) that minimizes the empirical convergence time/rate, while simultaneously providing a high-probability bound on the true average performance on unseen instances. This is particularly useful when you have limited data (e.g., few medical images) and need to know how well your chosen hyperparameters will generalize, not just how they performed on the training set.

4. Trajectory-Aware Loss Functions for Training Optimizers

  • Improvement: Modify the training loss of the learned optimizer to be a function of the entire trajectory (e.g., the rate function in Definition 36), not just the final iterate or a single step. This is enabled by the Markovian model of the trajectory distribution (Theorem 19).

  • What it can do: The learned optimizer is trained to directly minimize a meaningful, trajectory-level objective (e.g., geometric mean contraction per iteration, or time-to-converge), rather than a proxy like sum of losses. This leads to optimizers that are genuinely faster on average, not just better at reducing loss in the first few steps. The system can also be trained to avoid pathological behaviors like oscillating or plateauing, because these are explicitly penalized in the trajectory-level loss.

5. Conditional Independence for Scalable Multi-Problem Learning

  • Improvement: Leverage the factorization result (Lemma 21) to train the optimizer on a large number of problem instances (N) in a way that is computationally efficient and statistically sound.

  • What it can do: The system can process N different problem instances (e.g., N different blurring kernels in image deblurring) in parallel, knowing that the trajectories are conditionally independent given the hyperparameters. This allows for: (a) efficient batch training using standard parallel hardware, (b) exact computation of the empirical risk (average over N) without approximation, and (c) tighter generalization bounds because the effective sample size is N, not 1.

6. Probabilistic Stopping-Time Prediction for Resource Allocation

  • Improvement: Use the distribution of the stopping time τ (Section 4.3) to build a predictive model that, given a problem instance, outputs a full probability distribution over the number of iterations needed, not just a point estimate.

  • What it can do: In a cloud computing or job scheduling environment, the system can predict, with a confidence interval, how long a given optimization task will take. This enables: (a) better resource allocation (e.g., reserving GPU time), (b) dynamic load balancing (e.g., if one task is predicted to take longer, start another), and (c) setting realistic timeouts for users. The PAC-Bayesian bound ensures that the predicted distribution is not overconfident.

7. Safe Deployment of Learned Optimizers in Non-Convex Problems

  • Improvement: Apply the framework to non-convex problems (as in the neural network training experiments, Section 6.4) with a convergence set defined by a combination of gradient norm and loss value.

  • What it can do: The system can be deployed to train neural networks with a certificate that, with high probability, it will reach a region where the gradient norm is small (e.g., < 0.75) and the loss is below a threshold (e.g., < 0.75) within a specified number of iterations. This is valuable for training models in production where you need to guarantee that training will finish in a reasonable time and not get stuck in a bad local minimum. The certificate is based on the observed behavior on a training set of similar problems, not on theoretical assumptions about the loss landscape.

Abstract

We present a probabilistic model for stochastic iterative algorithms with the use case of optimization algorithms in mind. Based on this model, we present PAC-Bayesian generalization bounds for functions that are defined on the trajectory of the learned algorithm, for example, the expected (non-asymptotic) convergence rate and the expected time to reach the stopping criterion. Thus, not only does this model allow for learning stochastic algorithms based on their empirical performance, it also yields results about their actual convergence rate and their actual convergence time. We stress that, since the model is valid in a more general setting than learning-to-optimize, it is of interest for other fields of application, too. Finally, we conduct five practically relevant experiments, showing the validity of our claims.

Sources

Related papers