LLM et smartphones : l'exemple MobileDiffusion

Générer une image de 512 x 512 pixels en 0,2 seconde sur un iPhone 15 ? Quatre chercheurs de chez Google en ont fait leur « produit d’appel » pour MobileDiffusion.

Ce modèle expérimental est taillé pour une exécution sur les appareils mobiles. Il a, pour cela, fait l’objet de multiples optimisations sur la base de l’architecture UNet.

Certaines de ces optimisations s’inspirent de travaux préalables. Par exemple, pour la mise à l’échelle des blocs dits transformeurs. Ceux-ci sont, dans les grandes lignes, responsables de la compréhension du texte qui permet de générer les images. Ils incluent une couche d’autoattention particulièrement coûteuse en ressources de calcul.

Les travaux préalables en question introduisent une architecture UViT. Laquelle réduit les « coûts d’attention » en privilégiant le placement de ces blocs sur le segment basse résolution de UNet.
Les chercheurs ont exploité cette approche pour relocaliser des blocs depuis des segments haute résolution. Pour maintenir la cohérence du nombre de paramètres, ils ont réduit la taille des canaux aux plus basses résolutions. Une méthode qui dégrade moins les performances que de réduire la taille même des blocs transformeurs. À nombre de paramètres comparable, le coût d’inférence diminue de 26 % sans baisse de qualité.

Autre technique utilisée : sur les hautes résolutions, découpler la couche d’autoattention et la couche de cross-attention pour ne conserver que cette dernière. Elle est effectivement bien moins gourmande, ne dépendant pas d’une relation quadratique entre résolution et longueur de séquence. Gain annoncé : 15 % en performance, sans perte de qualité.

Au final, en considérant ces deux techniques, les chercheurs ont :

– Supprimé intégralement les blocs transformeurs à la plus haute résolution (64 x 64)
– Éliminé la couche d’autoattention dans les blocs transformeurs intermédiaires (32 x 32 et 16 x 16 « extérieur »)
– Conservé intégralement les blocs transformeurs au cœur de UNet

L’approche GAN appliquée à MobileDiffusion

Toujours au niveau des blocs transformeurs, la mise en place d’un partage des projections de clés et de valeurs (tous deux dérivés du même input) a permis de réduire d’environ 5 % le nombre de paramètres. On en a supprimé 10 % de plus en réduisant le ratio d’expansion – de 8 à 6 – pour les couches à propagation avant.

Les blocs de convolution, qui servent notamment à extraire des caractéristiques, ont aussi fait l’objet d’optimisations. D’une part, une « modularisation » à presque tous les niveaux, à renfort de couches « légères » (résultat : -10 % de paramètres). De l’autre, la réduction du nombre de blocs résiduels (+19 % de performance en inférence, -15 % de paramètres).

Deux plafonds arbitraires ont conditionné le développement de MobileDiffusion : 400 millions de paramètres et 200 Gflops. En complément aux travaux sur l’architecture, les chercheurs ont optimisé le finetuning. En particulier à travers la distillation progressive et l’implémentation de la méthode « hybride » UFOGen, qui intègre une dimension d’entraînement antagoniste.

MobileDiffusion – ainsi qu’une variante Lite – a été comparé à trois autres modèles de diffusion à espace latent. Son VAE (autoencodeur variationnel, qui compresse les images dans cet espace) a été adapté pour encoder sur trois canaux et non quatre.

Les exemples ci-dessous donnent une idée des capacités de MobileDiffusion. En matière de poids d’image, par rapport à Stable Diffusion 1.5, on atteint un taux de compression de 46,4 % (29,9 % pour la version Lite).