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

📄 arXiv: 2605.07100v1 📥 PDF

作者: Zhenhan Fang, Aixin Tan, Jian Huang

分类: stat.ML, cs.LG

发布日期: 2026-05-08

备注: 22 pages, 5 figures and 5 tables

💡 一句话要点

提出TRACE框架：利用扩散与流匹配模型实现基于传输对齐的共形预测

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 共形预测 扩散模型 流匹配 不确定性量化 生成式建模 传输对齐

📋 核心要点

  1. 现有共形预测方法受限于几何假设或需显式似然计算，难以处理复杂生成模型中的多维输出。
  2. TRACE通过扩散与流匹配模型，利用随机传输轨迹上的误差对齐来构建非一致性评分，无需似然评估。
  3. 实验证明该方法在合成与真实数据集上均能实现有效覆盖，并能灵活适应多模态与非凸条件分布。

📝 摘要（中文）

为多维输出构建有效且信息丰富的共形预测区域仍是当前的核心挑战。尽管共形预测提供了有限样本下的无分布覆盖保证，但其性能高度依赖于非一致性评分（nonconformity score）的选择。现有方法往往依赖于严格的几何假设，或需要显式的似然评估与可逆变换，这限制了其在复杂生成场景下的应用。本文提出了TRACE（传输对齐共形估计），这是一种通过扩散和流匹配模型中的传输对齐来定义非一致性的共形预测框架。该方法不依赖似然评估，而是通过计算沿随机传输轨迹的去噪或速度匹配误差的平均值，来衡量候选输出与学习到的生成动力学之间的对齐程度。由此产生的基于传输的评分是标量值，可通过拆分共形预测进行校准，从而在可交换性假设下实现有效的边际覆盖。实验表明，该方法生成的预测区域能自然适应多模态及非凸条件分布。

🔬 方法详解

问题定义：论文旨在解决高维输出空间中，如何构建既满足统计覆盖保证、又能适应复杂非凸条件分布的共形预测区域。现有方法在处理非参数化或复杂生成模型时，常因无法获取精确似然或难以满足几何假设而失效。

核心思路：TRACE的核心思想是将共形预测与生成模型的动力学过程相结合。它不再依赖于传统的密度估计，而是通过衡量候选点在生成模型（如扩散模型或流匹配模型）的传输轨迹上的“一致性”来定义非一致性评分。

技术框架：该框架首先训练一个生成模型（扩散或流匹配），学习从噪声到数据的传输路径。在推理阶段，对于给定的输入，TRACE通过计算候选输出在模型动力学下的去噪误差或速度匹配误差，将其转化为标量评分，并利用拆分共形预测（Split Conformal Prediction）进行校准。

关键创新：最重要的创新在于引入了“传输对齐”作为非一致性度量。这种方法将生成模型的学习目标（去噪/速度匹配）直接转化为统计推断的度量标准，从而绕过了对复杂分布进行显式似然建模的难题。

关键设计：关键设计在于对随机传输轨迹的积分处理。通过对轨迹上的误差进行平均，模型能够有效捕捉候选点是否位于学习到的数据流形上。此外，该方法对计算预算具有良好的可扩展性，允许在精度与计算成本之间进行权衡。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，TRACE在多个合成数据集及复杂真实数据集上均达到了预期的边际覆盖率。与基线方法相比，TRACE生成的预测区域在处理多模态和非凸分布时表现出更强的几何适应性，且在不同计算预算下均保持了良好的统计稳健性，证明了其在复杂生成模型中的通用性。

🎯 应用场景

TRACE适用于需要高置信度预测的复杂生成任务，如医学图像分割中的不确定性量化、自动驾驶中的轨迹预测，以及科学计算中的多维参数估计。其对多模态分布的适应性使其在处理具有多种可能结果的决策系统时具有显著的实际价值。

📄 摘要（原文）

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.