L'entraînement des grands modèles d'intelligence artificielle nécessite aujourd'hui des ressources considérables, principalement en raison de la mémoire requise pour stocker les états intermédiaires de l'ensemble du réseau. Une nouvelle approche, baptisée DiffusionBlocks, propose un changement de paradigme : diviser le réseau en blocs et les entraîner indépendamment les uns des autres, réduisant ainsi la mémoire nécessaire à celle d'un seul bloc.

Présenté à la conférence ICLR 2026, ce travail est le fruit d'une collaboration entre Makoto Shing et Takuya Akiba (Sakana AI) et Masanori Koyama (Université de Tokyo). Les chercheurs publient leurs résultats sur arXiv et OpenReview.

Le goulot d'étranglement de la mémoire

L'entraînement des réseaux de neurones modernes, y compris les Transformers, repose sur la rétropropagation du gradient à travers l'ensemble du réseau (end-to-end). Cette méthode oblige à conserver en mémoire toutes les activations intermédiaires, ce qui conduit à une consommation mémoire qui croît linéairement avec la profondeur du modèle. Alors que les modèles comptent désormais des centaines de milliards de paramètres, seules quelques organisations disposent des ressources matérielles nécessaires pour les entraîner.

L'idée d'un entraînement par blocs (block-wise training) n'est pas nouvelle en soi, mais les tentatives antérieures se sont surtout concentrées sur la classification d'images et ont généralement accusé un retard de performance par rapport à l'entraînement classique. L'extension de cette approche aux architectures Transformer modernes et aux tâches génératives restait un défi ouvert.

DiffusionBlocks : un cadre théorique inspiré des modèles de diffusion

DiffusionBlocks s'appuie sur une observation clé : les connexions résiduelles d'un réseau peuvent être interprétées comme les pas discrets d'une équation différentielle ordinaire, la même classe d'équations qui sous-tend les modèles de diffusion. Les chercheurs attribuent alors à chaque bloc un rôle précis : celui de se rapprocher progressivement de la cible, à l'image du processus inverse (débruitage) d'un modèle de diffusion.

Concrètement, plutôt que de traiter le réseau comme une boîte noire qui doit apprendre une cartographie entrée-sortie de façon indifférenciée, DiffusionBlocks assigne à chaque bloc un segment du chemin allant du bruit à la cible. Chaque bloc peut alors être entraîné indépendamment, sans dépendre des sorties des autres blocs pendant l'apprentissage.

La conversion d'un réseau existant en un modèle entraînable par blocs s'effectue en trois modifications :

  • Partition : diviser les L couches en B blocs.
  • Attribution de plages de bruit : chaque bloc est responsable d'une plage de « proximité à la cible ».
  • Conditionnement : un module de conditionnement permet à chaque bloc de reconnaître la plage qui lui est attribuée.

Lors de l'entraînement, un seul bloc est échantillonné aléatoirement à chaque itération. Les autres blocs n'étant pas calculés, la consommation mémoire est réduite à environ 1/B de celle de l'entraînement classique.

Résultats expérimentaux

Les expérimentations menées par l'équipe couvrent plusieurs architectures : classification d'images, génération d'images et génération de texte. Selon les auteurs, DiffusionBlocks atteint des performances comparables à celles de l'optimisation de bout en bout, tout en offrant une réduction significative de l'empreinte mémoire.

Ces travaux s'inscrivent dans la continuité de NoProp, un travail antérieur partageant l'idée générale d'utiliser un cadre de diffusion pour l'entraînement par blocs, mais dont les résultats se limitaient à la classification d'images. DiffusionBlocks étend cette approche aux Transformers modernes et aux tâches génératives.

Implications pour l'accessibilité de l'IA

En réduisant la mémoire nécessaire à l'entraînement des grands modèles, DiffusionBlocks pourrait contribuer à démocratiser la recherche en intelligence artificielle. Les auteurs soulignent que leur motivation est de permettre l'entraînement de modèles de toute taille, indépendamment du matériel disponible, et de rendre le développement de l'IA plus accessible.