发表机构
National University of Singapore; RIKEN AIP; The University of Tokyo; A*STAR CFAR; Nanyang Technological University（新加坡国立大学; 日本理化学研究所先进智能研究中心; 东京大学; 新加坡科技研究局计算与人工智能研究中心; 南洋理工大学）
机构由 AI 辅助整理，请以论文原文为准。AI 中文总结
本文理论上证明扩散变换器（DiT）在预训练后，能从有限演示中学习并生成捕获任务不确定性的预测分布，并在Hölder类任务上达到极小极大最优速率，从而成为统计上最优的上下文生成器。
AI 中文摘要
生成式基础模型因其能够在推理时根据演示生成所需输出而无需更新参数，引起了人们的兴趣。然而，由于少量演示无法唯一确定预期任务，挑战在于如何学习并采样反映这种任务不确定性的输出分布。在这项工作中，我们从理论上分析了在不同任务上预训练的扩散变换器（DiT）如何从演示中学习并为新查询生成预测分布。我们首先表明，从有限演示中生成的自然目标不是通过估计单个任务得到的输出，而是捕获观察演示后剩余任务不确定性的预测分布。然后我们证明，DiT可以通过分数估计学习这种预测分布，利用注意力从演示中聚合信息，并通过扩散生成样本。凭借这一特性，在足够的预训练资源和扩散采样步骤下，所得到的DiT在测试时任务的Hölder类上达到极小极大最优速率。这些结果表明，DiT充当了一个具有统计基础的上下文生成器，能够生成适应新任务的分布，同时保留有限演示中固有的不确定性。
英文摘要
Generative foundation models are attracting interest for their ability to produce desired outputs from demonstrations given at inference time, without updating parameters. However, since a few demonstrations cannot uniquely identify the intended task, the challenge is how to learn and sample from an output distribution that reflects this task uncertainty. In this work, we theoretically analyze how a Diffusion Transformer (DiT), pretrained across diverse tasks, learns and generates predictive distributions for a new query from demonstrations. We first show that the natural target to generate from finite demonstrations is not an output derived from estimating a single task, but rather a predictive distribution that captures the task uncertainty remaining after observing the demonstrations. We then prove that a DiT can learn this predictive distribution through score estimation, using attention to aggregate information from demonstrations and diffusion to generate samples. Owing to this property, with sufficient pretraining resources and diffusion sampling steps, the resulting DiT achieves the minimax optimal rate over a Hölder class of test-time tasks. These results imply that DiT acts as a statistically grounded in-context generator capable of generating distributions adapted to new tasks while retaining the uncertainty inherent in finite demonstrations.
