CaTok: Taming Mean Flows for One-Dimensional Causal Image Tokenization

Each language version is independently generated for its own context, not a direct translation.

Imaginez que vous essayez d'apprendre à un ordinateur à dessiner des images, exactement comme un humain dessine : d'abord un croquis grossier, puis des détails, puis des couleurs, toujours dans le bon ordre. C'est ce que les modèles de langage (comme ceux qui écrivent des textes) font très bien : ils prévoient le mot suivant, un par un.

Mais pour les images, c'est plus compliqué. Les ordinateurs ont du mal à comprendre l'ordre naturel d'une image. Le nouveau papier que vous avez partagé, CaTok, propose une solution élégante pour résoudre ce problème.

Voici l'explication simple, avec quelques analogies :

1. Le Problème : Le casse-tête de l'ordre

Imaginez que vous devez décrire une photo de chat à un ami au téléphone.

Les anciennes méthodes (2D) : Elles prennent la photo, la découpent en mille petits carrés, et les mélangent dans un sac. L'ordinateur doit deviner quel carré va avec quel autre sans ordre précis. C'est comme essayer de reconstruire un puzzle en jetant les pièces au hasard sur la table.
Les méthodes "causales" (comme le texte) : Elles veulent que l'ordinateur dessine ligne par ligne, de gauche à droite. Mais les images sont complexes : comment décider si on dessine d'abord l'œil ou la moustache ?

2. La Solution de CaTok : Le "Train à Vitesse Moyenne"

CaTok transforme l'image en une liste unique de 256 petits morceaux (des "jetons" ou tokens), comme des wagons d'un train. Mais la vraie magie réside dans la façon dont ils sont assemblés.

L'équipe utilise une technique appelée MeanFlow (Flux Moyen). Voici l'analogie :

L'ancien problème : Imaginez que vous conduisez une voiture. Si vous ne regardez que votre vitesse exacte à l'instant T, vous risquez de faire des à-coups ou de rater un virage. C'est ce que faisaient les anciens modèles : ils regardaient chaque instant séparément, ce qui créait des erreurs.
L'approche CaTok : Au lieu de regarder la vitesse instantanée, CaTok regarde la vitesse moyenne sur un trajet.
- Imaginez que vous devez aller du point A (le bruit, une image floue) au point B (l'image nette).
- CaTok ne dit pas "allez vite maintenant, ralentissez après". Il dit : "Sur ce tronçon de route (entre le temps r et le temps t), voici la direction moyenne à prendre".
- Cela permet au modèle de faire le trajet en une seule étape (très rapide !) ou en plusieurs étapes (très précis), selon ce dont on a besoin.

3. La Causalité : Construire l'image brique par brique

Le grand atout de CaTok est qu'il apprend la causalité.

Analogie du livre : Quand vous lisez un livre, vous ne pouvez pas comprendre la page 100 sans avoir lu la page 1. Les mots ont un ordre.
Chez CaTok : Les 256 "wagons" de l'image sont aussi dans un ordre strict. Le premier wagon contient les grandes formes (le ciel, le sol), le suivant ajoute des contours, et les derniers ajoutent les détails fins (les poils du chat, les reflets).
Si vous ne gardez que les 16 premiers wagons, l'ordinateur voit une image floue mais reconnaissable. Si vous ajoutez les suivants, l'image devient de plus en plus nette. C'est comme si l'image se révélait progressivement, du grossier au fin.

4. L'Entraînement : Le "Professeur de Dessin" (REPA-A)

Pour apprendre à l'ordinateur à faire cela, les auteurs ont ajouté un petit truc astucieux appelé REPA-A.

L'analogie : Imaginez un élève qui apprend à dessiner. Au début, ses dessins sont moches. S'il regarde un tableau de maître (un modèle d'intelligence très puissant appelé "Fondation Visuelle"), il peut comparer son dessin au tableau et se corriger.
Ce que fait REPA-A : C'est ce "tableau de maître". Il aide l'ordinateur à comprendre ce qu'est une "vraie" image pendant l'entraînement, ce qui rend l'apprentissage beaucoup plus rapide et stable.

En résumé, pourquoi c'est génial ?

Vitesse : Grâce à la "vitesse moyenne", CaTok peut générer une image en une seule seconde (une étape) sans perdre trop de qualité. C'est comme passer du "dessin au crayon" au "téléchargement instantané".
Qualité : Si on prend le temps de faire plusieurs étapes, l'image est magnifique (très nette).
Ordre naturel : Il a enfin réussi à donner aux images le même ordre logique que le texte, ce qui ouvre la porte à l'utilisation de modèles de langage pour créer des images (comme si on écrivait une histoire pour générer un film).

En une phrase : CaTok est comme un chef d'orchestre qui apprend à un ordinateur à dessiner une image brique par brique, dans le bon ordre, en utilisant une "moyenne" intelligente pour aller vite quand on veut, ou lentement quand on veut la perfection.

Each language version is independently generated for its own context, not a direct translation.

1. Problématique

Le paradigme des modèles de langage autoregressifs (AR) a révolutionné le traitement du texte en tokenisant les séquences en une chaîne causale 1D, permettant une prédiction du "prochain token". Cependant, l'extension de cette approche à la vision par ordinateur reste un défi majeur.

Les méthodes actuelles de tokenisation visuelle présentent des limitations fondamentales :

Tokenisation 2D (ex: VQGAN) : Les images sont découpées en patches 2D puis aplaties en séquences 1D (raster ou aléatoire). Cette approche brise la causalité naturelle entre les tokens successifs.
Tokenisation Multi-échelle (ex: VAR) : Bien qu'elles garantissent une causalité via une prédiction de "prochaine échelle", elles s'éloignent du motif standard de prédiction du "prochain token" des LLM.
Auto-encodeurs de diffusion (ex: FlexTok, Semanticist) :
- Les décodeurs "naïfs" conditionnent sur tous les tokens simultanément, supprimant la causalité.
- Les décodeurs de consistance utilisent un "nested dropout" (conditionnement sur les $k$ premiers tokens). Cela introduit un déséquilibre : les tokens précoces sont sur-représentés, ce qui nuit à la qualité de la génération autoregressive.
- L'application de mécanismes de dropout hiérarchique crée souvent des biais qui dégradent les performances de génération.

L'objectif est donc de concevoir un tokeniseur d'image 1D causal qui respecte la structure de prédiction autoregressive tout en permettant une génération rapide (échantillonnage en une étape) et de haute fidélité.

2. Méthodologie : CaTok

Les auteurs proposent CaTok, un auto-encodeur de diffusion basé sur une architecture MeanFlow, conçue pour apprendre des représentations visuelles 1D causales et équilibrées.

A. Architecture

Encodeur : Un Vision Transformer (ViT) causal avec des "registers" (registres) appris. Il extrait des caractéristiques visuelles riches et les compresse en une séquence de tokens 1D. Un masque d'attention causal est appliqué pour garantir que chaque token ne dépend que des tokens précédents.
Décodeur : Un MeanFlow Diffusion Transformer (DiT). Contrairement aux décodeurs classiques qui prédisent la vitesse instantanée, ce décodeur est entraîné pour prédire le champ de vitesse moyen sur un intervalle de temps.

B. Le Cœur de l'Approche : MeanFlow et Sélection d'Intervalle

La contribution centrale réside dans la manière dont le décodeur est conditionné :

Sélection d'Intervalle Temporel : Au lieu d'utiliser tous les tokens ou seulement les premiers, CaTok sélectionne dynamiquement un sous-ensemble de tokens correspondant à un intervalle de temps $[r, t]$ échantillonné aléatoirement.
Objectif MeanFlow : Le décodeur prédit la vitesse moyenne $u$ $u$ sur cet intervalle $[r, t]$ $[r, t]$ en fonction des tokens sélectionnés dans cet intervalle.
- Cela permet de modéliser le champ de vitesse moyen le long du sous-chemin de la diffusion.
- Causalité : En liant les tokens à l'intervalle de temps, la structure causale est préservée naturellement.
- Équilibre : Cette méthode évite le biais vers les tokens précoces (contrairement au nested dropout), assurant que chaque token contribue de manière équitable à la reconstruction.
- Échantillonnage en une étape : La nature de l'objectif MeanFlow permet une génération directe en une seule étape ( $z_0 = \epsilon - u_\theta(\epsilon, 0, 1)$ ) sans perte significative de qualité.

C. Régularisation REPA-A

Pour stabiliser et accélérer l'entraînement, les auteurs proposent REPA-A (Representation Alignment - A).

Contrairement aux méthodes précédentes qui alignent les caractéristiques du décodeur, REPA-A aligne les caractéristiques de l'encodeur avec celles d'un modèle de fondation visuel (VFM, ex: DINOv2).
Cela force l'encodeur à produire des représentations sémantiques de haute qualité, accélérant la convergence et améliorant la capacité des tokens à capturer du contenu visuel discriminatif.

3. Contributions Clés

Architecture Novel : Proposition d'un tokeniseur d'image 1D causal basé sur des auto-encodeurs de diffusion avec un objectif MeanFlow.
Échantillonnage Unifié : Combinaison fluide d'un encodeur causal et d'un décodeur à flux rectifié permettant à la fois un échantillonnage rapide en une étape et un échantillonnage multi-étapes de haute fidélité.
Régularisation REPA-A : Une technique avancée d'alignement de représentations utilisant des modèles de fondation visuels pour stabiliser l'entraînement des auto-encodeurs conditionnels.
Résultats SOTA : Atteinte de performances state-of-the-art sur la reconstruction ImageNet avec un nombre d'époques d'entraînement réduit.

4. Résultats Expérimentaux

Les expériences ont été menées sur ImageNet-1K (256x256).

Reconstruction

CaTok-L-256 atteint des performances exceptionnelles :
- rFID : 0.75 (meilleur résultat parmi les tokeniseurs 1D, comparable aux meilleurs modèles 2D).
- PSNR : 22.53.
- SSIM : 0.674.
Efficacité : CaTok-B-256 atteint des résultats comparables en seulement 80 époques, contre 300+ pour d'autres méthodes (ex: Semanticist, FlexTok).
Flexibilité : Le modèle supporte un échantillonnage en une étape (rapide) tout en permettant une reconstruction de haute qualité via un échantillonnage multi-étapes (25 étapes).

Génération Autoregressive (AR)

En utilisant un modèle AR standard (LlamaGen) entraîné avec une perte de diffusion sur les tokens CaTok :
- gFID : 2.95 (pour 128 tokens), comparable aux approches leaders.
- Les ablations montrent que la sélection de tokens dans l'intervalle $[r, t]$ est cruciale : les méthodes sans causalité ou avec déséquilibre (tokens "First k") obtiennent des gFID nettement pires (ex: 13.54 vs 4.91).

Analyse Ablative

Causalité et Équilibre : La sélection dans $[r, t]$ est supérieure à l'utilisation de tous les tokens (non causal) ou des premiers $k$ tokens (déséquilibré).
Impact de REPA-A : L'ajout de REPA-A stabilise l'entraînement (réduit les pics de perte lors de l'introduction du terme MeanFlow) et améliore la qualité des features de l'encodeur, comme le montre la visualisation PCA.

5. Signification et Impact

Ce travail est significatif car il comble le fossé entre les modèles de langage autoregressifs et les modèles de vision génératifs.

Paradigme Unifié : CaTok démontre qu'il est possible d'appliquer le paradigme "next-token prediction" pur aux images sans sacrifier la qualité de reconstruction ou la vitesse de génération.
Efficacité : La capacité à atteindre des résultats SOTA avec moins d'époques d'entraînement et la possibilité d'une génération en une étape rendent cette approche très attractive pour les applications temps réel.
Causalité Naturelle : En résolvant le problème du déséquilibre inhérent aux méthodes de dropout imbriqué, CaTok permet une modélisation plus juste de la dépendance temporelle dans les images, ouvrant la voie à des modèles de génération visuelle plus scalables et performants, similaires aux LLM.

En résumé, CaTok propose une nouvelle voie pour la tokenisation visuelle, transformant les images en séquences 1D causales équilibrées, compatibles avec l'architecture des grands modèles de langage.