Résumé technique : Diffusion par blocs de triplets RWKV (B3D-RWKV)

Énoncé du problème

Les modèles de langage Transformer causaux font face à deux limitations structurelles principales : un décodage strictement séquentiel, qui empêche la parallélisation, et des coûts d'attention quadratiques par étape, qui rendent l'inférence en contexte long coûteuse en calcul. Bien que les modèles causaux à temps linéaire (par exemple, RWKV) et les modèles de diffusion discrets abordent chacun ces faiblesses indépendamment, leur intégration a été intrinsèquement incohérente. Les modèles de diffusion nécessitent une attention bidirectionnelle pour débruiter des blocs de tokens en parallèle, alors que les modèles causaux sont strictement unidirectionnels. Les tentatives existantes pour les combiner requièrent souvent des modifications architecturales qui compromettent l'efficacité du socle à temps linéaire ou nécessitent un entraînement à partir de zéro.

Méthodologie

Les auteurs proposent B3D-RWKV, une variante RWKV de style diffusion qui unifie l'efficacité d'inférence $O(L)$ avec une diffusion discrète bidirectionnelle et parallèle, sans altérer le socle causal sous-jacent. L'innovation centrale est la Disposition par blocs de triplets, une stratégie d'entraînement et d'inférence conçue pour simuler un accès bidirectionnel au sein d'une architecture strictement causale.

1. Disposition par blocs de triplets (Entraînement)

Pour permettre la diffusion sur un modèle causal, chaque bloc logique de génération de taille $B$ est étendu en trois blocs physiques contigus au sein de l'échantillon d'entraînement :

• $b_1$ (Copie masquée) : Une version masquée du bloc logique.
• $b_2$ (Copie masquée soumise à la perte) : Une copie masquée identique où la perte de débruitage est calculée.
• $b_3$ (Copie propre de vérité terrain) : Une version non masquée qui rafraîchit l'état récurrent pour le bloc suivant.

Accès pseudo-bidirectionnel :
Comme le socle causal traite les tokens strictement de gauche à droite, l'état caché arrivant à n'importe quelle position masquée dans $b_2$ a déjà traité l'intégralité de $b_1$. Puisque $b_1$ et $b_2$ partagent le même motif de masquage, les tokens non masqués dans $b_1$ sont efficacement « absorbés » dans l'état caché avant que le modèle ne tente de prédire les tokens masqués dans $b_2$. Cela confère au modèle un accès pseudo-bidirectionnel à son propre contexte non masqué au sein du bloc logique, satisfaisant ainsi les exigences bidirectionnelles de la diffusion tout en maintenant la nature unidirectionnelle du socle RNN/SSM.

Objectif d'entraînement :
Le modèle est entraîné à l'aide d'un objectif de prédiction de masques. La perte se compose de :

• Perte d'entropie croisée ($L_{CE}$) : Calculée sur les positions supervisées où les tokens sont masqués.
• Perte parallèle consciente de la confiance ($L_{CAP}$) : Suivant LLaDA-2.0, ce terme affine la distribution de probabilité sur les positions déjà correctement prédites, encourageant l'échantillonneur d'inférence à valider plus de tokens par étape.

2. Inférence : Débruitage itératif par blocs

Lors de l'inférence, le modèle génère un bloc logique à la fois. Le processus implique :

1. Initialiser un bloc avec tous les tokens masqués.
2. Exécuter des étapes itératives de débruitage où le modèle prédit les probabilités top-1 pour les positions masquées.
3. Valider les tokens dont la probabilité top-1 dépasse un seuil $\tau$.
4. Un mécanisme de repli assure la progression en validant les $k_{min}$ tokens les plus confiants si moins de $k_{min}$ dépassent le seuil.
5. Une fois toutes les positions validées, le bloc propre est ajouté au préfixe, et le prochain bloc logique commence.

Contributions clés

1. B3D-RWKV-7.2B : Le premier modèle de langage RNN à temps linéaire de style diffusion, entraîné à l'échelle 7B avec un objectif de prédiction de masques. Il intègre la sélection parallèle de tokens dans le socle RWKV-7 sans modifier ses paramètres originaux.
2. Cadre par blocs de triplets : Une recette d'entraînement universelle qui convertit des socles strictement causaux (satisfaisant la causalité stricte et les propriétés de propagation avant de l'état) en modèles de diffusion par blocs. Cette approche évite la complexité architecturale de la modification du socle en un bloc bidirectionnel.
3. Efficacité et performance : La méthode réalise des améliorations significatives du débit par rapport au décodage causal standard, tout en maintenant des capacités de raisonnement compétitives.

Résultats expérimentaux

Le modèle a été évalué sur une suite de 8 tâches incluant des connaissances générales (MMLU, ARC, PIQA, RACE) et des tâches de raisonnement/mathématiques (GSM8K, MATH, GPQA).

• Précision : B3D-RWKV-7.2B atteint une précision comparable à la ligne de base RWKV-7 et à d'autres modèles de diffusion de taille similaire (par exemple, LLaDA-8B, Dream-7B). Il surpasse notablement les lignes de base sur des tâches de raisonnement comme ARC-Challenge et RACE, ce qui est attribué à la perception pseudo-bidirectionnelle améliorant le raisonnement.
• Débit : Le modèle démontre un accélération moyenne de 1,6× du débit de décodage par rapport à la ligne de base RWKV-7 sur des longueurs de contexte allant de 1K à 512K tokens. Avec des paramètres d'échantillonnage ajustés, une accélération de 2,02× est réalisable avec un léger compromis sur la qualité.
• Raisonnement mathématique : L'article note une légère baisse de performance sur des tâches mathématiques complexes (par exemple, MATH) par rapport à la ligne de base causale. Les auteurs l'attribuent à la nature du décodage parallèle, qui peine avec la vérification hautement structurée et étape par étape requise pour les problèmes mathématiques symboliques où aucun crédit partiel n'est accordé.

Signification et revendications

L'article revendique offrir une voie efficace pour transformer des modèles de langage causaux pré-entraînés en modèles de langage de diffusion sans changements architecturaux. En tirant parti de la disposition par blocs de triplets, B3D-RWKV unifie avec succès l'efficacité d'inférence $O(L)$ des modèles à temps linéaire avec les capacités de décodage parallèle de la diffusion.

Les auteurs positionnent ce travail comme une solution structurelle à l'incompatibilité entre les paradigmes causal et de diffusion. Ils reconnaissent des limitations, notamment la validation actuelle sur un seul socle (RWKV-7), le coût de longueur de séquence physique 3× pendant l'entraînement, et l'utilisation d'un ensemble de données SFT relativement petit sans alignement RL, ce qui peut contribuer à des régressions de précision sur des tâches spécifiques. Cependant, la contribution centrale reste la démonstration qu'une génération de diffusion parallèle à haut débit est réalisable sur des socles strictement causaux et à temps linéaire par la disposition des données plutôt que par une réingénierie architecturale.