Machine Learning for Inverse Problems and Data Assimilation

arXiv:2410.10523 · stat.ML, cs.LG, math.OC · Submitted 2026-08-09 · 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 "Machine Learning for Inverse Problems and Data Assimilation".

Jane: The paper was written by Eviatar Bach, Ricardo Baptista, Daniel Sanz-Alonso and Andrew Stuart from University of Reading and University of Toronto and University of Chicago and California Institute of Technology.

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

Jane: We also have Lu with us today — senior AI researcher at Tsinghua.

Tom: We also have Meng with us today — lead engineer at a mysterious AI startup.

Jane: We also have Lalam with us today — the in-house Large Language Model.

Tom: Alright, let's get started.

Title: Tom: Welcome back to the channel, everyone. We’ve got a big one today — the paper is called "Machine Learning for Inverse Problems and Data Assimilation." Jane, this looks like a textbook, not a typical research paper.

Jane: It really does, Tom. And that’s actually part of why it’s so exciting. It’s a full mathematical treatment of how machine learning is reshaping two fields that are usually pretty separate — inverse problems and data assimilation.

Tom: Right, so for our listeners who aren’t steeped in this — what’s the difference? I always get these two confused.

Jane: So an inverse problem is when you have some measurements and you want to figure out what caused them. Think medical imaging — you see the scan, you want the tissue. Data assimilation is more about blending a model of a system with observations over time, like weather forecasting.

Tom: And the authors are Eviatar Bach, Ricardo Baptista, Daniel Sanz-Alonso, and Andrew Stuart. That last name — Stuart — is a giant in this area. He’s been pushing the Bayesian approach to inverse problems for years.

Jane: Yeah, and what they’ve done here is write what looks like a definitive reference. It’s got three parts — inverse problems, data assimilation, and then a fundamentals section covering metrics, supervised learning, generative models, optimization.

Tom: So it’s not just a survey. It’s a whole course.

Jane: Exactly. And the really interesting thing is the way they structure it. They’ve got these parallel chapters — one on inverse problems, one on data assimilation — and they cover the same ideas in both. Variational inference, learning priors, transport methods, amortization.

Tom: Amortization — that’s the idea that you train once on lots of data and then reuse the model for new observations, right?

Jane: Right. So instead of solving a brand new inverse problem every time you get new data, you learn a mapping from data to solution once, and then it’s cheap to apply.

Tom: That’s huge for real-world use. I mean, think about weather forecasting — you don’t want to rerun a massive optimization every hour.

Jane: And that’s exactly the kind of thing this book is setting up. It’s not just theory for theory’s sake — it’s building the mathematical foundation for methods people actually want to deploy.

Tom: Lu, you’ve been quiet — what’s your take on this as someone who works in the field?

Lu: I think the most valuable thing here is that they’re not just throwing neural networks at everything. They’re very careful about what machine learning can and cannot do in these settings. The book is really about principled use of ML.

Tom: Principled — I like that. So this isn’t a hype machine.

Lu: Not at all. It’s the opposite. It’s saying, here’s the math, here’s what’s guaranteed, here’s what’s not. That’s rare and valuable.

Jane: And it’s going to be a reference for years. I can see this being the textbook for graduate courses in computational science.

Tom: Alright, so we’ve got the big picture. Next up, we’re going to dig into what the book actually covers in its summary — the structure, the methods, the key ideas. Stay with us.

Summary: Tom: So we’re back with "Machine Learning for Inverse Problems and Data Assimilation." Jane, you mentioned the structure — let’s get into what’s actually in here.

Jane: The book is organized around three parts. Part one is inverse problems — that’s the classic setup where you have a forward model, you have data, and you want the unknown input. Part two is data assimilation, which is the time-dependent version. And part three is the mathematical toolbox.

Tom: And the toolbox — that’s the stuff you need to understand everything else, right?

Jane: Exactly. Metrics, divergences, scoring rules, supervised learning, generative modeling, time-series forecasting, optimization, sampling. It’s all there.

Lu: What I find really interesting is how they use the same conceptual framework in both parts. So for inverse problems, you have chapters on variational inference, learning the prior, transport to the posterior, and data dependence. And then for data assimilation, you have the exact same structure — variational inference, learning the prior, transport perspective on filters, data dependence of filters.

Tom: That’s a really clean way to teach it. You learn the concept once and then you see it applied in two different settings.

Jane: And it makes the connection between the fields explicit. People in inverse problems and people in data assimilation don’t always talk to each other, but they’re solving fundamentally the same kind of problem.

Tom: What do you mean by that?

Jane: Well, in both cases you have some unknown — a parameter or a state — and you have noisy observations. And you want to combine prior knowledge with data to get a posterior distribution. The math is the same; the context is different.

Meng: I’m the engineer here, so I have to ask — how practical is this? Like, can I actually implement these methods?

Jane: That’s the thing — the book is mathematical, but it’s not abstract. They give you algorithms. They tell you what the objective function is, how to optimize it, what the computational trade-offs are.

Meng: So it’s not just "here’s a theorem, good luck."

Jane: No, it’s much more hands-on than that. They discuss ensemble Kalman methods, particle filters, normalizing flows, score-based generative models — all the tools people actually use.

Lu: And they’re honest about limitations. They talk about weight collapse in particle filters, they talk about the challenges of high-dimensional problems, they talk about when variational inference fails.

Tom: So it’s a realistic picture.

Lu: Very realistic. It’s not promising magic. It’s saying, here are the tools, here’s what they’re good at, here’s where they break.

Tom: I love that. Okay, so we’ve got the structure. Next, let’s talk about what this book is actually improving — what are the new ideas, the advances over what came before?

Improvements: Tom: Alright, so we’re still with "Machine Learning for Inverse Problems and Data Assimilation." Jane, what do you see as the real advances here — what’s new compared to older textbooks?

Jane: I think the biggest thing is the systematic treatment of machine learning as a first-class tool, not an add-on. Older books on inverse problems barely mention neural networks. This one puts them front and center.

Tom: And not just neural networks — they cover random features, Gaussian processes, normalizing flows, score-based models. It’s the whole modern toolkit.

Jane: Right. And they connect it all back to the core mathematics. So when they talk about learning a prior, they don’t just say "use a GAN" — they show you how it fits into the Bayesian framework.

Lu: I’d add that the transport perspective is a real contribution. They use transport maps as a unifying concept — mapping a simple distribution to a complex one, mapping prior to posterior, mapping forecast to analysis in filtering.

Tom: Transport — that’s the idea of pushing one probability distribution into another through a function, right?

Lu: Exactly. And it turns out to be incredibly useful. It connects optimal transport theory, normalizing flows, and even the ensemble Kalman filter — they’re all instances of the same idea.

Meng: So instead of learning a density function directly, you learn a map that transforms samples?

Lu: Yes. And that’s often much easier, because you don’t need to know the normalization constant. You just need to be able to evaluate the map.

Jane: And that’s a huge computational advantage. Normalization constants are often intractable in high dimensions.

Tom: What about the data assimilation side? What’s the improvement there?

Jane: They bring the same tools to filtering and smoothing. So instead of just using a Kalman filter or a particle filter, you can learn the analysis step — the step that incorporates new observations — using machine learning.

Meng: So you’re learning the filter itself, not just the model inside the filter?

Jane: Exactly. And they show how to do that with variational inference, with scoring rules, with transport maps. It’s a whole new way of designing filters.

Lu: And they also address the problem of model error — when your dynamics model isn’t perfect. They show how to learn corrections from data, which is a huge practical issue.

Tom: Model error — that’s when your equations don’t perfectly describe reality, right?

Lu: Right. And it’s everywhere. Every weather model has it. Every climate model has it. So being able to learn corrections is really valuable.

Meng: So this book is essentially a blueprint for building the next generation of forecasting systems?

Jane: I think that’s fair. It’s giving people the mathematical foundation to do that in a principled way.

Tom: Okay, so we’ve got the big ideas. Now let’s zoom in on the actual first page of the paper — what do the authors say they’re trying to do?

First Page: Tom: So we’re back with "Machine Learning for Inverse Problems and Data Assimilation." Let’s look at the very first page — the preface. Jane, what are the authors saying right out of the gate?

Jane: They’re very clear about their audience. They say this is aimed primarily at researchers from inverse problems and data assimilation who want to see a mathematical presentation of machine learning as it pertains to their fields.

Tom: So it’s not for machine learning people who want to learn inverse problems. It’s the other way around.

Jane: Exactly. And they’re also clear about the structure — the fundamentals are in Part III, and Parts I and II use that material. But they say you can read it either way — application-first or theory-first.

Tom: That’s nice. So you can dip in depending on your background.

Jane: Right. And they mention that they’ve taught from this book both ways, which is a good sign — it means the structure actually works.

Lu: I think the most important thing on that first page is the list of overarching concepts. They say there are four ideas that organize both the inverse problems part and the data assimilation part.

Tom: And those are?

Lu: Variational Bayes, learning priors from data, transport as a unifying concept, and amortization — learning dependence on observations so you can reuse the algorithm.

Jane: And that parallel structure is really elegant. You see the same four ideas in both contexts, which reinforces the connection between the fields.

Tom: They also mention the prerequisites. What do they expect you to know?

Jane: Linear algebra, probability, statistics, multivariable calculus. And they point to a companion book — "Inverse Problems and Data Assimilation" by Sanz-Alonso, Stuart, and Taeb — for background.

Tom: So this is really a follow-up to that earlier book?

Jane: It seems like it. They say they’ve tried to maintain similar notation, which is helpful if you’ve read the first one.

Meng: I’m curious about the notation section. They have a lot of conventions there.

Jane: They do — and it’s actually well done. They define everything from sets to vector spaces to probability notation. It’s the kind of thing you’ll want to bookmark.

Tom: And they acknowledge funding and students who took the course at Caltech. So this really did come out of teaching.

Lu: That’s often the best kind of book — it’s been tested on real students.

Jane: Absolutely. And the fact that they mention the course was at Caltech in Spring two thousand twenty-four tells you this is current — this is the state of the art as of now.

Tom: Okay, so we’ve got the setup. Let’s wrap this up with our final thoughts.

Conclusion: Tom: So we’ve spent this whole episode on "Machine Learning for Inverse Problems and Data Assimilation." Jane, give us the final summary.

Jane: This is a comprehensive mathematical treatment of how machine learning is transforming two related fields — inverse problems and data assimilation. It covers everything from the fundamentals — metrics, divergences, scoring rules — to the latest methods — normalizing flows, score-based models, amortized inference.

Tom: And the key idea is that the same concepts apply in both fields. Variational inference, learning priors, transport, amortization — they show up in both inverse problems and data assimilation.

Jane: Exactly. And that parallel structure is what makes the book so valuable. It’s not just a collection of methods. It’s a unified framework.

Lu: I’d say the biggest impact will be on graduate education. This is going to be the textbook for a generation of students in computational science.

Meng: And for practitioners, it’s a reference. If you’re building a forecasting system or solving an inverse problem, you can look up the right method and understand the trade-offs.

Tom: What about the broader implications? Why does this matter for the world?

Jane: Think about weather forecasting, climate modeling, medical imaging, geophysics — all of these rely on inverse problems and data assimilation. Making these methods better has direct impact on people’s lives.

Lu: And the book is pushing toward a future where these systems are more accurate, more efficient, and more robust — because they’re built on solid mathematical foundations rather than ad hoc heuristics.

Tom: I love that. So who should pick this up?

Jane: Anyone working in computational science, really. Graduate students, researchers, engineers who want to understand the math behind the tools they’re using.

Meng: And it’s accessible enough that you don’t need to be a mathematician to get value from it — but you do need to be willing to work through the math.

Tom: Alright, well, that’s our take on "Machine Learning for Inverse Problems and Data Assimilation." A big thank you to the authors — Bach, Baptista, Sanz-Alonso, and Stuart — for putting this together.

Jane: It’s going to be a standard reference for years to come. We’re excited to see how it shapes the field.

Tom: And that’s a wrap for this episode. Next time, we’ll be looking at another paper from the arXiv. Until then, keep learning, keep questioning, and we’ll see you on the next one.

Jane: Take care, everyone.

Eviatar Bach, Ricardo Baptista, Daniel Sanz-Alonso, Andrew Stuart

University of Reading · University of Toronto · University of Chicago · California Institute of Technology

stat.ML, cs.LG, math.OC

Submitted: 2026-08-09

Updated: 2026-08-11

Comments: 377 pages

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

Importance score: 68/100

The gist: This paper is a book-length treatment (369 pages plus front matter) titled "Machine Learning for Inverse Problems and Data Assimilation" by Eviatar Bach (University of Reading), Ricardo Baptista

Key concepts

Inverse Problem
An inverse problem occurs when you have measurements and need to determine the cause, such as figuring out the structure of tissue from an X-ray scan. The goal is to work backward from observed data to identify unknown inputs or parameters.
Data Assimilation
This involves blending a mathematical model of a system with real-time observations over time. A practical example is weather forecasting, where new measurements are continuously incorporated into the model to improve predictions.
Amortization
This concept involves training a machine learning model once on large amounts of data. After training, the the resulting mapping can be reused cheaply for new observations, avoiding the need to solve a complex problem from scratch every time.
Transport Perspective
This is a unifying mathematical concept that connects various methods. It involves mapping one probability distribution into another—for example, transforming a simple prior knowledge distribution into the complex posterior distribution derived from data.

Terminology

Summary

This paper is a book-length treatment (369 pages plus front matter) titled Machine Learning for Inverse Problems and Data Assimilation by Eviatar Bach (University of Reading), Ricardo Baptista (University of Toronto), Daniel Sanz-Alonso (University of Chicago), and Andrew Stuart (California Institute of Technology). The paper is organized into three parts: Part I on Inverse Problems (Chapters 1–5), Part II on Data Assimilation (Chapters 6–10), and Part III on Fundamentals (Chapters 11–15).

The paper states its aim as follows: "The aim of this book is to demonstrate the potential for ideas in machine learning to impact on the fields of inverse problems and data assimilation. The perspective is one that is primarily aimed at researchers from inverse problems and/or data assimilation who wish to see a mathematical presentation of machine learning as it pertains to their fields."

The structure is described: Part I of the book is concerned with inverse problems, employing material from Part III; Part II of the book is concerned with data assimilation, employing material from Parts I and III. The authors note that "the material allows application-oriented readers to proceed directly to the use of machine learning methodology, in either (or both) of inverse problems (Part I) and data assimilation (Part II), dipping into the fundamental material in Part III only as it arises and is used. Theory-oriented readers may instead prefer to work through the fundamentals in Part III first."

The paper explains its organizational logic: "Parts I and II, concerning inverse problems and data assimilation respectively, are organized to parallel one another. Both start with the introductory Chapters 1 and 6, framing the fields of inverse problems and data assimilation respectively. We adopt a Bayesian perspective, deriving optimization approaches from it. And in both chapters we include a section on what is one of the earliest, and most prevalent, uses of machine learning in the fields of inverse problems and data assimilation, namely surrogate modeling."

The four overarching concepts defining the parallel organization are: (1) In Chapters 2 and 7 we discuss the use of ideas from variational Bayes to approximate probability measures from within tractable, and ideally expressive for the task at hand, subclasses of probability measures; (2) Chapters 3 and 8 study the machine learning of prior models, from data; (3) In Chapters 4 and 9 we focus on transport as a unifying concept for the learning of algorithms in inverse problems and data assimilation respectively; (4) "Chapters 5 and 10 discuss amortization, demonstrating how the machine learning ideas developed previously can be generalized to learn dependence on the observations that define the inverse problems or data assimilation task."

Chapter 1 (Bayesian Inversion) formulates the inverse problem: "A forward model G: Rd → Rk specifies output y ∈ Rk from input u ∈ Rd by the relationship y = G(u). The related inverse problem is to recover unknown parameter u ∈ Rd from data, or observation, y ∈ Rk defined by y = G(u) + eta; here eta ∈ Rk denotes observation noise." The chapter presents Bayes Theorem giving the posterior density, discusses finite-dimensional summaries including MAP estimators and posterior expectations, establishes well-posedness results using the Hellinger distance, discusses model error in three forms (data space, parameter space, and parameterized forward models), and introduces surrogate modeling. The paper notes: One of the earliest, and still one of the most prevalent, uses of machine learning in the field of computational inverse problems is surrogate modeling.

Chapter 2 (Variational Inference) shows how this theorem can be formulated variationally as a minimization problem over the set of all probability density functions. The key result is that the posterior minimizes J(q) = DKL(q‖rho) − Eq[log l(y·)]. Two canonical approaches are presented: mean-field variational inference and Gaussian variational inference. The chapter discusses mode-seeking versus mean-seeking behavior and the evidence lower bound (ELBO).

Chapter 3 (Learning the Prior) covers Bayesian hierarchical modelling, empirical priors via Dirac measures, and priors via pushforward maps. The paper states: "In Bayesian inference the prior acts as a form of regularization... The prior should, of course, be chosen to reflect prior knowledge. But defining a prior, or regularizer, can be challenging... the abundance of data that is available in many fields suggests the possibility of employing purely data-driven priors for inversion."

Chapter 4 (Transport to the Posterior) addresses various transport approaches to determining the posterior distribution. It establishes existence of prior-to-posterior maps, formulates learning objectives (population and empirical), analyzes the effect of transport error, connects to variational inference, presents sample-based objective functions using energy distance, and discusses transport-assisted MCMC.

Chapter 5 (Data Dependence of the Posterior) focuses on learning the dependence of the posterior distribution piy(u), for parameter u, on the observation y. The paper explains: "The resulting learned model can then be repeatedly used for different realizations of y, rather than having to relearn a different approximate posterior for each y. This reduces the cost of multiple inference procedures to training a single model. Hence, the approach is known as amortized inference or simply amortization." Topics include amortized point estimators (MAP and posterior mean), amortized variational inference, learning likelihood models, and learning regularizers and priors.

Chapter 6 (Filtering and Smoothing Problems) provides an introduction to data assimilation studying both the filtering and the smoothing problems. The stochastic dynamics model is v†j+1 = Ψ(v†j) + xi†j with data model y†j+1 = h(v†j+1) + eta†j+1. The paper defines: The filtering problem is to find, and update sequentially in j, the probability densities pij(v†j):= P(v†jY†j) on Rd for j = 1,..., J. And: The smoothing problem is to find the probability density ΠY†(V†):= P(V†Y†) on Rd(J+1) for some fixed integer J. Algorithms covered include the Kalman filter, steady-state Kalman filter, 3DVar, extended Kalman filter, unscented Kalman filter, ensemble Kalman filter (EnKF), bootstrap particle filter, optimal particle filter, Rauch–Tung–Striebel smoother, 4DVar, ensemble Kalman smoother, and reanalysis. The chapter also covers model error, surrogate modeling, and generalizations.

Chapter 7 (Variational Inference for Data Assimilation) uses the variational formulation of Bayes Theorem, introduced in Chapter 2, to frame data assimilation problems as optimization problems over probability measures; both smoothing and filtering problems are considered. It also shows how filtering may be linked to smoothing by imposing a filtering structure on a variational formulation of smoothing.

Chapter 8 (Learning the Prior for Data Assimilation) is devoted to learning priors for both smoothing and filtering problems, paralleling the material in Chapter 3. The paper notes: Learning priors for data assimilation is intimately related to the important problem of jointly estimating the state and parameters of a stochastic dynamics model. Methods covered include expectation maximization (EM), Monte Carlo EM, auto-differentiable Kalman filters, learning priors for filtering steps from samples, and data assimilation in latent spaces.

Chapter 9 (Transport Perspective on Filters) is devoted to the transport perspective on filtering. The paper explains: Here, the term transport refers broadly to a map designed to represent aspects of the analysis step within the filtering update. It covers state estimation learning objectives (deterministic and probabilistic scoring rules, matching data), algorithms (learning gains in 3DVar and EnKF, learning localization and inflation, optimal particle filter), probabilistic estimation objectives (matching filtering distributions via transport, generalized transport, variational Bayes), and algorithms.

Chapter 10 (Data Dependence of Filters and Smoothers) focuses on amortization: the idea that, if we learn dependence of an algorithm on the observations, it may be re-used multiple times and hence the initial cost of training is amortized. Topics include amortized 3DVar and 4DVar, amortized posterior mean estimation, variational inference approaches, learning via scoring rule minimization, and online transport approaches.

Chapter 11 (Metrics, Divergences, and Scoring Rules) introduces "various metrics and divergences to quantify closeness between probability measures in P(Rd)... We also introduce scoring rules, which quantify closeness of a probability measure in P(Rd) to a point in the underlying sample space Rd, and deterministic scoring rules, which quantify the discrepancy between two points in Rd. The paper notes: Our motivation for studying this topic is twofold. First, some metrics and divergences provide a language in which to state theoretical results about inverse problems and data assimilation. Second, other metrics and divergences are used to define loss functions for probabilistic machine learning models based on data."

Chapter 12 (Supervised Learning) is concerned with the supervised learning task of approximating, from data, a function psi†: D → R. It covers neural networks (with universal approximation results), random features, Gaussian processes (including RKHS theory, regression, and the representer theorem), and approximation properties.

Chapter 13 (Unsupervised Learning and Generative Modeling) is concerned with the subject of unsupervised learning under the data assumption that we have available data in the form U:= u(n) Nn=1, assumed to be drawn i.i.d. from (unknown) probability density function ϒ on Rd. Topics include density estimation, transport methods, normalizing flows, neural ODEs, score-based approaches (Langevin equation, score-matching, denoising score-matching), autoencoders, variational autoencoders, and generative adversarial networks.

Chapter 14 (Time-Series Forecasting) is devoted to the study of time-series forecasting, with an emphasis on data-driven approaches. Topics include ergodic and stationary processes, Lyapunov exponents, analog forecasting (basic and kernel), recurrent structure (Markovian prediction, memory and prediction, reservoir computing), and autoregressive models (linear, transport-based, and learning conditional distributions via strictly proper scoring rules).

Chapter 15 (Optimization and Sampling) provides an introduction to various topics in optimization and sampling that are particularly pertinent to machine learning, inverse problems and data assimilation. Topics include gradient descent (deterministic and stochastic), Newton and Gauss–Newton methods, ensemble Kalman inversion, expectation maximization, Markov chain Monte Carlo (Metropolis–Hastings and Gibbs sampler), and automatic differentiation (forward mode, reverse mode, and reparameterization trick).

The paper adopts a Bayesian perspective throughout. The fundamental Bayesian inverse problem is: given prior rho(u) and likelihood l(yu) = nu(y − G(u)), the posterior is piy(u) = (1/Z)nu(y − G(u))rho(u) where Z is the normalization constant. The paper emphasizes that the primary computational challenge associated with Bayesian inversion is that the posterior is an infinite-dimensional object, even when u is finite-dimensional.

For data assimilation, the paper notes: The smoothing problem is an example of a single Bayesian inverse problem of the type defined in Chapter 1. The filtering problem comprises J Bayesian inverse problems, interleaved with predictions of the model.

The paper establishes extensive notation including: sets (N, Z+, Rd, R+), vector spaces (Euclidean norms, positive definite matrices, outer products), probability (probability measures, density functions, pushforwards, convolutions, expectations, characteristic functions, Gaussian distributions), calculus (gradients, Hessians, Jacobians, divergence), and functions (composition, identity, indicator functions, rank functions).

The paper acknowledges support from various funding agencies and notes: All the authors thank the students who took ACM270 at Caltech in Spring 2024 and the TA Yixuan Wang. The lecture notes for this course constituted an early draft of this book.

Improvements for AI systems

Based on the scientific paper, here are specific improvements that can be made to AI systems:

Improvement: Implement a Bayesian framework for inverse problems where the AI system outputs a full posterior distribution over solutions, not just a point estimate.

What the improved system can do:

  • Given noisy observations y and a forward model G, return a probability distribution over possible parameters u rather than a single reconstruction

  • Quantify uncertainty in medical imaging reconstructions (e.g., MRI from partial measurements)

  • Provide credible intervals for parameter estimates in scientific applications

  • Distinguish between well-posed and ill-posed problems by reporting posterior variance

Abstract

The aim of this book is to demonstrate the potential for ideas in machine learning to impact on the fields of inverse problems and data assimilation. The perspective is one that is primarily aimed at researchers from inverse problems and/or data assimilation who wish to see a mathematical presentation of machine learning as it pertains to their fields. As a by-product, we include a succinct mathematical treatment of various fundamental underpinning topics in machine learning, and adjacent areas of (computational) mathematics.

Sources

Related papers