Active learning improves conditional generative modeling for distributional data
Active Learning of Conditional Generative Models via the Transport Neural Tangent Kernel
Machine Learning
Summary
Many scientific experiments produce complex output data that cannot be summarized by simple averages. The authors explain why current methods that focus only on average outputs miss important details about the shape of the output data. They introduce a new mathematical tool called the transport neural tangent kernel (tNTK) that better captures these distributional differences. Using this tool, they show how to select new experiments more efficiently to learn these complex data patterns. Their tests show that this approach outperforms older methods on various tasks including biology and language data.
What this means in practice
- •For biomedical data teams: Design experiments that predict cellular responses to new perturbations more efficiently by selecting input conditions sensitive to distributional changes.
- •For machine learning engineers: Improve training of conditional generative models for text or distributional outputs by actively selecting data points that reduce uncertainty in predicted output distributions.
Authors
Jayoung Ryu, Kyunghyun Cho, Romain Lopez
Abstract
Many high-throughput scientific experiments produce distributional outputs, requiring conditional generative models to interpret the data, and to predict outcomes at new inputs. Acquiring experimental data at new conditions (e.g., for a new antigen or a new cell type) is often expensive, motivating the need for automated experimental design methods. However, existing active learning frameworks are limited as they assume scalar or vector-valued outputs. We first show that the Wasserstein test risk of a conditional generative task decomposes into a mean term and a shape term, demonstrating why mean-based active learning is blind to distributional shape. To capture both, we propose the transport Neural Tangent Kernel (tNTK), a tractable kernel that measures how sensitive a generative model's predicted distribution is with respect to its parameters. The tNTK recovers the standard NTK in the deterministic limit and upper-bounds the posterior Fréchet variance of the prediction under small parameter perturbations. It admits a closed-form expression for affine Gaussian transport and tractable Monte Carlo estimators otherwise. Treating the tNTK as the covariance function of a Gaussian process, we greedily select inputs that minimize the total GP posterior variance across the pool, yielding a distribution-aware active learning strategy. We evaluate this strategy across four classes of conditional generative models, including conditional flow matching and conditional Monge gap, on synthetic distributions, intent-labeled sentences, and single-cell perturbation responses. In these experiments, the tNTK outperforms the baseline NTK computed from the conditional mean, demonstrating the value of distribution-aware active learning for generative-modeling tasks.