Augmenter la taille du codebook de VQGAN à 100 000 avec un taux d’utilisation de 99 %

Dans le domaine de la quantification d’images, illustrée par VQGAN, le processus consiste à encoder les images en des jetons discrets extraits d’un codebook de taille prédéfinie. Les avancées récentes, notamment avec LLAMA 3, montrent que l’agrandissement significatif du codebook améliore notablement les performances du modèle. Toutefois, VQGAN et ses dérivés, tels que VQGAN-FC (Codes factorisés) et VQGAN-EMA, continuent de faire face à des défis liés à l’expansion de la taille du codebook et à l’amélioration de son utilisation. Par exemple, VQGAN-FC est limité à l’apprentissage d’un codebook dont la taille maximale est de 16 384, tout en maintenant un taux d’utilisation généralement inférieur à 12 % sur ImageNet. Dans ce travail, nous proposons un nouveau modèle de quantification d’images nommé VQGAN-LC (Large Codebook), qui étend la taille du codebook à 100 000 éléments, atteignant un taux d’utilisation dépassant 99 %. Contrairement aux méthodes précédentes qui optimisent chaque entrée du codebook individuellement, notre approche part d’un codebook initialement constitué de 100 000 caractéristiques extraites par un encodeur vision pré-entraîné. L’optimisation se concentre ensuite sur l’entraînement d’un projecteur qui aligne l’ensemble du codebook avec les distributions de caractéristiques de l’encodeur dans le cadre de VQGAN-LC. Nous démontrons la supériorité de notre modèle par rapport à ses prédécesseurs sur diverses tâches, notamment la reconstruction d’images, la classification d’images, la génération d’images auto-régressive à l’aide de GPT, ainsi que la création d’images par des modèles génératifs basés sur la diffusion ou les flux. Le code source et les modèles sont disponibles à l’adresse suivante : https://github.com/zh460045050/VQGAN-LC.