Discrete Variational Autoencoding via Policy Search

Technische Samenvatting: Discrete Variational Autoencoding via Policy Search (DAPS)

1. Het Probleem: De uitdaging van discrete latente ruimtes

In Variational Autoencoders (VAEs) is het gebruik van discrete latente variabelen zeer wenselijk voor toepassingen zoals telecommunicatie, robotica en compressie, omdat ze efficiënt zijn en multimodale verdelingen kunnen modelleren. Echter, discrete variabelen vormen een fundamenteel probleem voor gradiëntgebaseerde optimalisatie: de operaties die nodig zijn om discrete keuzes te maken (zoals argmax of afronding) zijn niet differentieerbaar.

De huidige oplossingen hebben significante nadelen:

• Gumbel-Softmax & Straight-Through Estimators: Deze maken gebruik van benaderingen. Gumbel-Softmax is zeer gevoelig voor de temperatuur-hyperparameter (een trade-off tussen bias en variantie), en de benadering wordt minder nauwkeurig bij grotere bottlenecks.
• Vector Quantization (VQ-VAE/FSQ): Deze vermijden de reparameterisatie-problemen, maar gebruiken vaak de straight-through estimator en hebben een onhandelbare latente verdeling, waardoor het moeilijk is om de entropie (informatiedichtheid) expliciet te reguleren.
• Score-function estimators (REINFORCE): Hoewel ze onbevooroordeeld zijn, lijden ze onder een extreem hoge variantie in de gradiënten, wat training op complexe taken (zoals hoge-resolutie afbeeldingen) bijna onmogelijk maakt.

2. Methodologie: DAPS (Discrete Autoencoding via Policy Search)

De auteurs stellen DAPS voor, een framework dat de training van discrete VAE's benadert als een policy search-probleem uit de Reinforcement Learning (RL) hoek. In plaats van te proberen de gradiënt door de discrete steekproef te "forceren" via benaderingen, optimaliseert DAPS de encoder via weighted maximum likelihood.

De kerncomponenten zijn:

• KL-geregulariseerde Policy Search: Het optimaliseren van de Evidence Lower Bound (ELBO) wordt geformuleerd als een probleem waarbij een niet-parametrische doelverdeling ($q^*$) wordt geconstrueerd. De parametrische encoder ($q_\theta$) wordt vervolgens bijgewerkt om deze doelverdeling te benaderen.
• Geen Reparameterisatie nodig: Omdat de update gebeurt via gewogen maximum likelihood (importance sampling), is er geen backpropagation door de discrete steekproef nodig. Dit maakt het gebruik van krachtige autoregressieve Transformers als encoders mogelijk zonder de instabiliteit van backpropagation through time.
• Variantie-reductie via Advantages: Om de hoge variantie van score-function methoden te beheersen, gebruikt DAPS een "soft-maximum" van de returns als baseline, wat fungeert als een optimistische schatting van de waarde van een staat.
• Automatische Stapgrootte Adaptatie (ESS): Om de stabiliteit te waarborgen, wordt een trust-region parameter ($\eta$) automatisch aangepast met behulp van de Effective Sample Size (ESS). Dit zorgt ervoor dat de stapgrootte tijdens de training stabiel blijft, ongeacht de schaal van de loss.

3. Belangrijkste Bijdragen

1. Nieuw Trainingsframework: Een robuust alternatief voor Gumbel-Softmax en VQ-VAE dat direct de ELBO optimaliseert voor discrete, autoregressieve encoders.
2. Schaalbaarheid: Het vermogen om autoregressieve Transformers te gebruiken voor zeer grote latente sequenties (bijv. 1024 tokens voor ImageNet), wat voorheen computationeel en numeriek onstabiel was.
3. Expliciete Controle: In tegenstelling tot VQ-methoden biedt DAPS directe controle over de entropie en bit-rate via de $\beta$-parameter in de ELBO.
4. Robotica Toepassing: De demonstratie dat de geleerde latente ruimte effectief kan worden gebruikt voor complexe, hoog-dimensionale bewegingscontrole (robotica).

4. Resultaten

De experimenten op diverse datasets laten superieure prestaties zien:

• MNIST & CIFAR-10: DAPS presteert beter dan Gumbel-Softmax en GR-MCK op zowel reconstructiekwaliteit (PSNR) als de ELBO-score.
• ImageNet-256: DAPS schaalt naar zeer hoge resoluties en overtreft traditionele methoden (zoals FSQ en VQ-VAE) aanzienlijk in reconstructiekwaliteit en FID-score, terwijl het een compacte latente ruimte behoudt.
• LAFAN (Robotica): DAPS is in staat om menselijke bewegingen accuraat te reconstrueren en te coderen. In een downstream taak (goal-conditioned robot control) kan een robot complexe bewegingen uitvoeren op basis van de door DAPS geleerde latente codes.
• Latente Benutting: Visualisaties tonen aan dat DAPS de volledige vocabulaire van de codebook veel efficiënter benut dan VQ-VAE of FSQ, wat duidt op een gezondere en rijkere latente representatie.

5. Betekenis en Conclusie

DAPS overbrugt de kloof tussen Variational Inference en Policy Search. De belangrijkste wetenschappelijke betekenis is dat het aantoont dat de problemen van discrete gradiënten niet noodzakelijkerwijs opgelost hoeven te worden door de discretisering te "verzachten" (zoals Gumbel-Softmax doet), maar dat men de optimalisatie kan herformuleren naar een vorm die de discretisering respecteert.

Dit opent de deur naar het gebruik van extreem krachtige, autoregressieve modellen voor discrete representaties, wat essentieel is voor de volgende generatie generatieve AI-modellen die werken met discrete tokens (zoals Large Language Models, maar dan voor visuele of fysieke data).