🤖 machine learning

Learned Relay Representations for Forward-Thinking Discrete Diffusion Models

Ce papier présente les Représentations de Relais Appris (Relay), une méthode permettant aux Modèles de Diffusion Masqués de propager des informations latentes à travers les étapes de débruitage via un canal différentiable par token, améliorant ainsi les performances sur les tâches de planification et de codage tout en réduisant considérablement la latence d'inférence.

Auteurs originaux : Benjamin Rozonoyer, Jacopo Minniti, Dhruvesh Patel, Neil Band, Avishek Joey Bose, Tim G. J. Rudner, Andrew McCallum

Publié 2026-05-25

📖 5 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Benjamin Rozonoyer, Jacopo Minniti, Dhruvesh Patel, Neil Band, Avishek Joey Bose, Tim G. J. Rudner, Andrew McCallum

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Le Problème : L'Amnésie du « Réinitialisation Totale »

Imaginez que vous essayez de résoudre une énigme complexe, comme un Sudoku ou l'écriture d'un morceau de code. Vous utilisez un assistant très intelligent (un modèle d'IA) pour vous aider.

Actuellement, cet assistant fonctionne comme une personne atteinte d'amnésie qui oublie tout dès qu'elle termine une seule phrase.

L'assistant regarde une page blanche remplie de points d'interrogation.
Il réfléchit intensément, effectue des calculs mentaux complexes et décide de remplir une ou quelques lettres.
La Réinitialisation Totale : Dès qu'il écrit ces lettres, il jette immédiatement tous les pensées complexes, les calculs et le « ressenti » qu'il avait concernant le reste de la page.
Pour l'étape suivante, il repart de zéro, ne regardant que les lettres qu'il vient d'écrire, oubliant complètement le travail interne riche qu'il venait d'accomplir.

Ceci est appelé le problème de la « Réinitialisation Totale ». Le papier soutient que cela est inefficace. C'est comme un chef qui goûte une soupe, note « il faut du sel », puis oublie immédiatement le goût du bouillon avant d'ajouter le sel. Il doit régoûter toute la soupe depuis le début à chaque fois qu'il ajoute un ingrédient.

La Solution : Le Témoin du « Relais »

Les auteurs proposent une nouvelle méthode appelée Relais.

Imaginez une course de relais. Au lieu que le coureur s'arrête, oublie la course et reparte de zéro, il passe un témoin au coureur suivant. Ce témoin transporte l'élan, la stratégie et la fatigue du coureur précédent afin que le suivant puisse reprendre exactement là où il s'est arrêté.

Dans le monde de l'IA, le « témoin » est un flux continu d'informations (un état caché) que le modèle transporte vers l'avant.

Avant le Relais : Le modèle calcule, écrit un token et supprime le calcul.
Avec le Relais : Le modèle calcule, écrit un token, et passe une « note » à lui-même pour l'étape suivante. Cette note dit : « Hé, je réfléchissais à ce motif spécifique, et je suis sûr à 80 % de cette prochaine partie. »

Comment Cela Fonctionne : Apprendre à Passer le Témoin

Le papier introduit quelques astuces clés pour rendre cela fonctionnel :

L'Entraînement « Prévisionnel » :
Habituellement, les modèles d'IA sont entraînés à simplement obtenir la bonne réponse actuelle. Les auteurs ont entraîné le modèle à être « prévisionnel ». Ils ont enseigné au modèle : « Ne te contente pas d'écrire la bonne lettre aujourd'hui ; écris une note (le relais) qui t'aidera à écrire la bonne lettre demain. »
Rétropropagation Tronquée (BPTT) :
C'est un terme mathématique sophistiqué pour « regarder en avant pour apprendre ». Imaginez que vous pratiquez une routine de danse. Au lieu de pratiquer seulement un mouvement, vous pratiquez une séquence de trois mouvements. Si vous trébuchez sur le troisième mouvement, vous réalisez : « Oh, j'ai raté le premier mouvement parce que je ne me suis pas préparé pour le troisième. »
Le papier utilise une méthode appelée BPTT tronquée pour permettre au modèle de regarder quelques étapes en avant pendant l'entraînement. Il apprend à transmettre des informations vers l'avant spécifiquement pour faciliter les étapes futures.
Le Canal « Doux » :
L'information transmise vers l'avant n'est pas seulement un mot ou une lettre ; c'est un nombre « doux » (une valeur continue). Pensez-y comme à un gradateur plutôt qu'à un interrupteur marche/arrêt. Cela permet au modèle de transporter des indices subtils et des probabilités qui n'ont pas encore été transformés en une décision finale.

Les Résultats : Plus Rapide et Plus Intelligent

Les auteurs ont testé cela sur deux choses :

Énigmes Sudoku : Ils ont traité le Sudoku comme une tâche de planification. Le modèle « Relais » a résolu les énigmes avec moins de tentatives (moins de « passages en avant ») et a fait moins d'erreurs que le modèle standard. Il était meilleur pour garder la « vue d'ensemble » à l'esprit tout en remplissant la grille.
Programmation (Écriture de Logiciels) : Ils ont pris une IA de pointe pour le code (Fast-dLLM v2) et y ont ajouté le système Relais.
- Meilleure Précision : Il a écrit un meilleur code.
- Vitesse Accrue : Il a dû exécuter son « moteur de réflexion » 32 % de moins pour obtenir le même résultat.

Pourquoi Cela Compte

Le papier affirme qu'en arrêtant la « Réinitialisation Totale » et en permettant au modèle de transporter ses propres « pensées » vers l'avant comme un témoin de relais, nous pouvons rendre les modèles d'IA :

Plus Intelligents : Ils peuvent raisonner mieux sur de longues séquences (comme du code complexe ou des énigmes logiques).
Plus Rapides : Ils n'ont pas à tout recalculer depuis le début à chaque fois, ce qui économise du temps et de l'énergie.

En bref, le papier enseigne aux modèles d'IA à cesser d'oublier leur propre travail et à commencer à s'appuyer dessus, étape par étape, tout comme un humain le fait lorsqu'il résout un problème difficile.

Résumé Technique : Représentations de Relais Apprises pour les Modèles de Diffusion Discrète à Pensée Prospective

1. Énoncé du Problème : Le « Réinitialisation Dure » dans les Modèles de Diffusion Masquée

Les Modèles de Diffusion Masquée (MDM) génèrent des séquences discrètes par débruitage itératif. À partir d'une séquence entièrement masquée, le modèle dévoile itérativement une fraction des positions. Cependant, une inefficacité critique existe dans l'inférence standard des MDM : après chaque passage avant, le modèle rejette les calculs internes continus riches (états cachés) générés pour les positions masquées. L'étape suivante commence uniquement à partir des tokens discrets nouvellement engagés.

Les auteurs qualifient cela de problème de réinitialisation dure. Bien que les MDM effectuent plusieurs passages avant par génération, l'incapacité de réutiliser le calcul continu intermédiaire empêche le modèle d'amortir efficacement le raisonnement à travers les étapes. Cela contraste avec le calcul récurrent, où le déroulement d'un modèle à paramètres fixes sur plusieurs étapes a été montré pour élargir la classe de fonctions que le modèle peut approximer, améliorant ainsi les performances sur des tâches de raisonnement complexes. La question centrale abordée est de savoir comment permettre aux MDM de prendre en charge un calcul récurrent qui transporte des informations plus riches à travers les étapes de dévoilement séquentielles sans rejeter les états latents intermédiaires.

2. Méthodologie : Représentations de Relais Apprises (Relay)

La solution proposée, Relay, introduit un état latent continu et différentiable transmis vers l'avant à travers les étapes de décodage, permettant au modèle d'être « à pensée prospective ».

Mécanisme Central

Trajectoire d'État Augmentée : À chaque étape d'inférence $k$ , le modèle maintient un état augmenté $s_k = (x_{t_k}, h_k)$ , où $x_{t_k}$ est la séquence partiellement masquée actuelle et $h_k$ est un état de relais appris (vecteur latent continu).
Architecture : Le modèle se compose d'un socle $f_\theta$ $f_{θ}$ , d'embeddings de tokens $Emb_\theta$ $E m b_{θ}$ , d'une tête de dé-embedding $UnEmb_\theta$ $U n E m b_{θ}$ , et d'un module de relais spécifique $R_\theta$ $R_{θ}$ .
- Le prochain état caché est calculé comme suit : $h_{k+1} = f_\theta(Emb_\theta(x_{t_k}) + R_\theta(h_k))$ .
- L'état de relais $h_k$ est projeté et ajouté aux tokens embarqués avant d'entrer dans le socle.
- Les logits de sortie $\ell_k$ sont dérivés de $h_{k+1}$ via $UnEmb_\theta$ .
Entraînement par Rétropropagation Tronquée dans le Temps (BPTT) : Contrairement aux MDM standards entraînés avec des étapes indépendantes, Relay est entraîné en utilisant la Rétropropagation Tronquée dans le Temps (BPTT).
- Le modèle effectue un déploiement de $K$ étapes (par exemple, $K=2$ ) à partir d'une séquence entièrement masquée.
- Pendant le déploiement, les positions sont dévoilées selon une politique (par exemple, des seuils de confiance), et les tokens engagés sont « forcés par l'enseignant » vers les valeurs de vérité terrain pour empêcher l'accumulation d'erreurs pendant le déploiement.
- La perte est la somme des pertes d'entropie croisée sur les $K$ étapes.
- Les gradients sont rétropropagés à travers le chemin du relais ( $h_k \to h_{k+1}$ ), permettant au modèle d'apprendre non seulement à prédire l'étape actuelle, mais à encoder dans $h_k$ des informations maximales utiles pour les étapes futures.

Inférence

Au moment de l'inférence, le processus reste largement inchangé en termes de calendrier de dévoilement ou de politique d'échantillonnage. La seule addition est que le modèle transporte l'état de relais $h_k$ vers l'avant parallèlement aux tokens engagés. L'état de relais sert de canal mémoire différentiable pour les informations qui ne sont pas encore engagées dans des tokens discrets.

3. Contributions Clés

Cadre Relay : Une méthode générale pour intégrer un calcul récurrent dans les MDM en entraînant le modèle à transmettre un relais latent appris vers l'avant à travers les étapes de décodage via une BPTT tronquée. Elle est agnostique à l'architecture et compatible avec les techniques d'accélération MDM existantes comme la diffusion par blocs et la mise en cache KV.
Évolutivité et Adaptation : Validation de Relay à l'échelle des Grands Modèles de Langage (LLM) par l'adaptation de paramètres complets de Fast-dLLM v2 (1,5B). La méthode adapte avec succès un Modèle de Langage Discrète (DLM) de pointe avec un nombre minimal de FLOPs d'entraînement supplémentaires.
Améliorations Performance-Latence : Démonstration empirique que Relay fait progresser la frontière de Pareto performance-latence, surpassant le fine-tuning supervisé (SFT) standard sur des tâches de codage tout en réduisant significativement la latence d'inférence.
Analyse de l'Espace de Conception : Études d'ablation extensives sur une tâche de planification basée sur le Sudoku, difficile, qui isolent les contributions de la procédure de déploiement, du canal de relais et de l'objectif d'entraînement BPTT.

4. Résultats Expérimentaux

Tâche de Planification Sudoku

En utilisant l'ensemble de données Sudoku-Extreme, les auteurs ont comparé quatre objectifs d'entraînement :

MLM : Diffusion masquée uniforme standard (étape unique, sans relais).
Rollout : Déploiement multi-étapes avec forçage par l'enseignant mais sans canal de relais ( $R_\theta \equiv 0$ ).
Relay (sg) : Inclut le chemin de relais mais utilise des gradients arrêtés (pas de BPTT à travers le relais).
Relay : Méthode complète avec $K=2$ BPTT à travers le relais.

Constats :

Frontière Précision-NFE : Relay a atteint le meilleur compromis entre précision et Nombre d'Évaluations Avant (NFE).
Impact des Composantes : Le passage de MLM à Rollout a amélioré les performances. L'ajout du canal de relais (Relay sg) a fourni un saut significatif. Enfin, l'activation de la BPTT (Relay) a produit les meilleurs résultats, démontrant que le modèle apprend à maintenir le plateau partiel cohérent avec lui-même sous un dévoilement plus agressif.
Efficacité : À un seuil de confiance fixe, Relay a engagé plus de cellules par passage avant tout en maintenant la légalité (aucune violation de ligne/colonne/boîte), atteignant la même précision en moins de passages avant au total par rapport aux références.

Adaptation de Modèle Préentraîné (Fast-dLLM v2)

Les auteurs ont adapté Fast-dLLM v2 (1,5B) en utilisant Relay sur un mélange d'ensembles de données de code et de mathématiques.

Performance en Codage : Sur les benchmarks HumanEval et MBPP, Relay a surpassé le SFT standard.
- HumanEval : Relay a atteint 42,1 % de précision (contre 38,4 % pour SFT) avec 88,3 NFE (contre 130,7 pour SFT).
- MBPP : Relay a atteint 46,6 % de précision (contre 43,9 % pour SFT) avec 78,8 NFE (contre 84,8 pour SFT).
Latence : Relay a réduit la latence d'inférence d'jusqu'à 32 % par rapport au SFT standard tout en améliorant la précision.
Surcharge Mémoire : Le profilage sur un GPU A100 a montré que, bien que Relay augmente la mémoire d'activation active pendant le deuxième passage avant, l'utilisation maximale de la mémoire reste presque identique au SFT standard (environ 20,1 GiB contre 21,2 GiB) car le goulot d'étranglement est le passage arrière de la projection du vocabulaire, et non le stockage de l'état de relais.

5. Importance et Revendications

L'article revendique que Relay change fondamentalement la nature de la génération itérative non-autoregressive en permettant aux MDM d'accumuler un calcul continu intermédiaire.

Pensée Prospective : La méthode permet aux modèles d'optimiser explicitement les représentations intermédiaires pour les décisions de dévoilement futures, plutôt que de traiter chaque étape comme un problème de prédiction frais.
Avancement de la Frontière de Pareto : En réduisant le nombre de passages avant nécessaires pour atteindre une précision cible, Relay fait progresser la frontière performance-latence, rendant les modèles de diffusion discrète plus efficaces pour les tâches complexes de raisonnement et de génération.
Compatibilité : L'approche est présentée comme une amélioration intégrable qui fonctionne avec les DLM de pointe existants et les techniques d'accélération (par exemple, diffusion par blocs, mise en cache KV) sans nécessiter de refontes architecturales.

Les auteurs concluent que Relay représente une étape significative vers une génération non-greedy et à pensée prospective, offrant un substrat continu pour le raisonnement latent qui peut être davantage exploré pour l'interprétabilité et l'application dans d'autres modalités (par exemple, images, molécules).