Cette Découverte Pourrait Diviser la VRAM par 6 (DiffusionBlocks)

Diffusion Blocks : une méthode japonaise pour réduire la mémoire d’entraînement des IA

Le problème de mémoire dans l’entraînement des grands modèles

milliers de machines connectées entre elles. Mais une équipe de chercheurs japonais vient de présenter une méthode qui pourrait réduire la mémoire nécessaire par 3, 4, voire davantage.

• L’introduction de la vidéo rappelle que l’entraînement des grands modèles d’intelligence artificielle repose sur des milliers de GPU interconnectés. Cette architecture entraîne des coûts matériels et énergétiques considérables, car chaque GPU doit constamment échanger des informations avec les autres. C’est dans ce contexte que le laboratoire japonais Sakana AI propose Diffusion Blocks, une approche consistant à découper le modèle en blocs et à entraîner chaque bloc indépendamment. L’idée paraît presque trop simple sur le papier, mais les premiers résultats indiquent que les performances restent très proches de l’entraînement classique, parfois même meilleures.
• Le problème central est la mémoire. Lors d’une utilisation normale, un modèle comme ChatGPT fait une passe avant : le texte traverse toutes les couches du réseau et le modèle génère une réponse. Pendant l’entraînement, le processus est beaucoup plus lourd. Le modèle fait d’abord une passe avant pour produire un résultat, puis il compare ce résultat avec la réponse attendue et repart dans l’autre sens pour déterminer quelles parties du réseau doivent être corrigées. C’est la rétropropagation.
• Pour pouvoir revenir en arrière, le système doit conserver en mémoire les résultats intermédiaires produits par chaque couche, appelés activations. Il faut également stocker les poids du modèle, les gradients utilisés pour les corriger et les états de l’optimiseur qui suit l’évolution de chaque paramètre. Tous ces éléments s’additionnent et font exploser la consommation mémoire, bien au-delà de ce que peut contenir une seule carte graphique.
• L’exemple d’un modèle de 10 milliards de paramètres illustre l’ampleur du problème : les poids seuls occupent environ 40 Go en FP32. En ajoutant les gradients, les états de l’optimiseur et les activations, la mémoire nécessaire peut rapidement dépasser les 100 Go. C’est pour cette raison que l’entraînement nécessite plusieurs GPU très puissants. Et lorsque plusieurs GPU travaillent ensemble, ils doivent échanger leurs informations en permanence, ce qui ajoute encore un goulot d’étranglement. Diffusion Blocks cible précisément ces deux difficultés : réduire la mémoire en entraînant un bloc à la fois et limiter la communication entre machines.

La rétropropagation et les activations : le vrai gouffre mémoire

Le modèle effectue d'abord une passe vers l'avant pour produire un résultat. Ensuite, il compare ce résultat avec la réponse attendue et repart dans l'autre sens pour déterminer quelle partie du réseau doivent être corrigé.

• La rétropropagation est le mécanisme clé de l’apprentissage supervisé. Pendant la passe avant, l’information circule couche par couche et le modèle produit une prédiction. La fonction de perte compare cette prédiction à la cible. La passe arrière calcule ensuite le gradient de la perte par rapport aux paramètres. Pour pouvoir revenir en arrière, il est indispensable de connaître l’état intermédiaire de chaque couche : ce sont les activations. Les stocker toutes en mémoire est extrêmement coûteux, et c’est ce coût qui limite la taille des modèles que l’on peut entraîner.
• Outre les activations, il faut également conserver les poids du réseau, les gradients qui indiquent comment ajuster ces poids, et les états de l’optimiseur. Ce dernier est souvent négligé dans les explications, mais il joue un rôle crucial : des algorithmes comme Adam maintiennent une moyenne historique des gradients pour chaque paramètre. Un modèle de 10 milliards de paramètres peut ainsi nécessiter plusieurs centaines de Go, bien au-delà de la mémoire d’un seul GPU.
• La solution classique consiste à répartir le modèle sur plusieurs GPU, mais cela génère des communications intensives. Chaque couche dépend des précédentes, donc les GPU doivent synchroniser leurs calculs. Cette synchronisation devient un facteur limitant, tout comme le coût des interconnexions rapides. Diffusion Blocks propose une autre voie : si chaque bloc sait exactement quelle transformation il doit effectuer, il peut être entraîné séparément, sans avoir à garder tout le réseau en mémoire en permanence.

Diffusion Blocks : découper le Transformer en blocs indépendants

Au lieu de charger les 12 couches et de calculer les gradients de l'ensemble du modèle, le système entraîne uniquement quatre couches à la fois.

• Concrètement, un Transformer de 12 couches peut être divisé en trois blocs de quatre couches. Pendant l’entraînement, on charge seulement un bloc dans la mémoire GPU et on calcule les gradients pour ce bloc. Une fois le bloc entraîné, on passe au suivant. La mémoire nécessaire dépend alors principalement de la taille du plus grand bloc, et non plus de celle du modèle complet. Avec trois blocs, la réduction théorique est d’environ trois fois ; avec quatre blocs, d’environ quatre fois.
• Cette réduction n’est possible que si chaque bloc apprend précisément ce qu’il doit faire. Pour cela, les chercheurs s’appuient sur une analogie avec les modèles de diffusion. Dans la génération d’images, le modèle part du bruit et le retire progressivement : d’abord les grandes formes, puis la structure, enfin les détails. De la même façon, les couches d’un Transformer peuvent être vues comme une cascade de débruitage : chaque groupe de couches reçoit une information imparfaite, effectue une petite correction, puis transmet un état légèrement plus propre au groupe suivant.
• Avec Diffusion Blocks, le premier bloc devient spécialiste des informations fortement bruitées, le deuxième traite une version déjà structurée, et le dernier produit le résultat final. C’est une chaîne de production où chaque équipe a une mission précise. Comme les équipes n’ont pas besoin de voir toute la chaîne en même temps, il devient possible d’entraîner les blocs indépendamment. Cette indépendance réduit aussi les communications entre GPU, puisque chaque machine peut travailler sur un bloc différent sans échanger en permanence ses activations et ses gradients.

Un apprentissage calqué sur le débruitage progressif

Chaque groupe de couche reçoit une information imparfaite, effectue une petite correction puis transmet un état légèrement plus propre au groupe suivant.

• Le lien entre Transformers et modèles de diffusion est au cœur de l’innovation. Dans un modèle de diffusion classique, le débruiteur commence par reconstruire la structure globale, puis affine progressivement les détails. Les chercheurs ont remarqué que les couches d’un Transformer font quelque chose d’analogue : les premières couches extraient des représentations grossières, les couches intermédiaires construisent la structure, et les dernières couches produisent une sortie précise. Diffusion Blocks transforme cette analogie en principe d’entraînement.
• Chaque bloc étant spécialisé dans une étape du débruitage, il peut apprendre sa tâche sans attendre les résultats des blocs voisins. Cela rappelle une chaîne de production : une première équipe s’occupe de la forme générale, une autre construit la structure, la dernière ajoute les finitions. La conséquence est importante : il n’est plus nécessaire de conserver toute la chaîne de production en mémoire en permanence, puisque chaque équipe s’entraîne séparément.
• Cette spécialisation automatique est rendue possible par la fonction objectif du modèle, qui pousse chaque bloc à minimiser son erreur sur l’étape qu’il doit accomplir. Les auteurs parlent de « diffusion » parce que le passage de l’information brute à l’information propre suit un processus de débruitage en plusieurs étapes. C’est cette vision qui distingue Diffusion Blocks de simples méthodes de partitionnement, et qui explique pourquoi les blocs peuvent être entraînés en parallèle sans perte majeure de performance.

Des résultats expérimentaux compétitifs, parfois supérieurs

Le modèle classique obtenait environ 60,5 % de précision. Diffusion Blocks atteignait 59,3 % avec une empreinte mémoire théorique environ trois fois plus faible pour l'entraînement.

• Les chercheurs ont testé Diffusion Blocks sur plusieurs types de Transformers : classification d’images, génération d’images et génération de texte. En classification d’images, avec trois blocs, le modèle classique atteignait environ 60,5 % de précision, contre 59,3 % pour Diffusion Blocks. La différence est donc très faible, alors que la mémoire d’entraînement théorique est réduite d’un facteur trois. Ce résultat montre que le découpage en blocs n’a pas fait perdre de capacité de raisonnement au modèle.
• En génération d’images, le résultat est encore plus surprenant. Sur ImageNet, le modèle entraîné normalement obtenait un score FID de 12,09, tandis que Diffusion Blocks atteignait 10,63. Le score FID mesure la proximité entre les images générées et les images réelles, en tenant compte de la qualité et de la diversité ; plus il est bas, meilleur est le résultat. La méthode utilisant moins de mémoire a donc légèrement surpassé l’entraînement classique, ce qui contredit l’intuition selon laquelle réduire la mémoire devrait dégrader la performance.
• Enfin, un petit modèle de langage inspiré de Llama 2, divisé en quatre blocs, a conservé des performances proches de celles d’un modèle entraîné normalement. Cette troisième expérience est importante, car elle suggère que la méthode n’est pas réservée à la vision par ordinateur. Même si les modèles testés restent modestes par rapport aux géants comme GPT, Claude ou DeepSeek, la généralité du résultat est encourageante. Il reste cependant à prouver que l’approche fonctionne à l’échelle de modèles de plusieurs milliards de paramètres.

Limites, conséquences et avenir de Diffusion Blocks

Il faut quand même rester prudent parce que les expériences qui sont démontrées ici, bah elles sont quand même encore très loin de l'échelle de modèles comme GPT, Claude ou encore Dipsic.

• La première limite est le nombre de blocs. Lorsqu’on augmente trop le découpage, chaque bloc contient un nombre trop restreint de couches et perd sa capacité à comprendre la transformation demandée. En génération d’images, deux ou trois blocs donnent les meilleurs résultats ; avec six blocs, la qualité se dégrade nettement. Il existe donc un compromis : plus on crée de blocs, plus on économise de mémoire, mais si les blocs deviennent trop petits, ils ne sont plus assez puissants pour apprendre correctement leur étape.
• Si Diffusion Blocks parvenait à passer à l’échelle, les conséquences économiques seraient considérables. Des laboratoires moins dotés pourraient entraîner ou affiner des modèles beaucoup plus grands avec des GPU plus accessibles. L’entraînement distribué serait aussi simplifié, car les blocs indépendants réduisent le besoin de communication entre machines. On pourrait alors s’affranchir en partie des interconnexions ultra-rapides comme NVLink ou InfiniBand, et utiliser des infrastructures moins coûteuses et plus flexibles.
• Il faut toutefois rester prudent. Les expériences de Sakana AI sont encore loin de l’échelle des très grands modèles de fondation. Les chercheurs reconnaissent eux-mêmes que le passage à plus grande échelle reste à démontrer. De plus, Diffusion Blocks ne réduit pas la mémoire nécessaire à l’inférence : pour générer du texte avec un modèle autorégressif, il faut toujours charger l’ensemble des couches en mémoire. La méthode ne permet donc pas de faire fonctionner un modèle géant sur une petite carte graphique.
• En conclusion, Diffusion Blocks remet en question une règle fondamentale de l’intelligence artificielle moderne : l’idée que toutes les couches d’un réseau doivent obligatoirement être entraînées ensemble. En transformant chaque groupe de couches en spécialiste d’une étape de débruitage, Sakana AI montre qu’un entraînement bloc par bloc reste compétitif. C’est une recherche prometteuse, pas encore une révolution industrielle, mais elle pourrait rendre l’entraînement des intelligences artificielles beaucoup plus accessible et bouleverser une partie de l’économie actuelle des GPU.