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

Omni-Diffusion introduceert het eerste multimodale taalmodel dat volledig is gebaseerd op een gemaskerde discrete diffusie-architectuur en zo het begrijpen en genereren van tekst, spraak en beelden in één unified model verenigt.

De kern

Stel je voor dat je een super-intelligente robot bouwt die niet alleen kan lezen en schrijven, maar ook kan kijken, luisteren en praten. Tot nu toe waren de slimste robots die we hadden, eigenlijk als een zeer snelle schrijver: ze schreven woord voor woord, letter voor letter, van links naar rechts. Dit noemen we "autoregressief". Het werkt goed, maar het is alsof je een schilderij maakt door één penseelstreek tegelijk te doen, terwijl je de rest van het doek nog niet ziet.

De onderzoekers van Omni-Diffusion hebben een heel andere aanpak bedacht. Ze zeggen: "Waarom niet alles tegelijk doen?"

Hier is een simpele uitleg van wat ze hebben gedaan, met behulp van een paar creatieve vergelijkingen:

1. De "Magische Sieradenkist" (In plaats van een schrijfmachine)

Stel je voor dat je een grote doos hebt vol met losse letters, plaatjes en geluiden.

• De oude manier (Autoregressief): De robot pakt één letter, schrijft die op, pakt dan de volgende, en zo gaat het door. Het is een lange, lineaire lijn.
• De nieuwe manier (Omni-Diffusion): De robot begint met een doos die volledig vol zit met vraagtekens (in de vakjargon: "mask tokens"). Het ziet eruit als een raadsel. De robot kijkt naar de vraagtekens en denkt: "Aha, hier hoort waarschijnlijk een 'k' te staan, en daar een plaatje van een kat."
• Het proces: In plaats van één voor één te schrijven, kijkt de robot naar alle vraagtekens tegelijk. Hij vult er een paar in, kijkt weer, en vult er nog een paar in. Hij herhaalt dit proces totdat alle vraagtekens zijn vervangen door het juiste antwoord.

Dit is als het maken van een puzzel. Je begint met een lege puzzel en vult stukjes in die het beste passen bij de rest, totdat het hele plaatje helder is. Dit gaat veel sneller omdat je niet hoeft te wachten tot het vorige stukje klaar is om aan het volgende te beginnen.

2. De "Universele Vertaler"

Tot nu toe hadden we vaak aparte robots voor verschillende taken: één die goed kon tekenen, één die goed kon praten, en één die goed kon lezen. Als je de teken-robot iets wilde laten praten, moest je eerst een tussenstap maken.

Omni-Diffusion is als een universele vertaler die alles in één taal spreekt: "Tokens".

• Of je nu een foto, een stemopname of een tekst invoert, de robot zet alles om in dezelfde soort bouwstenen.
• Omdat alles in dezelfde taal is, begrijpt de robot dat een "gelach" in een stemopname en een "lachend gezicht" op een foto eigenlijk hetzelfde gevoel oproepen. Ze zitten in hetzelfde "gevoelsgebied" in zijn hersenen.
• Hierdoor kan hij moeiteloos schakelen: "Ik hoorde een stem die vraagt om een tekening van een kat" -> Bamm -> Hij maakt de tekening. Of: "Ik zie een foto van een hond" -> Bamm -> Hij vertelt je er een verhaal over.

3. De "Slimme Regels" om fouten te voorkomen

Omdat de robot alles tegelijk probeert in te vullen, kan hij soms in de war raken. De onderzoekers hebben daarom een paar slimme trucs bedacht:

• De "Geen Herhaling"-regel (Voor plaatjes): Soms tekent de robot twee keer dezelfde boom of twee keer dezelfde oogbol, omdat hij aan beide kanten van het plaatje tegelijk begint. De onderzoekers hebben een regel toegevoegd: "Hé, begin niet te snel aan de randen, wacht even met de uiteinden." Dit zorgt voor mooiere, meer natuurlijke plaatjes.
• De "Vooraf invullen"-truc (Voor spreken): Als de robot een tekst moet omzetten naar spraak, helpt hij zichzelf door eerst een speciaal teken te zetten dat zegt: "Hier komt eerst de tekst, en daarna de stem." Dit zorgt ervoor dat de stemlogica klopt en hij niet halverwege de zin begint te stotteren.
• De "Niet te lang"-regel: Soms maakt de robot te veel "leegte" (vulwoorden) aan het einde van een zin. De onderzoekers hebben de robot geleerd om die lege plekken minder vaak te maskeren, zodat hij zich meer focust op de echte inhoud.

Waarom is dit belangrijk?

Deze nieuwe robot (Omni-Diffusion) is niet alleen slimmer in het begrijpen van alles tegelijk, maar ook sneller. Omdat hij niet woord voor woord hoeft te wachten, kan hij in één keer een hele zin of een heel plaatje "ontdekken".

Het is alsof je van een robot die een briefje voor je schrijft, bent gegaan naar een robot die een compleet gesprek voert, een schilderij maakt en een liedje zingt, allemaal tegelijk, zonder te struikelen. Dit is een enorme stap voorwaarts voor de toekomst van kunstmatige intelligentie die echt "menselijk" aanvoelt in zijn interactie.

Technische samenvatting

Probleemstelling

Recente vooruitgangen in multimodale grote taalmodellen (MLLM's) hebben voornamelijk vertrouwd op autoregressieve architecturen (waarbij tokens sequentieel worden gegenereerd). Hoewel deze modellen indrukwekkende prestaties leveren, laten ze ruimte voor alternatieve probabilistische modellering. Autoregressieve modellen hebben beperkingen, zoals het gebrek aan parallelle decoding (wat de efficiëntie belemmert) en minder controle over de semantische structuur en output-indeling tijdens het generatieproces.

Daarnaast zijn bestaande multimodale systemen vaak gefragmenteerd: ze gebruiken een taalmodel voor tekst en koppelen daar extra modellen aan voor andere modaliteiten (zoals afbeeldingen of spraak). Dit leidt tot een gebrek aan een intrinsiek uitgelijnde semantische representatieruimte tussen de verschillende modaliteiten. Er is een behoefte aan een unificatie van begrip en generatie voor tekst, spraak en afbeeldingen binnen één enkel architectonisch raamwerk dat niet afhankelijk is van autoregressie.

Methodologie: Omni-Diffusion

Omni-Diffusion is het eerste "any-to-any" multimodale taalmodel dat volledig is gebouwd op een mask-based discrete diffusion model. In plaats van tokens sequentieel te voorspellen, leert het model de gezamenlijke verdeling van discrete multimodale tokens door maskering en reconstructie.

1. Architectuur en Tokenisatie

• Unificatie: Het model behandelt tekst, spraak en afbeeldingen als een enkele reeks discrete tokens.
• Tokenizers:
  • Afbeeldingen: Gebruik van MAGVIT-v2 om afbeeldingen te comprimeren en te kwantiseren naar 8192 discrete tokens.
  • Spraak: Gebruik van SenseVoiceSmall voor encoding en GLM-4-Voice decoder voor decoding, met een codebook van 16384 tokens.
  • Tekst: Bestaande tokenizers voor tekst.
• Backbone: Het model is gebaseerd op Dream-7B, een voorgeïmplementeerde discrete diffusion taalmodel. Het vocabulaire is uitgebreid om de nieuwe spraak- en beeldtokens te accommoderen, maar de kernarchitectuur blijft behouden.

2. Training Strategie

Om een stabiel en effectief trainingsproces te garanderen, wordt een driefasen progressieve trainingspipeline gebruikt:

1. Fase 1 (Visueel-Taal Vooruitlijning): Optimalisatie op tekst-naar-afbeelding en afbeelding-captioning taken om de visuele modality uit te lijnen met de semantische ruimte van het taalmodel.
2. Fase 2 (Spraak-Visie-Taal Gezamenlijke Uitlijning): Introductie van Automatic Speech Recognition (ASR) en Text-to-Speech (TTS) data om de uitlijning tussen tekst en spraak te verbeteren, terwijl visuele data behouden blijft.
3. Fase 3 (Verbetering van Spraak-gedreven Visuele Interactie): Fijnafstelling op een nieuw samengesteld dataset (SDVI), dat bestaat uit mondelinge visuele vraag-antwoordparen en spraak-naar-afbeelding generatie. Dit vereist gezamenlijke verwerking van spraak en visuele data.

Specifieke Trainingstechnieken:

• Attenuated Tail-Pad Masking: Om variabele lengte generatie mogelijk te maken, worden padding-tokens toegevoegd. Om overfitting op deze padding-tokens te voorkomen, wordt de maskeringsratio voor deze tokens verlaagd (verzwakt) met een factor $\gamma < 1$.
• SDVI Dataset: Een dataset van >30.000 samples voor mondelinge visuele interactie, gegenereerd met Cosyvoice2, waarbij vragen en antwoorden zijn omgezet naar spraak met gevarieerde stemmen.

3. Inference (Afleiding) Strategieën

Om de kwaliteit en efficiëntie te maximaliseren, worden specifieke technieken toegepast:

• Entropie-gebaseerde Decoding: Tokens worden geselecteerd op basis van hun entropie (onzekerheid), met integratie van herhalingsstraffen en classifier-free guidance.
• Positie Penalty (voor Afbeeldingen): Om repetitieve patronen in gegenereerde afbeeldingen te voorkomen (waarbij het model vaak van de randen naar het midden decodeert), wordt de logit van de laatste $N$ tokens verlaagd. Dit dwingt een meer gebalanceerde decodeervolgorde af zonder de flexibiliteit van autoregressie te vereisen.
• Special Token Pre-Infilling (voor Spraak): Bij spraakgeneratie wordt een speciaal token [begin-of-speech] op een specifiek punt in de maskerreeks ingevoegd. Dit leidt het model om eerst tekst en vervolgens spraak te genereren, wat de logica en coherentie verbetert.
• Adaptieve Token Lengte: Voor ASR en TTS wordt de initiële lengte van de maskerreeks dynamisch bepaald op basis van de lengte van de tekst (bijv. 3.5x de tekstlengte voor TTS), wat de sampling versnelt.

Belangrijkste Resultaten

Omni-Diffusion is uitgebreid geëvalueerd op diverse benchmarks en presteert gelijkwaardig aan of beter dan bestaande autoregressieve systemen:

• Spraaktaken (ASR & TTS): Op de LibriSpeech en LibriTTS benchmarks behaalt Omni-Diffusion een lagere woordfoutenratio (WER) dan het "any-to-any" model AnyGPT en presteert het vergelijkbaar met gespecialiseerde TTS-modellen (zoals CosyVoice), terwijl het significant beter presteert dan spraak-specifieke LLM's.
• Visuele Taken (VQA & Text-to-Image):
  • VQA: Presteert op niveau met gespecialiseerde visuele LLM's (zoals LLaVA en InstructBLIP) op benchmarks zoals POPE en MME-Perception.
  • Text-to-Image: Bereikt een betere tekst-afbeelding uitlijning dan andere "any-to-any" modellen en visuele kwaliteit vergelijkbaar met methoden die externe diffusion-modellen gebruiken.
• Cross-Modal Uitlijning: Het model toont sterke uitlijning bij taken die meerdere modaliteiten combineren, zoals spraak-naar-afbeelding generatie, waarbij de kwaliteit vergelijkbaar is met tekst-naar-afbeelding.
• Efficiëntie: Dankzij de parallelle decoding van diffusion-modellen behoudt Omni-Diffusion hoge kwaliteit zelfs bij zeer weinig tijdstappen (bijv. 10 stappen voor afbeeldingen), wat een groot voordeel is ten opzichte van autoregressieve modellen die veel stappen nodig hebben.

Bijdragen

1. Eerste Any-to-Any Diffusion Model: Introductie van Omni-Diffusion, het eerste multimodale model dat volledig gebaseerd is op een mask-based discrete diffusion architectuur voor zowel begrip als generatie.
2. Unificatie van Representaties: Door de gezamenlijke verdeling van discrete tokens direct te modelleren, creëert het model een intrinsiek uitgelijnde semantische ruimte voor tekst, spraak en beeld, zonder noodzaak voor extra output-modellen.
3. Gespecialiseerde Technieken: Ontwikkeling van trainings- en inferencetechnieken specifiek voor discrete diffusion, zoals attenuated tail-pad masking, positie penalties voor beeldkwaliteit en pre-infilling voor spraakcoherentie.
4. Empirisch Bewijs: Uitgebreide experimenten tonen aan dat diffusion-modellen een levensvatbaar en potentieel superieur alternatief zijn voor autoregressieve modellen in de volgende generatie multimodale foundation modellen.

Significantie

Dit paper markeert een paradigmaverschuiving in multimodale AI. Het demonstreert dat discrete diffusion-modellen niet alleen geschikt zijn voor generatie, maar ook voor eenheid van begrip en generatie over diverse modaliteiten. De resultaten suggereren dat diffusion-architecturen, dankzij hun parallelle aard en flexibiliteit in het sturen van de generatie, de basis kunnen vormen voor de volgende generatie efficiëntere en krachtigere multimodale systemen. Het opent de deur voor "any-to-any" interacties waarbij gebruikers kunnen schakelen tussen spraak, tekst en beeld in één vloeiend gesprek, zonder de beperkingen van sequentiële generatie.