StableVQ: Linee guida pratiche per addestrare tokenizer a quantizzazione vettoriale in modo stabile

Abstract degli autori

La quantizzazione vettoriale (VQ) è fondamentale per i tokenizer visivi discreti alla base dei moderni modelli autoregressivi e di generazione di immagini con mascheramento. Sebbene i recenti metodi di codebook a proiezione condivisa abbiano migliorato notevolmente l’utilizzo del codebook, la stabilità dell’addestramento resta una sfida cruciale e poco esplorata. Sosteniamo che la causa principale risieda nell’intreccio tra l’addestramento del codificatore--decodificatore e quello del codebook: poiché nessuno dei due moduli può svolgere in modo affidabile la propria funzione in isolamento, il sistema riesce a funzionare solo quando i due sottosistemi finiscono per collaborare, una condizione fragile che viene meno proprio quando l’addestramento è sottoposto al massimo stress. Proponiamo StableVQ, che riesamina il corretto obiettivo di apprendimento di ciascun modulo e risolve i problemi che emergono quando ognuno viene addestrato a svolgere autonomamente il proprio ruolo. Nello specifico, (1) Dynamic STE corregge l’instabilità nell’obiettivo di apprendimento del codificatore, consentendogli di ottimizzare in modo robusto lo spazio di ricostruzione con una regolarizzazione discreta, anche quando l’utilizzo del codebook è basso. (2) Region VQ Loss riformula l’obiettivo di apprendimento del codebook, così che possa garantire autonomamente il pieno tracciamento della distribuzione degli output del codificatore, senza fare affidamento sulle oscillazioni del codificatore per innescarne l’attivazione. (3) Decoupled Schedule riconosce che le diverse funzioni del codificatore--decodificatore e del codebook richiedono dinamiche di ottimizzazione distinte e assegna a ciascuno una pianificazione indipendente del tasso di apprendimento, per garantire un comportamento robusto a livello di sistema. Basato su codebook a proiezione condivisa, StableVQ è leggero e non introduce parametri apprendibili. Gli esperimenti su ImageNet dimostrano miglioramenti costanti nella stabilità dell’addestramento, nell’utilizzo del codebook e nella qualità della ricostruzione, con codebook di dimensioni e impostazioni di inizializzazione diverse.