Alinhamento de Representação para Transformers de Imagem Justos não é Mais Fácil do que Você Imagina

Representation Alignment for Just Image Transformers is not Easier than You Think

March 15, 2026

Autores: Jaeyo Shin, Jiwook Kim, Hyunjung Shim

cs.AI

Resumo

O Alinhamento de Representações (REPA) surgiu como uma forma simples de acelerar o treinamento de Diffusion Transformers no espaço latente. Paralelamente, diffusion transformers no espaço de pixels, como os Just Image Transformers (JiT), têm atraído crescente atenção por eliminarem a dependência de um tokenizador pré-treinado e, assim, evitarem o gargalo de reconstrução da difusão latente. Este artigo demonstra que o REPA pode falhar para JiT. O REPA produz FID pior para JiT à medida que o treinamento prossegue e entra em colapso de diversidade em subconjuntos de imagens que estão fortemente agrupados no espaço de representação do codificador semântico pré-treinado no ImageNet. Rastreamos a falha até uma assimetria de informação: a remoção de ruído ocorre no espaço de imagem de alta dimensão, enquanto o alvo semântico é fortemente comprimido, tornando a regressão direta um objetivo de atalho. Propomos o PixelREPA, que transforma o alvo de alinhamento e restringe o alinhamento com um Adaptador de Transformer Mascarado que combina um adaptador de transformer raso com o mascaramento parcial de tokens. O PixelREPA melhora tanto a convergência do treinamento quanto a qualidade final. O PixelREPA reduz o FID de 3,66 para 3,17 no JiT-B/16 e melhora o Inception Score (IS) de 275,1 para 284,6 no ImageNet 256x256, ao mesmo tempo que alcança uma convergência >2 vezes mais rápida. Por fim, o PixelREPA-H/16 alcança FID=1,81 e IS=317,2. Nosso código está disponível em https://github.com/kaist-cvml/PixelREPA.

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 > 2times 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.