APEX compresse les latents de diffusion d'un facteur 128 et surpasse un modèle à 256 tokens sur ImageNet-512

Des chercheurs présentent sur arXiv la méthode APEX, qui s'attaque au principal problème de la diffusion à bas coût : plus un latent est compressé, plus il est difficile à générer. Sur ImageNet-512, APEX surpasse les méthodes existantes pour chaque nombre de canaux testé, aux taux de compression 64x et 128x. À coût d'entraînement égal et avec du guided sampling, il devance également un modèle à 256 tokens sur tous les budgets testés, et l'écart se creuse à mesure que le calcul diminue. Le code est publié sur la page du projet.

Le dilemme à l'origine du projet

Le coût de la diffusion latente dépend directement du nombre de tokens traités par le débruiteur. Moins de tokens, c'est un entraînement comme une inférence plus rapides et moins chers. Mais l'information contenue dans l'image ne disparaît pas. Si les tokens sont peu nombreux, il faut la loger dans les canaux de chacun d'eux.

C'est là que surgit la contradiction que les auteurs appellent reconstruction-generation dilemma. Davantage de canaux permet à l'autoencodeur de restituer l'image plus fidèlement. Mais un latent aussi dense est plus difficile à modéliser pour le modèle de diffusion. Le gain en reconstruction est annulé par la perte en génération.

Ce que propose APEX

Le nom se déploie en Asynchronous Prefix denoising for EXtreme compression. La solution repose sur deux composants qui fonctionnent ensemble.

Le premier est un espace latent hiérarchique. L'autoencodeur ordonne les canaux par préfixes : d'abord ceux qui portent la structure de l'image, puis ceux qui gèrent les détails. Le début du vecteur de canaux décrit l'image d'ensemble, la suite l'affine.

Le second est un débruitage asynchrone par groupes de canaux. Un seul débruiteur commun génère les différents groupes à des instants locaux différents. Le nombre de tokens n'augmente pas, et l'architecture ne se complexifie pas avec des modèles distincts pour chaque groupe. Selon les auteurs, ce schéma permet de générer la structure et les détails de façon différenciée, au lieu de traiter tout le latent comme une seule tâche.

Les chiffres de l'article

• Benchmark : ImageNet-512.
• Taux de compression : 64x et 128x.
• Comparaison avec les méthodes existantes : APEX fait mieux pour chaque nombre de canaux testé.
• Comparaison avec le modèle à 256 tokens : APEX fait mieux sur chaque budget testé, à coût d'entraînement et guided sampling égaux.
• Dynamique : plus le budget de calcul est bas, plus le gain est important.
• Tokens supplémentaires : non requis.

Le point sur le budget importe plus que les autres. Si l'écart se creuse précisément là où le calcul est limité, la méthode intéresse en priorité ceux pour qui l'économie est la première exigence.

Que faire de cela en pratique

Pour les équipes qui construisent leurs propres modèles génératifs d'images, APEX fournit une bonne raison de réévaluer le compromis entre nombre de tokens et qualité. Le code est ouvert : il est donc possible de reproduire le résultat sur ImageNet-512 sans avoir à rétroconcevoir l'article.

Il faut aussi rester prudent. L'abstract ne revendique de résultats que sur ImageNet-512. On n'y voit pas si l'approche fonctionne en text-to-image sur de grands jeux de données, avec de la vidéo ou avec d'autres architectures de débruiteur. Le coût de l'autoencodeur à hiérarchie de canaux n'est pas non plus indiqué : combien coûte son entraînement et dans quelle mesure il est sensible aux réglages.

Il vaut la peine de lire le texte complet pour répondre à trois questions. Comment le calendrier des instants asynchrones est-il construit, exactement ? Quel gain réel en temps d'inférence apporte la compression 128x ? L'avantage se maintient-il lorsque le modèle monte en échelle ? Ces réponses détermineront si APEX deviendra un outil de travail ou restera un travail académique soigné.