Omni-Diffusion: Unified Multimodal Understanding and Generation with Masked Discrete Diffusion

Le papier présente Omni-Diffusion, le premier modèle de langage multimodal universel entièrement basé sur un processus de diffusion discrète masquée, qui unifie la compréhension et la génération de texte, de parole et d'images en capturant directement la distribution conjointe de leurs tokens discrets.

L'essentiel

Imaginez un chef cuisinier ultra-polyvalent nommé Omni-Diffusion. Jusqu'à présent, la plupart des chefs intelligents (les modèles d'IA actuels) fonctionnent comme des écrivains qui écrivent un livre mot par mot, de gauche à droite. S'ils font une erreur au début, ils doivent tout effacer et recommencer. C'est lent et rigide.

Omni-Diffusion, lui, fonctionne différemment. Il utilise une technique appelée "Diffusion Discrète Masquée". Voici comment on peut l'expliquer simplement :

1. Le concept de base : Le Puzzle Flou

Imaginez que vous avez une image, un texte ou un enregistrement vocal, mais que tout est caché sous un voile de brouillard (des masques).

• Les autres chefs (Modèles Autogressifs) : Ils regardent le brouillard et devinent le premier mot, puis le deuxième, puis le troisième, un par un.
• Omni-Diffusion : Il regarde tout le puzzle flou d'un coup. Il devine plusieurs pièces manquantes en même temps, les place, puis regarde à nouveau le puzzle pour corriger les erreurs. Il répète ce processus très vite jusqu'à ce que l'image, la voix ou le texte soit parfaitement clair.

C'est comme si vous aviez un tableau blanc rempli de taches de peinture floues, et que vous deviez deviner ce qu'il y a dessous. Au lieu de peindre ligne par ligne, vous ajoutez de la peinture par endroits, regardez le résultat, et ajustez. C'est beaucoup plus rapide et flexible !

2. Le Super-Héros "Tout-en-Un"

La grande innovation d'Omni-Diffusion est qu'il ne fait pas de différence entre les langues, les images et les voix.

• Les autres modèles : Ils ont souvent un cerveau pour le texte, un autre pour les images, et un troisième pour la voix. Ils doivent passer les informations d'un cerveau à l'autre, comme un jeu de "téléphone arabe" où le message se dégrade.
• Omni-Diffusion : Il a un seul cerveau universel. Pour lui, une image, un mot et une note de musique sont tous la même chose : des briques de Lego (des tokens).
  • Vous pouvez lui dire : "Regarde cette photo de chat et raconte-moi une histoire à voix haute."
  • Il comprend la photo, invente l'histoire, et la chante en même temps, car tout est mélangé dans son cerveau unique.

3. Comment il apprend (L'entraînement en 3 étapes)

Pour devenir aussi doué, le chef a suivi un entraînement progressif, comme un athlète :

1. Étape 1 (Texte + Image) : Il apprend d'abord à associer des mots à des images (comme décrire une photo).
2. Étape 2 (Ajout de la Voix) : On lui apprend à associer les mots aux sons de voix.
3. Étape 3 (Le Grand Mix) : On lui donne des exercices complexes où il doit tout faire en même temps, comme répondre à une question posée à voix haute en regardant une image, puis générer une nouvelle image en réponse.

4. Ses super-pouvoirs spéciaux

Pour éviter les erreurs, les chercheurs ont ajouté des astuces intelligentes :

• Le "Masque Atténué" : Parfois, le modèle a tendance à ajouter trop de "rien" (des espaces vides) à la fin de ses réponses. Omni-Diffusion apprend à ne pas trop masquer ces espaces vides pendant l'entraînement, pour qu'il sache exactement quand s'arrêter.
• La "Pénalité de Position" : Pour les images, le modèle avait tendance à répéter les mêmes motifs (comme des rayures partout). Les chercheurs lui ont dit : "Attention, ne devine pas le début et la fin de l'image en même temps, sinon tu vas faire des doublons !". Cela force le modèle à créer des images plus naturelles.
• Le "Pré-remplissage" : Pour la voix, il sait qu'avant de parler, il faut souvent penser à ce qu'on va dire. Il prépare donc un petit texte mental avant de générer la voix, ce qui rend la conversation plus logique.

En résumé

Omni-Diffusion est le premier modèle capable de comprendre et de créer n'importe quoi (texte, image, voix) en utilisant une seule et même méthode magique : deviner et corriger des pièces manquantes dans un puzzle géant, toutes en même temps.

C'est comme passer d'un écrivain qui écrit lettre par lettre à un sculpteur qui voit la statue finale dans le bloc de marbre et enlève simplement ce qui ne va pas, rapidement et avec une grande précision. Cela ouvre la porte à des assistants IA beaucoup plus rapides, plus créatifs et capables de dialoguer avec nous de manière totalement naturelle, comme des humains.

Résumé technique

1. Problématique et Contexte

Les modèles de langage multimodaux (MLLM) actuels reposent majoritairement sur une architecture autoregressive (génération token par token, de gauche à droite). Bien que performants, ces modèles présentent des limitations :

• Génération séquentielle : Ils ne peuvent pas générer plusieurs tokens en parallèle, ce qui limite l'efficacité de la génération.
• Contrôle limité : Il est difficile de contrôler la structure sémantique ou le format de sortie dès le début du processus de génération.
• Architectures fragmentées : La plupart des systèmes unifiés utilisent un LLM pour le texte et des modèles de sortie supplémentaires (décodeurs) pour les autres modalités (images, audio), créant un espace de représentation sémantique non intrinsèquement aligné.

Le papier propose d'explorer une alternative : les modèles de diffusion discrets basés sur le masquage (Masked Discrete Diffusion Models - MDM). Ces modèles offrent la possibilité de décoder en parallèle et de contrôler le processus de génération, mais leur application à un système "any-to-any" (toutes modalités vers toutes modalités) unifié reste une tâche non résolue.

2. Méthodologie : Omni-Diffusion

Omni-Diffusion est le premier modèle de langage multimodal capable de gérer n'importe quelle combinaison d'entrées et de sorties (texte, image, parole) entièrement basé sur un modèle de diffusion discret masqué.

A. Formulation Probabiliste Unifiée

Au lieu d'utiliser un LLM pour le texte et des décodeurs séparés pour les autres modalités, Omni-Diffusion modélise directement la distribution conjointe de tous les tokens discrets multimodaux.

• Tokenisation : Le texte, la parole et les images sont convertis en séquences de tokens discrets.
  • Image : Utilise MAGVIT-v2 (encodeur visuel + quantiseur).
  • Parole : Utilise SenseVoiceSmall pour l'encodage et GLM-4-Voice pour le décodage (quantisation scalaire finie).
  • Texte : Basé sur le vocabulaire du modèle de base.
• Apprentissage : Le modèle prend une séquence de tokens corrompus (où certains tokens sont remplacés par un token spécial [MASK]) et prédit les tokens originaux non masqués. La perte est une entropie croisée calculée uniquement sur les tokens masqués.

B. Architecture du Modèle

• Backbone : Basé sur Dream-7B, un modèle de langage pré-entraîné par diffusion discrète.
• Extension du vocabulaire : Le vocabulaire est étendu pour inclure 16 384 tokens de parole et 8 192 tokens d'image, tout en conservant l'architecture de base inchangée.

C. Stratégies d'Entraînement

Pour stabiliser l'entraînement d'un modèle unifié sur des distributions de données hétérogènes, les auteurs proposent un pipeline d'entraînement progressif en trois étapes :

1. Pré-alignement Visuel-Langage : Alignement du modèle de diffusion pré-entraîné avec la modalité visuelle (tâches de légendage et génération image-texte).
2. Alignement Joint Parole-Vision-Langage : Introduction des données ASR (reconnaissance) et TTS (synthèse) pour aligner la parole avec le texte et l'image.
3. Amélioration de l'Interaction Visuelle Pilotée par la Parole : Affinement sur un nouveau jeu de données SDVI (Speech-Driven Visual Interaction) contenant des questions/réponses orales sur des images et de la génération d'images à partir de la parole.

Techniques spécifiques d'entraînement :

• Masquage atténué en fin de séquence (Attenuated Tail-Pad Masking) : Pour éviter le surapprentissage sur les tokens de remplissage (pad) et permettre des générations de longueurs variables, le taux de masquage des tokens de remplissage est réduit par un facteur d'échelle $\gamma < 1$.

D. Techniques d'Inférence

Pour optimiser la qualité et la cohérence, plusieurs stratégies sont introduites :

• Décodage basé sur l'entropie : Sélection des tokens les plus confiants pour le démasquage itératif.
• Pénalité de position (Position Penalty) : Pour l'image, les logits des tokens en début et fin de séquence sont réduits pour empêcher le modèle de générer simultanément des motifs répétitifs aux extrémités (un problème courant dans la diffusion discrète).
• Pré-remplissage de tokens spéciaux (Special Token Pre-Infilling) : Pour la parole, un token spécial [begin-of-speech] est inséré à 25% de la séquence pour guider le modèle à générer d'abord le texte, puis la parole correspondante, améliorant la cohérence logique.
• Initialisation adaptative de la longueur : La longueur initiale des tokens masqués est ajustée dynamiquement (ex: 3.5x la longueur du texte pour le TTS) pour accélérer l'échantillonnage.

3. Contributions Clés

1. Premier modèle "Any-to-Any" basé sur la diffusion : Omni-Diffusion est le premier système unifié capable de comprendre et de générer du texte, de la parole et des images sans utiliser d'architecture autoregressive, mais via un modèle de diffusion discret masqué.
2. Espace de représentation unifié : En modélisant la distribution conjointe des tokens, le modèle crée un espace sémantique intrinsèquement aligné, éliminant le besoin de décodeurs externes pour les modalités non textuelles.
3. Nouvelles techniques d'entraînement et d'inférence : Développement de stratégies spécifiques (masquage atténué, pénalité de position, pré-remplissage) qui résolvent les défis de la génération multimodale avec la diffusion.
4. Jeu de données SDVI : Création d'un ensemble de données pour l'interaction visuelle pilotée par la parole (questions/réponses orales sur images, génération d'images par la parole).

4. Résultats Expérimentaux

Les évaluations ont été menées sur plusieurs benchmarks standards :

• Tâches de Parole (ASR et TTS) :
  • Sur LibriSpeech et LibriTTS, Omni-Diffusion surpasse le modèle "Any-to-Any" autoregressif (AnyGPT) et montre des performances comparables aux modèles experts en synthèse vocale (CosyVoice), tout en étant un modèle unique.
• Tâches Visuelles (VQA et Image-to-Text) :
  • Sur les benchmarks POPE, MME-Perception et Seed-2-Plus, le modèle atteint des performances comparables aux LLM visuels spécialisés (comme LLaVA ou InstructBLIP), tout en supportant également la génération.
• Génération d'Images (Text-to-Image et Speech-to-Image) :
  • Le modèle obtient des scores CLIP-T et CLIP-I supérieurs aux autres modèles "Any-to-Any" et comparables aux modèles utilisant des décodeurs de diffusion externes pré-entraînés.
  • Il démontre une forte cohérence sémantique entre la parole d'entrée et l'image générée.
• Efficacité d'Échantillonnage :
  • Grâce au décodage parallèle, Omni-Diffusion maintient une haute qualité de génération même avec un nombre très réduit d'étapes de temps (jusqu'à 10 étapes pour l'image et 0.25L pour la parole), surpassant largement l'efficacité des modèles autoregressifs.
• Inpainting :
  • Le modèle peut effectuer de l'inpainting (remplissage de zones manquantes) sans fine-tuning supplémentaire, simplement en masquant les régions à reconstruire.

5. Signification et Impact

Ce travail démontre que les modèles de diffusion discrets sont une alternative viable, et potentiellement supérieure, aux architectures autoregressives pour les fondations de l'IA multimodale.

• Efficacité : La capacité de génération parallèle offre un avantage majeur en termes de vitesse d'inférence.
• Flexibilité : La nature "any-to-any" unifiée simplifie l'architecture des systèmes multimodaux futurs, remplaçant des pipelines complexes par un seul modèle.
• Contrôle : Les mécanismes de diffusion permettent un contrôle plus fin sur la structure et le format de la sortie.

Omni-Diffusion ouvre la voie à une nouvelle génération de modèles de base multimodaux capables d'interagir de manière fluide et naturelle à travers le texte, la parole et l'image, tout en surmontant les limitations de génération séquentielle.