苹果查明扩散Transformer中"异常token"问题的成因

发现了什么

苹果ML Research于8月公开了一篇正面探讨图像生成中所用扩散Transformer(Diffusion Transformer,DiT)内部"异常token(outlier token)"现象的论文。研究团队确认,在基于表征自编码器(RAE)的DiT流程中,预训练视觉Transformer(ViT)编码器和DiT本身都会出现这种现象。尤其是在DiT内部的中间层,这种token表现得尤为突出。

研究团队发现的有趣之处在于,单纯遮蔽或删除数值较大的异常token,并不能提升性能。这意味着问题并非源于少数几个极端值,而是该位置的局部图像块信息本身已经受损。为解决这一问题,研究团队提出了"双阶段寄存器(Dual-Stage Registers,DSR)"。其方式是:若可使用已训练好的寄存器,则加以利用;若没有,则在推理阶段反复应用寄存器;同时在去噪网络(denoiser)上额外附加专门的扩散寄存器。

为什么"异常token"是个问题

视觉Transformer将图像切分为小块(patch),把每一块都当作一个"token"处理,这与语言模型将句子按词切分处理的方式相同。然而随着训练的进行,已有报告指出:部分token会出现异常大的数值(高范数,high-norm),而该位置实际的图像信息却几乎没有被保留,反而过度吸引了其他token的注意力(attention)。在ViT研究中,已知可以为这类token专门设置"寄存器"位置来吸收这种现象,但在生成模型DiT中这一现象具体如何运作,此前一直未被真正厘清。

DiT是Stable Diffusion 3、FLUX系列等最新图像生成模型采用的架构,用以取代原有的U-Net结构。苹果此前曾持续开展扩散模型系列研究,比如比较基于扩散的方式与自回归方式的性能,或将扩散语言模型扩展至17亿参数规模等。这篇论文可以视为该研究脉络的延伸,找出了侵蚀图像质量的细微缺陷的根本成因。

那么会带来什么改变

异常token问题此前一直被认为是生成图像中偶尔出现的局部色斑或图像块扭曲的原因之一。苹果此次的研究并未止步于临时性的遮蔽处理,而是从训练结构层面提出了应对方案,这意味着未来可能会为提升基于DiT的图像、视频生成模型的稳定性提供设计指南。不过,这一结果目前仅停留在研究论文层面的验证阶段,具体会以何种形式反映到实际商用模型中,还有待后续研究公开后进一步观察。