TRACE: Transport Alignment Conformal Prediction via Diffusion and Flow Matching Models

2026-05-07

AI Summary

This work addresses the reliance on strong geometric assumptions or explicit likelihood evaluations in conformal prediction for multidimensional outputs by proposing a novel nonconformity measure based on transport dynamics alignment. Leveraging stochastic transport trajectories from diffusion models and flow matching, the method constructs a scalar scoring mechanism through trajectory-averaged denoising or velocity matching errors, eliminating the need for likelihood computation or invertible transformations. The resulting framework is compatible with complex generative models and adaptively captures multimodal and non-convex conditional structures. Empirical evaluations on both synthetic and real-world datasets demonstrate valid marginal coverage, with predictive performance improving smoothly as computational budget increases.

Abstract

Constructing valid and informative conformal prediction regions for multi-dimensional outputs remains a fundamental challenge. While conformal prediction provides finite-sample, distribution-free coverage guarantees, its practical performance critically depends on the choice of nonconformity score. Existing approaches often rely on restrictive geometric assumptions or require explicit likelihood evaluation and invertible transformations, limiting their applicability in complex generative settings. In this work, we introduce TRACE (TRansport Alignment Conformal Estimation), a conformal prediction framework that defines nonconformity through transport alignment in diffusion and flow matching models. Rather than evaluating likelihoods, we measure how well a candidate output aligns with the learned generative dynamics by averaging denoising or velocity-matching errors along stochastic transport trajectories. The resulting transport-based scores are scalar-valued and can be calibrated using split conformal prediction, yielding valid marginal coverage under exchangeability. We further analyze the statistical properties of the proposed scores and their sensitivity to computational budget. Experiments on synthetic and real datasets demonstrate valid coverage and show that the resulting regions adapt naturally to multimodal and non-convex conditional distributions.

Problem

• conformal prediction
• multi-dimensional outputs
• nonconformity score
• generative models
• coverage guarantees

Innovation

• conformal prediction
• transport alignment
• diffusion models
• flow matching
• nonconformity score