跳到主要内容
返回时间线
arXiv来源发表:

掩码扩散模型不再只靠序列长度算样本量:隐结构决定训练样本,树混合下多项式即可

核心概要

基于Transformer的掩码扩散模型首次获得端到端统计保证:当数据分布可由隐马尔可夫树混合近似时,达到KL精度所需的训练样本量由隐结构复杂度决定,而非随序列空间规模指数增长。分析同时覆盖并行采样带来的并行化误差与有限样本下掩码预测器的估计误差,且Transformer架构与采样过程无需知道树拓扑或分量参数,能自动适应未知隐结构。

AI-generated editorial illustration: Transformer-Based Masked Diffusion Models Provably Learn Distributions with Latent Structure

深度剖析

掩码扩散模型的样本量可以按数据的隐结构来算,而不是按整个序列空间的大小来算。 此前的统计界要么依赖目标支撑集大小,要么依赖完整序列空间规模,即使数据存在紧凑隐表示,仍可能随序列长度指数增长。 在数据分布可由隐马尔可夫树混合良好近似、且并行化误差足够小的条件下,基于Transformer的掩码扩散模型以KL散度达到精度所需的训练样本量,其结构复杂度在序列长度、词表大小、隐状态数与混合分量数上均为多项式。

保证把训练阶段和采样阶段放进同一个误差分解里,同时算上并行解掩码带来的误差和有限样本训练带来的误差。 已有理论通常把掩码预测器视为已知或给定误差,只分析采样,或只分析得分估计,不刻画学习预测器的统计代价。 端到端采样误差由三项构成:只依赖数据分布与解掩码调度的一阶并行化误差,以及由掩码预测器估计误差乘以序列长度换算得到的联合KL项;并行化误差可用对偶总相关刻画,对星形分量结构该量可远小于序列长度。

掩码预测器用稀疏Transformer实现,一次前向就能给出所有被掩位置的边缘条件分布。 直接学习被掩位置的联合条件分布在计算上不可行,实际模型改为学习条件边缘分布并从其乘积中采样。 构造先对单个隐马尔可夫树用并行树收缩与重建在Transformer块内逼近条件对数几率,再用上下文算出的后验混合权重把各分量组合起来;架构预算只依赖序列长度、词表大小等参数,不依赖具体树拓扑与概率表。

启示与展望

该保证适用于数据分布可由隐马尔可夫树混合良好近似、且并行化误差足够小的情形,分析对象是固定基数解掩码调度下的Transformer掩码扩散模型。架构预算只依赖序列长度、词表大小、隐状态数与分量数等参数,不依赖具体树拓扑、分量参数或隐状态,因此模型能自动适应未知隐结构。对高度受限的树,另有基于向上向下消息传递的构造,其深度与树高成正比。

分析假设使用经验风险最小化得到的预测器,而实际训练只做有限步参数更新,优化误差尚未纳入保证。并行化误差项依赖数据分布与解掩码调度,其可计算性取决于对偶总相关等量的估计。向含环潜图模型、置信度解掩码等更一般结构的推广仍是开放问题。