JiT模型用传统对齐方法会失效,新方法让训练更快更准。
Representation Alignment for Just Image Transformers is not Easier than You Think
- 提出PixelREPA,用掩码适配器重构对齐目标
- FID从3.66降至3.17,收敛速度提升2倍以上
- 适合追求高效高质图像生成的研究者
Representation Alignment(REPA)虽能加速潜在空间扩散Transformer训练,但对像素空间的Just Image Transformers(JiT)却可能失效。随着训练进行,REPA导致JiT的FID上升,并在表示空间紧密聚集的图像子集上损失多样性。问题源于信息不对称:去噪在高维图像空间进行,而语义目标被强烈压缩,使直接回归成为捷径。为此,本文提出PixelREPA,通过掩码变压器适配器将对齐目标转换,并结合浅层变压器与部分令牌掩码实现约束。实验显示,PixelREPA使JiT-B/16的FID从3.66降至3.17,Inception Score(IS)从275.1升至284.6,且收敛速度超2倍;PixelREPA-H/16更达FID=1.81、IS=317.2。代码已开源。
原文摘要 · Abstract (English)
Representation Alignment (REPA) has emerged as a simple way to accelerate Diffusion Transformers training in latent space. At the same time, pixel-space diffusion transformers such as Just image Transformers (JiT) have attracted growing attention because they remove a dependency on a pretrained tokenizer, and then avoid the reconstruction bottleneck of latent diffusion. This paper shows that the REPA can fail for JiT. REPA yields worse FID for JiT as training proceeds and collapses diversity on image subsets that are tightly clustered in the representation space of pretrained semantic encoder on ImageNet. We trace the failure to an information asymmetry: denoising occurs in the high dimensional image space, while the semantic target is strongly compressed, making direct regression a shortcut objective. We propose PixelREPA, which transforms the alignment target and constrains alignment with a Masked Transformer Adapter that combines a shallow transformer adapter with partial token masking. PixelREPA improves both training convergence and final quality. PixelREPA reduces FID from 3.66 to 3.17 for JiT-B$/16$ and improves Inception Score (IS) from 275.1 to 284.6 on ImageNet $256 \times 256$, while achieving $> 2\times$ faster convergence. Finally, PixelREPA-H$/16$ achieves FID$=1.81$ and IS$=317.2$. Our code is available at https://github.com/kaist-cvml/PixelREPA.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。