Conferences

Abstract: Large language models are capable of in-context learning, the ability to perform new tasks at test time using a handful of input-output examples, without parameter updates. We develop a universal approximation theory to elucidate how transformers enable in-context learning. For a general class of functions (each representing a distinct task), we demonstrate how to construct a transformer that, without any further weight updates, can predict based on a few noisy in-context examples with vanishingly small risk. Unlike prior work that frames transformers as approximators of optimization algorithms (e.g., gradient descent) for statistical learning tasks, we integrate Barron’s universal function approximation theory with the algorithm approximator viewpoint. Our approach yields approximation guarantees that are not constrained by the effectiveness of the optimization algorithms being mimicked, extending far beyond convex problems like linear regression. The key is to show that (i) any target function can be nearly linearly represented, with small ℓ1-norm, over a set of universal features, and (ii) a transformer can be constructed to find the linear representation — akin to solving Lasso — at test time. This is joint work with Gen Li, Yuchen Jiao, Yu Huang, and Yuting Wei (arXiv:2506.05200).

Abstract: Diffusion model is a prevailing generative AI technology, and this talk will focus on understanding its performance, for the sake of certifying its results, choosing its hyperparameters, and assessing its applicability to downstream tasks. A major part of the talk will focus on quantifying diffusion model’s generation accuracy. The importance of this problem already led to a rich and substantial literature; however, most existing theoretical investigations assumed that an epsilon-accurate score function has already been oracle-given and focused on just the inference process of diffusion model. I will instead describe a first quantitative understanding of the end-to-end generative modeling protocol, including both score training (optimization) and inference (sampling). The resulting error analysis will lead to insights on how to design the training and inference processes for efficacious generation. Then, diffusion model’s generalization will be studied – when it is not memorizing the training data set, what exactly will it generate? This question is not only pertinent to privacy and copyright considerations, but also  important to the innovation of knowledge. Some quantitative results about diffusion model’s generation inductive biases will be described.

Abstract: Flow-based models have spurred a revolution in generative modeling, driving astounding advancements across diverse domains including high-resolution text to image synthesis and de-novo drug design. Yet despite their remarkable performance, inference in these models requires the solution of a differential equation, which is extremely costly for the large-scale neural network-based models used in practice. In this talk, we introduce a mathematical theory of flow maps, a new class of generative models that directly learn the solution operator for a flow-based model. By learning this operator, flow maps can generate data in 1-4 network evaluations, leading to orders of magnitude faster inference compared to standard flow-based models. We discuss several algorithms for efficiently learning flow maps in practice that emerge from our theory, and we show how many popular recent methods for accelerated inference — including consistency models, shortcut models, align your flow, and mean flow — can be viewed as particular cases of our formalism. We demonstrate the practical effectiveness of flow maps across several tasks including image synthesis, geometric data generation, and inference-time guidance of pre-trained text-to-image models.

Abstract: Probability divergences such as the Kullback–Leibler (KL) divergence provide an information-theoretic measure of discrepancy between two probability distributions, but they require absolute continuity of one distribution with respect to the other. This assumption often fails for empirical measures or distributions supported on low-dimensional structures. In contrast, Wasserstein metrics quantify the transport cost between distributions without requiring absolute continuity, but they can break down when one of the distributions is heavy-tailed. Lipschitz-regularized divergences, introduced in Dupuis and Mao (2022) as a special class of (f, \Gamma)-divergences proposed in Birrell et al. (2022), combine f-divergences and the Wasserstein-1 metric via their variational (dual) formulations and inherit the advantages of both. In this talk, I will present key theoretical properties of Lipschitz-regularized f-divergences and demonstrate how these results lead to robust generative modeling methods under minimal assumptions on the target data distribution, including cases with heavy tails, low-dimensional structure, or fractal-like support.

Abstract: Many datasets in modern applications, from cell gene expression and images to shapes and text documents, are naturally interpreted as probability measures, distributions, histograms, or point clouds. This perspective motivates the development of learning algorithms that operate directly in the space of probability measures. However, this space presents unique challenges: it is nonlinear and infinite-dimensional. Fortunately, it possesses a natural Riemannian-type geometry which enables meaningful learning algorithms.

This talk will provide an introduction to the space of probability measures and present approaches to unsupervised, supervised, and manifold learning within this framework. We will examine temporal evolutions on this space, including flows involving stochastic gradient descent and trajectory inference, with applications to analyzing gene expression in single cells. The proposed algorithms are furthermore demonstrated in pattern recognition tasks in imaging and medical applications.