通过非因子化预测加速离散数据生成,减少一半推理步数。
Self-Speculative Masked Diffusions
- 用因果注意力实现单次前向传播的非因子化预测
- 在GPT2和蛋白质序列生成中减少约50%的网络前向次数
- 适合追求高效文本与生物序列生成的研究者
我们提出自洽推测掩码扩散模型,一种用于离散数据的新一代掩码扩散生成模型,显著减少了生成样本所需的函数评估次数。标准掩码扩散模型对当前掩码位置进行因子化逻辑预测,但一次采样过多位置会导致样本质量下降。因此需大量模拟步骤和神经网络前向传播才能生成高质量数据。我们通过生成非因子化预测来降低计算负担:将Transformer最终注意力掩码从非因果改为因果,实现草稿令牌生成,并通过一种新型、模型内嵌的推测采样机制并行验证。这使得在单次前向传播中即可获得掩码位置的非因子化预测分布。我们将该方法应用于GPT2规模的文本建模和蛋白质序列生成,在保持高质量的前提下,相比标准掩码扩散模型,所需网络前向传播次数减少约2倍。
原文摘要 · Abstract (English)
We present self-speculative masked diffusions, a new class of masked diffusion generative models for discrete data that require significantly fewer function evaluations to generate samples. Standard masked diffusion models predict factorized logits over currently masked positions. A number of masked positions are then sampled, however, the factorization approximation means that sampling too many positions in one go leads to poor sample quality. As a result, many simulation steps and therefore neural network function evaluations are required to generate high-quality data. We reduce the computational burden by generating non-factorized predictions over masked positions. This is achieved by modifying the final transformer attention mask from non-causal to causal, enabling draft token generation and parallel validation via a novel, model-integrated speculative sampling mechanism. This results in a non-factorized predictive distribution over masked positions in a single forward pass. We apply our method to GPT2 scale text modelling and protein sequence generation, finding that we can achieve a ~2x reduction in the required number of network forward passes relative to standard masked diffusion models.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。