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

掩码扩散模型的样本量不再随序列长度爆炸:只要数据有隐结构,训练样本只随隐状态数多项式增长

导语

掩码扩散模型首次拿到端到端统计保证:当数据分布可被隐马尔可夫树混合逼近时,基于Transformer的掩码预测器用随隐状态数、混合分量数多项式增长的训练样本,就能在KL散度下达到给定精度,且无需知道树结构或分量参数。

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

正文

掩码扩散模型的采样误差现在可以整体算清,而不是把并行采样误差和有限样本估计误差分开处理。 此前的理论要么把掩码预测器当作已知,要么只给出随目标支撑集大小或整个序列空间大小增长的界,即使数据有紧凑的隐表示也可能随序列长度指数增长。 端到端保证把并行化误差与掩码预测器的有限样本估计误差合并进同一个KL散度上界,其中估计误差项由预测器参数数量与样本量之比控制。

保证的样本量由数据的隐结构决定,而不是由序列空间的大小决定。 数据分布用隐马尔可夫树混合来建模,它用分量内的隐依赖和分量间的异质性刻画结构,且数据分布不必精确属于这一类。 当数据分布可被该混合逼近、并行化误差足够小时,结构复杂度对序列长度、词表大小、隐状态数和混合分量数都是多项式的;在隐空间远小于词表与长度的情形下,达到给定KL精度所需的训练样本随隐状态数多项式增长。

Transformer掩码预测器和并行采样过程都不需要事先知道树拓扑或分量参数。 网络架构只依赖给定的预算参数,不依赖隐变量、树结构或概率表,同一网络类对所有此类分布一致地逼近被掩码位置的条件下分布。 训练目标只用观测数据,预算参数在逼近误差与估计误差之间做权衡;对分量树高度有界的结构,另一套构造把复杂度项换成与树高相关的量,深度正比于树高。

接下来

这套保证面向离散序列数据,训练与推理使用同一个固定基数解掩码调度,数据分布需可被隐马尔可夫树混合逼近。读者对象是研究离散扩散模型统计性质与Transformer逼近能力的研究者。对分量树高度有界的结构,另一套构造把复杂度项换成与树高相关的量,深度正比于树高。

分析假设使用经验风险最小化得到的预测器,而实际训练只做有限次参数更新。保证还要求并行化误差足够小,该误差只依赖数据分布与推理阶段的解掩码调度。把分析推广到带环的隐图模型、基于置信度的解掩码以及计入优化误差,被列为尚未解决的问题。

来源