Measure-to-measure Regression with Transformers

Abstract: Many learning problems require predicting how populations evolve under an unknown transformation. A natural representation for such populations is a probability measure, with point clouds as a key example. In this work, we study the measure-to-measure (M2M) regression problem, in which one seeks to learn a map between probability measures from a finite collection of observed input-output pairs. In contrast to classical regression, where individual samples are transformed independently, M2M regression treats entire distributions as the data points. This perspective is vital in certain scientific applications, for example, cellular and molecular biology, where cells are known to evolve not as independent data points but as a collection. However, few existing approaches address the problem of M2M regression with sufficient expressivity and scalability. We present a formalization of nonlinear M2M regression and introduce two easy-to-use, expressive, and scalable approaches to learn such operators: transformers as static M2M maps and transformers as dynamic M2M velocity fields. Our approach leverages the natural measure-dependent and mean-field structure of transformers to learn nonlinear M2M maps on the space of probability distributions. We illustrate the effectiveness of our proposed method to generalize to unseen measures on synthetic experiments, interacting particle systems, and a large-scale patient-derived organoid dataset for predicting treatment response in colorectal cancer.

• The paper introduces a novel framework where transformers serve as effective nonlinear regression operators between probability measures.
• Dynamic M2M regression, modeled via continuous flow matching, minimizes distributional losses and outperforms classical methods in diverse tasks.
• Empirical results show high correlation in biological applications and robust scalability across synthetic and high-dimensional population dynamics.

Measure-to-Measure Regression with Transformers: A Formal and Empirical Analysis

Problem Formalization and Motivation

The paper "Measure-to-measure Regression with Transformers" (2605.28075) introduces a principled framework to learn nonlinear regression operators between probability measures, focusing on the scenario where the fundamental object is a distribution rather than an individual sample. This abstraction is particularly relevant in scientific domains such as single-cell biology and high-dimensional physical systems, where the collective dynamics and responses are best characterized as transformations between population-level distributions. Classical regression approaches treat each observation independently, whereas measure-to-measure (M2M) regression models entire distributions as atomic units, demanding operators that map one probability measure to another.

This paradigm is formalized mathematically by defining the regression operator $$\mathcal{T}^*: \mathcal{P}(\mathbb{R}^d) \rightarrow \mathcal{P}(\mathbb{R}^d)$$, with the empirical data consisting of finite pairs of observed input-output measures $(\mu_i, \nu_i)$. The goal is to learn $\mathcal{T}^*$ from these samples so as to extrapolate to unseen distributions. The expressiveness and scalability of the operator class are central to the framework, particularly given the geometric and mean-field structure inherent in many real-world phenomena.

Transformers as Nonlinear M2M Operators

A key technical insight is that transformer architectures are uniquely suited to serve as nonlinear M2M operators due to their inherent ability to process sets (empirical measures) via permutation-invariant self-attention mechanisms. The paper leverages both recent theoretical results (universal approximation of nonlinear measure-dependent operators by deep transformers under mild regularity conditions) and practical architecture design to instantiate transformers as regression maps on $\mathcal{P}(\mathbb{R}^d)$.

Two distinct approaches are introduced:

1. Static M2M Regression: The operator is realized as a one-step pushforward map, $$F_\mu^\theta(\cdot): \mathbb{R}^d \rightarrow \mathbb{R}^d$$, parameterized by the input measure $\mu$. The transformer computes a new empirical measure via $F_\mu^\theta$, with training objectives minimizing distributional losses (e.g., $\mathcal{W}_1$, $\mathcal{W}_2$, MMD, ED) between source and target measures.
2. Dynamic M2M Regression (Transformer Flow Matching, M2M-TFM): In this more expressive paradigm, the operator is modeled as a continuous-time flow, where the transformer parameterizes a measure-dependent vector field $\chi_t^\theta[\mu(t)]$. The empirical measure evolves according to the continuity equation $(\mu_i, \nu_i)$0, with the learned operator integrating this vector field from source to target measures. Training is performed via extensions of flow matching and conditional flow matching, reducing the complexity of optimizing distributional losses.

Figure 1: Diffusion process – visualization of model predictions on multi-measure objects, demonstrating M2M-TFM's ability to reverse the dynamics for unseen empirical measures.

Empirical Evaluation and Numerical Results

The paper conducts rigorous empirical studies across synthetic and real-world tasks, evaluating both static and dynamic transformer-based M2M operators against baseline methods, including classical kernel-based regression, conditional flow matching (CFM), Meta Flow Matching (Meta-FM), and state-of-the-art generative alignment techniques.

Synthetic Multi-Measure Objects

• Setup: A dataset of 2D shapes (letters, numbers, objects) corrupted via distributional diffusion/kernels.
• Goal: Learn the reverse map (corruption $(\mu_i, \nu_i)$1 clean) for unseen test measures.
• Results: M2M-TFM achieves the lowest $(\mu_i, \nu_i)$2 and ED distances across all categories, outperforming Meta-FM, CFM, and WFM, with superior generalization to left-out shapes and compositions.

Figure 2: PDO-21 – empirical distribution of energy distances across measure pairs for a test split in the patient-derived organoid dataset.

High-Dimensional Population Dynamics

• McKean-Vlasov Systems: Measures evolve according to nonlinear, distribution-dependent SDEs (Kuramoto, FitzHugh-Nagumo, Atlas).
• Setup: Train/test splits over systems of different initial conditions and dimensions (2, 50, 100).
• Results: Dynamic M2M-TFM consistently yields the lowest errors in $(\mu_i, \nu_i)$3 and ED metrics in high dimensions, demonstrating scalability and robustness against baselines (NMKV, MF-Transformer, Meta-FM).

Biological Application: Patient-Derived Organoids

• Real-world Task: Predict cellular responses to treatments in unseen patients using high-dimensional mass cytometry data (44 markers, 10 patients).
• Results: Both static (loss-tailored) and dynamic M2M transformers achieve high correlation ($(\mu_i, \nu_i)$4) and low $(\mu_i, \nu_i)$5, ED distances, with M2M-TFM showing balanced performance across metrics and robust generalization despite limited patient observations.

Figure 3: Diffusion process – another visualization supporting the superiority of transformer-based regression for complex, distributional tasks.

Theoretical and Practical Implications

The study rigorously formalizes the class of nonlinear, measure-dependent pushforward operators, positioning transformers as a natural architectural choice due to their universality and scalability. Crucially, the dynamic approach aligns with geometric insights from Wasserstein space and mean-field analysis, ensuring maximal expressivity without incurring computational overload.

Empirically, the proposed frameworks outperform both classical and measure-dependent baselines, particularly in high-dimensional and extrapolative settings, highlighting their applicability to complex scientific domains such as single-cell perturbation screens and prediction of population dynamics under interventions.

Strong numerical results underscore transformer-based M2M regression as a scalable, expressive solution for learning distribution-to-distribution maps in settings where classical sample-based regression fails.

Future Directions

While the paper demonstrates strong empirical and theoretical performance, several limitations and open directions remain:

• Extension to non-gradient, stochastic, or unbalanced dynamics in physical/biological systems
• Theoretical guarantees for learning nonlinear operators via vector field models in the flow matching framework
• Integration of static M2M approaches with one-step generative modeling paradigms
• Application to larger or more heterogeneous perturbation datasets from genomics and systems biology

These avenues promise further advances in understanding and deploying measure-dependent architectures for large-scale scientific modeling.

Conclusion

"Measure-to-measure Regression with Transformers" (2605.28075) rigorously establishes transformer architectures as a versatile foundation for nonlinear regression between probability measures, encompassing both static pushforward and continuous-time flow matching regimes. Through formal analysis and strong empirical results across challenging regression and extrapolation tasks, the paper positions these methods as scalable and expressive tools with substantial theoretical and practical relevance for distribution-driven scientific modeling.