Technische Samenvatting: AURORA-LM

Probleemstelling

Taalgeneratie blijft een buitenbeentje in de moderne generatieve modellering. Terwijl afbeeldingen, video en audio steeds vaker worden gemodelleerd in continue latente ruimtes met behulp van diffusieprocessen, vertrouwt tekstgeneratie nog steeds grotendeels op discrete tokens. Bestaande pogingen om continue diffusie toe te passen op taal kampen met een fundamenteel ontwerpcompromis:

1. Geërfde Representaties: Methoden zoals Diffusion-LM of TEncDM erven embeddingruimtes van token-embeddings of vooraf getrainde encoders. Deze ruimtes zijn niet expliciet ontworpen voor de gezamenlijke vereisten van nauwkeurige tekstreconstructie en effectieve continue generatie.
2. Compressie-afwegingen: Op autoencoders gebaseerde benaderingen (bijv. LD4LG, COSMOS) mappen tekst naar continue latenten, maar comprimeren deze representaties vaak om de distributie van de diffusie gemakkelijker leerbaar te maken. Deze compressie offert het niveau van precisie bij tokens op, wat een knelpunt creëert waarbij het verbeteren van de generatieve hanteerbaarheid ten koste gaat van de decodeernauwkeurigheid, en vice versa.

De kernuitdaging bestaat uit het construeren van een continue tekstrepresentatie die een getrouwe, nauwkeurige decodering ondersteunt, terwijl er een generatief model wordt ontworpen dat in staat is de distributie ervan te leren zonder dat dit ten koste gaat van die representatieve capaciteit.

Methodologie: AURORA-LM Raamwerk

AURORA-LM (Autoencoding Unified Representation for Continuous-Latent Diffusion Language Modeling) adresseert deze uitdagingen door de constructie van de representatie los te koppelen van de distributiemodellering in twee afzonderlijke stadia.

1. Constructie van Continue Tekstlatenten (De Autoencoder)

In plaats van representaties over te nemen, construeert AURORA-LM expliciet een hoogwaardige, decodeerbare tekstlatent met behulp van een Query-gebaseerde Encoder-Decoder:

• Query-gebaseerde Encoder: Organiseert discrete tekst in een hoogwaardige, prefix-uitgelijnde latent sequentie. Het gebruikt $N$ latente queries om token-prefixes van toenemende lengte te aggregeren. Naarmate de query-index $i$ toeneemt, let de encoder op progressief langere token-prefixes ($w_{1:\lceil iL/N \rceil}$), waardoor een causale orde ontstaat waarbij latente posities overeenkomen met tekstuele prefixes.
• Query-gebaseerde Decoder: Mapt latente prefixes terug naar token logit's. Cruciaal is dat de decoder beperkt is tot het letten aan slechts de corresponderende latente prefix, wat zorgt voor een strikte causale correspondentie tussen de latente en de token sequenties.
• Hoogwaardig Ontwerp: In tegenstelling tot eerdere werken die latenten comprimeren voor meer eenvoudige generatie, behoudt AURORA-LM een hoge kanaalbreedte ($D$) en minimale sequentiecompressie ($c \approx 1$). Dit behoudt de informatie die nodig is voor exacte tokenherstel, waarbij de latente ruimte wordt behandeld als een getrouwe interface naar tekst in plaats van een vereenvoudigde prior.
• Training: De autoencoder wordt getraind via token-niveau cross-entropie en wordt vervolgens bevroren. Dit fixeert de doelwit-latente distributie, zodat het generatieve model onafhankelijk van de representatieleer kan worden bestudeerd.

2. Block-Causale Latente Diffusie (Het Generatieve Model)

Een Block-Causale Diffusie Transformer leert de distributie van de bevroren latente sequenties middels Flow Matching:

• Block-Causale Factorisatie: Om een balans te vinden tussen causale generatie en parallelle denoising, wordt de latente sequentie in blokken verdeeld. Het model genereert blokken van links naar rechts (causaal conditioneel op vorige blokken), maar voert denoising uit op alle posities binnen een blok in parallel.
• Full-Width Latente Voorspelling: Het model voorspelt de volledige breedte van de schone latente waarde ($x_0$), in plaats van een gecomprimeerde versie.
• Noisy-Input Knelpunt: Om de moeilijkheid van het leren van een hoogdimensionale distributie vanuit ruizige inputs aan te pakken, past het model alleen een laag-rang projectie toe op het ruizige inputpad ($W_{down} \in \mathbb{R}^{D_b \times D}$), terwijl de outputkop de volledige breedte van de schone target voorspelt. Dit stelt het model in staat om gecorrumpeerde inputs efficiënt te verwerken zonder de capaciteit gericht op de decoder te verminderen.
• Ruiskalibratie: De ruis-niveau distributie wordt gekalibreerd aan de latente breedte. Empirische resultaten tonen aan dat bredere latente representaties profiteren van noise schedules die meer nadruk leggen op inputs met een hoge ruis.
• Zelf-Trajectorie Consistentie: Om de kloof te overbruggen tussen trainen op onafhankelijk gesamplede ruizige toestanden en inferentie via iteratieve denoising, introduceert het model een consistentieverlies. Dit stemt de voorspellingen van zuivere latenten bij naburige staten langs de denoising traject samen, wat het generatieproces stabiliseert en kwalitatief hoogwaardige resultaten mogelijk maakt met minder sampling stappen.

3. Inferentie

Tijdens de inferentie worden latente blokken sequentieel van links naar rechts gegenereerd. De gegenereerde gestandaardiseerde latenten worden geïnverteerd naar de oorspronkelijke encoder-ruimte en door de bevroren decoder geleid om de uiteindelijke tokensequentie te herstellen. Het raamwerk ondersteunt zowel onconditionele generatie als prompt-geconditioneerde generatie (met gebruik van classifier-free guidance of self-conditioning guidance).

Belangrijkste Bijdragen

1. Verenigd Raamwerk: AURORA-LM scheidt het leren van de tekstrepresentatie van het modelleren van de generatieve distributie, door expliciet een continue representatie te construeren die het diffusiemodel moet genereren, in plaats van een bestaande vorm over te nemen van token-embeddings.
2. Interface met Hoge Capaciteit: Het behoudt de volledige kanaalcapaciteit van de latente representatie voor nauwkeurige tokendecodering, waarbij de diffusieprior zich aanpast om de resulterende complexe distributie te modelleren in plaats van de representatie te simplificeren.
3. Block-Causale Flow Matching: Introduceert een block-causale denoiser getraind met flow matching, die blokken van links naar rechts genereert terwijl binnen de blokken paralleel denoising vindt.
4. Architecturale Innovaties: De combinatie van een laag-rang noisy-input pad, ruis-niveau kalibratie naar latente breedte en zelf-trajectorie consistentie stelt het model in staat om effectief volledige brede latente distributies te leren.

Experimentele Resultaten

Experimenten werden uitgevoerd op Ascend NPU's.

• OpenWebText (Onconditionele Generatie): AURORA-LM-S (130M parameters) bereikte de laagste Generation Perplexity (Gen-PPL: 23.56) en de hoogste MAUVE (0.890) onder de geëvalueerde autoregressieve, discrete-diffusie en continue-embedding baselines.
• XSum (Conditionele Samenvatting): AURORA-LM-S presteerde beter dan alle baselines (waaronder ELF-B, AR, en MDLM) op ROUGE-1, ROUGE-2 en ROUGE-L metrieken.
• Opschaling: AURORA-LM-L (1.01B parameters, ~1.500 EFLOPs rekenkracht) vertoonde consistente prestatieverbeteringen en overtrof een groter publiek uitgebracht latent-diffusiemodel (Cola-DLM, ~1.8B parameters) onder een gelijkwaardig evaluatieprotocol.
• Ablatiestudies:
  • Latente Breedte: Bredere latenten ($D=1024$) boden betere robuustheid tegen corruptie en hogere generatieve prestaties.
  • Input Bottleneck: Een gemiddelde bottleneck breedte ($D_b=128$) voor ruizige inputs was optimaal, waardoor volledige breedte predictie mogelijk was zonder corrupte inputs op volledige dimensionaliteit te hoeven verwerken.
  • Trajectorie Consistentie: Zelf-trajecterie consistentie verbeterde de kwaliteit van weinig-stap samples significant.

Betekenis en Beweringen

Het artikel beweert dat continue taalgeneratie effectief de brug kan slaan tussen diffusie-gebaseerde generatieve modellering en discrete tokendecodering. De auteurs stellen dat het heersende ontwerpcompromis — het vereenvoudigen van representaties om generatieve modellen tegemoet te komen — onnodig is. Door in plaats daarvan een hoogwaardige, decodeerbare tekstlatent te bewaren en het diffusiemodel zo te ontwerpen dat het de distributie direct leert (via specifieke architecturale keuzes zoals block-causale factorisatie en trajectorie consistentie), is het mogelijk om state-of-the-art prestaties te leveren in zowel onconditionele als conditionele tekstgeneratietaken. De resultaten suggereren dat continue ruimtes kunnen dienen als een algemene interface voor taal, waarbij ze semantische trouw en syntactische regelmaat behouden terwijl zij de voordelen van diffusiemodellering benutten.