🤖 machine learning

Measure-to-measure Regression with Transformers

This paper introduces a formalization of nonlinear measure-to-measure regression and proposes two scalable transformer-based approaches—static maps and dynamic velocity fields—to learn transformations between probability distributions, demonstrating their effectiveness in predicting population evolution across synthetic, physical, and biological datasets.

Original authors: Matthew Vandergrift, Martha White, Yury Polyanskiy, Philippe Rigollet, Lazar Atanackovic

Published 2026-05-28

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Matthew Vandergrift, Martha White, Yury Polyanskiy, Philippe Rigollet, Lazar Atanackovic

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot how to predict how a crowd of people will move.

The Old Way (Classical Regression):
Usually, if you want to predict movement, you look at one person at a time. You say, "If Alice walks left, she will end up here. If Bob walks right, he will end up there." You treat everyone as an independent traveler.

The New Problem (Measure-to-Measure Regression):
But in many real-world situations—like a school of fish, a swarm of bees, or a crowd of cells in a body—people don't move alone. They move as a collective. If the whole crowd shifts, it's because of how they interact with each other. You can't just look at one person; you have to look at the entire "cloud" of people (the distribution) and predict how the whole cloud will reshape itself.

The paper calls this Measure-to-Measure (M2M) Regression. Instead of predicting where one point goes, you are predicting how an entire shape or cloud of points transforms into a new shape.

The Solution: The "Transformer" as a Crowd Manager

The authors propose using a specific type of AI called a Transformer (the same technology behind many modern chatbots) to solve this. But they use it in a special way.

Think of a standard Transformer as a super-organized manager who looks at a list of items and decides what to do with each one based on the others.

The Magic: In this paper, the Transformer doesn't just look at a list; it looks at the entire crowd as a single entity. It asks, "Given the current shape of this crowd, how should every single person move to turn this shape into that new shape?"

The paper introduces two main ways to do this:

1. The "Snap" Method (Static Transformers)

Imagine you have a photo of a crowd in a starting position and a photo of them in a finished position.

How it works: The AI looks at the starting photo and instantly calculates a "magic map" that tells every single person exactly where to jump to get to the destination in one giant leap.
The Analogy: It's like a director shouting, "Everyone, jump to your final spot NOW!" and the crowd instantly teleports to the new formation.
Pros/Cons: It's fast, but calculating that perfect "jump" for everyone at once is mathematically very hard and can be messy.

2. The "Flow" Method (Dynamic Transformers / M2M-TFM)

This is the paper's star performer. Instead of a giant leap, imagine the crowd moving like a fluid or a river.

How it works: The AI doesn't tell people where to jump. Instead, it acts like a wind or a current. It calculates a "velocity field"—a gentle breeze that pushes the crowd.
- At time 0, the breeze is weak.
- At time 0.5, the breeze pushes harder in a specific direction.
- At time 1, the crowd has flowed smoothly into the new shape.
The Analogy: Think of pouring water from a pitcher into a glass. You don't tell every water molecule where to go; you just control the flow of the stream, and the water naturally finds its way to fill the glass.
Why it's better: The paper shows this "flow" method is much more flexible and accurate. It learns the rules of movement rather than just memorizing the start and end points.

Where Did They Test This?

The authors didn't just talk about theory; they tested their "Crowd Manager" on three specific challenges:

Shapeshifting Letters: They took clouds of dots that formed letters (like "A" or "B") and corrupted them (scattered the dots). The AI had to learn how to push the scattered dots back into the correct letter shape. Their "Flow" method was the best at this.
Simulated Physics: They simulated complex systems where particles interact (like the "Kuramoto model," which is like a group of pendulums swinging together). The AI had to predict how the group would evolve over time. Again, the Flow method won.
Real-World Biology (Colorectal Cancer): This is the big one. They used data from patient-derived organoids (tiny, 3D clusters of cells grown from a patient's tumor).
- The Setup: They had data showing how a patient's cells looked before treatment and after treatment.
- The Goal: They trained the AI on 9 patients and asked it to predict the treatment response for a 10th patient it had never seen before.
- The Result: The AI successfully learned the "rules" of how these specific cancer cells react to drugs and could predict the outcome for the new patient better than previous methods.

The Bottom Line

The paper argues that when you are dealing with groups that move together (like cells, particles, or crowds), you shouldn't treat them as individuals. You should treat them as a single, shifting cloud.

By using Transformers to act as either a "Snap" map or a "Flow" current, they created a tool that can learn how to reshape these clouds. Their "Flow" approach (M2M-TFM) is particularly powerful because it breaks the complex problem of moving a whole crowd into a smooth, step-by-step journey, making it easier to learn and more accurate to predict, even for completely new situations.

Technical Summary: Measure-to-Measure Regression with Transformers

Problem Definition

The paper addresses the Measure-to-Measure (M2M) regression problem, which seeks to learn a mapping between probability measures rather than between individual data points. In many scientific domains, such as cellular and molecular biology, populations evolve as collectives rather than independent entities. Consequently, data is naturally represented as point clouds or distributions (measures) $\mu$ and $\nu$ . The goal is to learn an operator $T^*: \mathcal{P}(\mathbb{R}^d) \to \mathcal{P}(\mathbb{R}^d)$ such that $\nu = T^*(\mu)$ , given a finite set of observed input-output pairs $\{(\mu_i, \nu_i)\}_{i=1}^n$ .

Unlike classical regression where samples are transformed independently, M2M regression treats entire distributions as the fundamental data points. A critical challenge is that the learned operator must generalize to unseen measures, including those with different support sizes or smooth limits, rather than merely interpolating between the empirical supports seen during training. Existing approaches often rely on limited function classes (e.g., linear transport maps or kernel methods) that lack expressivity or scalability for complex, nonlinear dynamics.

Methodology

The authors propose leveraging Transformers as operators on probability measures, capitalizing on their inherent measure-dependent and mean-field structure. The self-attention mechanism allows a transformer to process a set of tokens (particles) as a collective, effectively conditioning its output on the empirical distribution of the input.

The paper formalizes nonlinear M2M operators as pushforward maps of the form $T(\mu) = (f_\mu)^\sharp \mu$ , where the transport map $f_\mu$ explicitly depends on the source measure $\mu$ . Two distinct transformer-based paradigms are introduced:

1. Static M2M Regression (One-Step Pushforward)

This approach models the transformation as a direct, single-step map $F_\mu: \mathbb{R}^d \to \mathbb{R}^d$ parameterized by a transformer.

Mechanism: The transformer takes a set of particles defining the input measure $\mu$ and outputs a transformed set of particles representing the target measure $\nu$ .
Training: The model is trained by minimizing a distributional loss $D(F_\mu(x), y)$ over couplings $\gamma_i$ between input and output measures. The authors utilize various losses, including Maximum Mean Discrepancy (MMD), Energy Distance (ED), and Wasserstein distances ( $W_1, W_2$ ), computed via the Sinkhorn algorithm.
Limitation: Optimizing distributional distances in high dimensions is statistically and computationally challenging.

2. Dynamic M2M Regression (Transformer Flow Matching - M2M-TFM)

To overcome the difficulties of direct distributional optimization, the authors introduce a dynamic approach that models the transformation as a continuous-time flow.

Mechanism: Instead of a single map, the transformation is defined by a time-dependent velocity field $\chi_t[\mu(t)]$ that depends on the current state of the measure. The evolution of the measure is governed by a continuity equation: $\partial_t \mu(t) + \nabla \cdot (\mu(t) \chi_t[\mu(t)]) = 0$ . The final output is obtained by integrating this Ordinary Differential Equation (ODE) from $t=0$ to $t=1$ .
Architecture: A time-dependent transformer parameterizes the velocity field $\chi_t$ . The transformer processes the current set of particles and the time step $t$ to predict the velocity vector for each particle.
Training (M2M-TFM): The method employs Conditional Flow Matching (CFM). It trains the velocity field to match the target velocity of a linear interpolation path $\phi_t(x, y) = (1-t)x + ty$ between source and target particles. The loss function minimizes the squared error between the predicted velocity and the target velocity $(y-x)$ , marginalized over the dataset of measure pairs. This converts the difficult distributional optimization problem into a tractable pointwise regression task.
Inference: Given a new input measure, the model simulates the ODE trajectory step-by-step to generate the predicted output measure.

Key Contributions

Formalization: The paper provides a rigorous formalization of nonlinear M2M regression as learning operators on the space of probability measures $\mathcal{P}(\mathbb{R}^d)$ , distinguishing between direct mappings and continuous-time dynamics.
Novel Frameworks: It introduces two scalable, expressive approaches:
- Static Transformers: Using standard transformers to approximate the M2M map directly.
- M2M-TFM: A dynamic framework that integrates flow matching with the measure-dependent nature of transformers, eliminating the need for explicit distributional loss optimization during training.
Empirical Validation: The authors demonstrate the effectiveness of these methods across three distinct settings:
- Synthetic Multi-Measure Objects: Predicting the reverse of diffusion and kernel interaction processes on letter shapes and multi-object compositions.
- McKean-Vlasov Systems: Predicting the evolution of population dynamics in 2D, 50D, and 100D systems (Kuramoto, FitzHugh-Nagumo, Mean-Field Atlas).
- Real-World Application: Predicting treatment response in colorectal cancer using a large-scale patient-derived organoid (PDO) dataset, where the model must generalize to unseen patients.

Results

The proposed methods, particularly the dynamic M2M-TFM, consistently outperform existing baselines across all experimental settings:

Synthetic Objects: M2M-TFM achieved the lowest Wasserstein ( $W_1$ ) and Energy Distance (ED) errors, significantly outperforming Conditional Flow Matching (CFM), Meta-Flow Matching (Meta-FM), and Wasserstein Flow Matching (WFM).
McKean-Vlasov Systems: In high-dimensional (100D) settings, M2M-TFM demonstrated superior generalization and expressivity compared to Neural McKean-Vlasov (NMKV) and Mean-Field Transformers (MF-Transformer), often reducing error by a substantial margin.
Biological Application: On the colorectal cancer dataset, M2M-TFM provided balanced performance across all metrics ( $W_1$ , ED, and $r^2$ ), whereas static methods tended to excel only in the metric corresponding to their specific loss function. M2M-TFM achieved the highest $r^2$ (0.9029) and competitive $W_1$ scores, indicating robust generalization to unseen patients.

Significance and Claims

The paper claims that Transformers are a natural and powerful architecture for M2M regression due to their intrinsic ability to model measure-dependent dynamics. By formalizing the problem through the lens of operators on $\mathcal{P}(\mathbb{R}^d)$ , the authors bridge the gap between theoretical mean-field limits of transformers and practical applications.

The primary significance lies in the M2M-TFM framework, which offers a scalable and expressive solution to the M2M regression problem. By reducing the learning of complex distributional operators to a pointwise regression on velocity fields, the method avoids the "curse of dimensionality" associated with optimizing distributional distances directly. The authors position this work as a versatile foundation for modeling collective dynamics in scientific applications, specifically highlighting its success in predicting cellular evolution and treatment responses where data is inherently distributed.

The paper remains modest regarding theoretical guarantees, noting that while deep transformers are known universal approximators for M2M maps, formal guarantees for learning nonlinear operators via vector field models under the flow matching framework remain an open direction for future work. Similarly, the current implementation assumes linear interpolation paths and deterministic dynamics, leaving extensions to nonlinear, stochastic, and unbalanced settings for future research.