Continuous pushforward models approximate complex input output distributions
Universal Approximation of Measure-to-Measure Operators by Pushforwards
Machine Learning
Summary
Many tasks in machine learning transform whole sets of data into other sets, rather than just changing individual points. The authors study how well models that move each input point to a new point, depending on the entire input set, can mimic any such transformation. They find that these models can't perfectly handle inputs that contain isolated points or atoms, but can approximate any continuous transformation between well-behaved input sets. This result helps explain how advanced models like transformers can universally approximate complex transformations on data distributions.
What this means in practice
- •For machine learning engineers: Design distribution-to-distribution models using continuous pushforwards that can approximate complex data transformations beyond standard pointwise methods.
- •For natural language processing developers: Improve transformer-based architectures by understanding their theoretical approximation limits of measure-to-measure operators with cross-attention mechanisms.
A theory result. No direct application yet.
Authors
Takashi Furuya, Nicholas H. Nelsen, Frank Cole
Abstract
Many learning tasks map an input distribution to an output distribution. A natural way to model such an operator is to transform each input sample using a continuous function that may depend on the entire input distribution, and then take the distribution of the transformed samples. This defines a measure-dependent pushforward model and includes measure-theoretic formulations of transformers. We ask when such models can approximate arbitrary continuous operators between spaces of probability measures. We first show that universal approximation fails when atomic inputs are allowed: some continuous measure-to-measure operators that split or redistribute atomic mass cannot be approximated arbitrarily well by deterministic pushforward models. We then introduce the uniform level set condition, which requires a continuous measure-dependent scalarization whose shrinking level set neighborhoods carry uniformly vanishing mass over the input family. This condition is satisfied, in particular, by compact families of absolutely continuous measures. On every compact family satisfying this condition, we prove that any continuous measure-to-measure operator with outputs of finite $p$-th moment can be uniformly approximated, in the $p$-Wasserstein distance, by continuous measure-dependent pushforwards. Combining our theorem with existing approximation results for measure-dependent in-context maps yields universal approximation by measure-theoretic transformers. We also extend the framework to continuously-varying source measures, yielding a corresponding universality result for a class of pushforward models that are closely aligned with cross-attention architectures.