Measure-to-measure Regression with Transformers

2026-05-27

AI Summary

Traditional regression methods struggle to model population dynamics where both inputs and outputs are probability measures, such as point clouds. This work addresses this challenge by treating entire measures as fundamental data units and introduces, for the first time, a systematic integration of Transformer architectures—leveraging their mean-field structure and measure-dependent properties—into measure-to-measure (M2M) regression. The proposed approach establishes two scalable and highly expressive nonlinear operator learning frameworks: one based on static mappings and the other on dynamic velocity fields. By fully exploiting the symmetry and mean-field characteristics inherent in Transformers, the method demonstrates strong generalization across synthetic datasets, interacting particle systems, and a large-scale dataset of patient-derived colorectal cancer organoids, successfully predicting population-level responses under previously unseen treatment conditions.

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.

Problem

• measure-to-measure regression
• probability measures
• transformers
• distributional learning
• population dynamics

Innovation

• measure-to-measure regression
• transformers
• probability measures
• mean-field structure
• distributional learning