Contexte et objectif

Les modèles de diffusion latente (LDM) séparent compression (VAE) et génération (DiT). Cette double chaîne impose un coût d’attention quadratique, surtout pour des résolutions élevées. Linum v2, basé sur un VAE 8×8, nécessitait 110 000 tokens pour un clip 720p de 5 s, ce qui ralentit fortement l’entraînement. L’objectif de JiT‑DDT est de supprimer le VAE, de réduire la fenêtre d’attention et d’augmenter la résolution d’image tout en diminuant le temps GPU.

Architecture du modèle JiT‑DDT

JiT‑DDT utilise un DiT actif en espace pixel, avec 2,5 M paramètres et 320 tokens d’image (64 pour l’encodeur, 256 pour le décodeur). La tokenisation repose sur une patchification 4×4 px → 48‑dim token, puis une projection linéaire 

W ∈ ℝ^{12×48}
qui compresse chaque patch à 12 dimensions. À 512 px, l’image devient 256 tokens, chaque token partant de 3 072 dimensions (32×32×3) avant la réduction. Cette «unified model» intègre la compression directement dans le DiT, éliminant le besoin d’un VAE séparé et autorisant une réduction de tokens de 32×32.

Analyse des performances

Comparé à Linum v2 (2,0 B, 256×256, VAE + DiT, 256 tokens latents), JiT‑DDT atteint 512×512 avec 4 fois plus de pixels tout en consommant 3,6 × moins d’heures GPU. Le tableau fourni indique que le nombre d’échantillons vus est 4,2 × inférieur, traduisant une efficacité d’apprentissage supérieure. Malgré l’augmentation de la résolution, le coût d’attention reste maîtrisé grâce à la réduction du nombre de tokens (256 vs 256 latent tokens) et à la suppression du passage VAE, où le gradient était bloqué.

Le gain provient de deux leviers : (i) la compression agressive intégrée, qui diminue la complexité quadratique O(N²) de l’attention, et (ii) la capacité du DiT à apprendre directement à partir de représentations pixel‑space, évitant la perte d’information introduite par la reconstruction VAE (loss ‖x‑x̂‖² + β·KL(q‖N)).

Limites et perspectives

Le modèle montre une perte de détails fins lorsqu entraîné sans les ajustements proposés, d’où l’ajout d’un décodeur dédié. La recherche indique que les DiT peinent à absorber des entrées de haute dimension, justifiant le besoin d’une architecture «encoder‑decoder» séparée. De plus, le code et les poids sont publiés sous licence Apache 2.0, mais l’auteur précise qu’il s’agit d’un artefact de recherche, non d’un produit final. Les prochains travaux (Linum v3) devront vérifier la scalabilité de cette approche sur des vidéos, où la fenêtre d’attention devient encore plus critique.