Measure-to-measure Regression with Transformers

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Vandergrift, Matthew, White, Martha, Polyanskiy, Yury, Rigollet, Philippe, Atanackovic, Lazar
Format: Preprint
Published: 2026
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866911723675451392
author Vandergrift, Matthew
White, Martha
Polyanskiy, Yury
Rigollet, Philippe
Atanackovic, Lazar
author_facet Vandergrift, Matthew
White, Martha
Polyanskiy, Yury
Rigollet, Philippe
Atanackovic, Lazar
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2605_28075
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle Measure-to-measure Regression with Transformers
Vandergrift, Matthew
White, Martha
Polyanskiy, Yury
Rigollet, Philippe
Atanackovic, Lazar
Machine Learning
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.
title Measure-to-measure Regression with Transformers
topic Machine Learning
url https://arxiv.org/abs/2605.28075