Adaptive Guidance for Retrieval-Augmented Masked Diffusion Models

Titre : Guidage Adaptatif pour les Modèles de Diffusion Masquée Augmentés par la Récupération (ARAM)

1. Problématique

Les modèles de langage (LLM) souffrent souvent d'hallucinations et de connaissances obsolètes lorsqu'ils s'appuient uniquement sur leurs connaissances paramétriques statiques. La Génération Augmentée par la Récupération (RAG) vise à résoudre ce problème en intégrant un contexte externe. Cependant, dans les scénarios réels, le contexte récupéré peut être bruyant, partiellement pertinent ou contradictoire par rapport aux connaissances internes du modèle. Cela crée un conflit de connaissances (knowledge conflict).

Bien que ce problème ait été étudié dans les modèles de langage autoregressifs, il reste largement inexploré dans les modèles de diffusion basés sur le langage, et plus spécifiquement les Modèles de Diffusion Masquée (MDM). Contrairement aux modèles autoregressifs, les MDM génèrent du texte via un processus itératif de débruitage (unmasking) qui permet des mises à jour parallèles et une raffinement itératif. L'application d'une stratégie de guidage statique (comme le Classifier-Free Guidance ou CFG) dans un contexte RAG pose deux dilemmes majeurs :

1. Sur-guidage (Over-guidance) : Si le contexte récupéré est bruyant ou faux, un guidage fort amplifie les erreurs et provoque des hallucinations.
2. Sous-guidage (Under-guidance) : Si le contexte est fiable mais que le guidage est trop faible, le modèle ne corrige pas correctement ses connaissances paramétriques erronées.

La question centrale est donc : Comment un modèle de diffusion linguistique peut-il déterminer dynamiquement quand faire confiance au contexte récupéré et quand s'appuyer sur ses connaissances paramétriques ?

2. Méthodologie : ARAM

Les auteurs proposent ARAM (Adaptive Retrieval-Augmented Masked Diffusion), un cadre de guidage adaptatif sans entraînement (training-free) conçu spécifiquement pour les MDM en configuration RAG.

Principes Fondamentaux :

• Interprétation Informationnelle : L'influence du contexte récupéré est analysée comme un déplacement de distribution (distributional shift) entre une distribution a priori (sans récupération) et une distribution conditionnelle (avec récupération).
• Rapport Signal-Bruit (SNR) : L'idée clé est de calibrer l'échelle de guidage ($\lambda_t$) à chaque étape de débruitage et pour chaque token en fonction d'un Rapport Signal-Bruit (SNR) dérivé théoriquement.
  • Signal : Mesure l'information gain apporté par le contexte (basé sur la divergence de Kullback-Leibler, $D_{KL}$, entre les distributions conditionnelle et a priori).
  • Bruit : Mesure l'incertitude ou la dispersion introduite par le contexte. Au lieu d'utiliser la variance (qui s'avère instable), les auteurs utilisent l'entropie conditionnelle ($H(p_{cond})$) comme proxy du bruit. Une entropie élevée indique un contexte ambigu ou contradictoire.

Formulation Mathématique :
L'échelle de guidage $\lambda_t$ est calculée dynamiquement selon la formule :
$$\lambda_t = \lambda_{max} \cdot \tanh\left(\beta \cdot \frac{\text{Signal}}{\text{Bruit} + \epsilon}\right)$$
Où :

• $\text{Signal} = D_{KL}(p_{cond} || p_{prior}) + D_{KL}(p_{prior} || p_{cond})$
• $\text{Bruit} = H(p_{cond})$ (Entropie de la distribution conditionnelle)
• $\lambda_{max}$ et $\beta$ sont des hyperparamètres de contrôle.

Comportement Adaptatif :

• Si le contexte est fiable (Signal fort, Entropie faible) : $\lambda_t$ augmente, renforçant l'influence du contexte pour corriger le modèle.
• Si le contexte est bruyant ou contradictoire (Signal faible ou Entropie élevée) : $\lambda_t$ diminue, protégeant le modèle contre les hallucinations en s'appuyant davantage sur son prior paramétrique.

3. Contributions Clés

1. Cadre ARAM : Première approche de guidage adaptatif pour les MDM en RAG, permettant un ajustement dynamique de l'échelle de guidage à l'échelle du token et de l'étape.
2. Analyse Théorique : Dérivation d'une règle de guidage basée sur une perspective informationnelle (SNR) pour résoudre les conflits de connaissances dans les modèles de diffusion discrets.
3. Validation Empirique : Démonstration que l'ajustement dynamique du guidage surpasse les méthodes statiques et les méthodes de décodage contrastif adaptées aux modèles autoregressifs.

4. Résultats Expérimentaux

Les expériences ont été menées sur cinq benchmarks de questions-réponses (QA) riches en connaissances : NQ, TriviaQA, MARCO QA, HotpotQA et T-REx.

• Performance Globale : ARAM surpasse systématiquement les modèles de base (LLaDA et Dream) et les méthodes de comparaison (RAG statique, CAD, A-CFG, SPREAD) en termes de Exact Match (EM) et de F1-score.
  • Exemple : Sur le modèle Dream, ARAM atteint un EM de 37,7 sur NQ contre 35,9 pour le RAG standard et 34,0 pour A-CFG.
• Analyse des Conflits :
  • ARAM augmente le taux de cas où le modèle passe d'une réponse incorrecte (sans RAG) à une réponse correcte (avec RAG) (Positive rate).
  • Il réduit significativement le taux de cas où le modèle passe d'une réponse correcte à une réponse incorrecte à cause d'un contexte nuisible (Negative rate).
• Robustesse au Bruit : Sur les sous-ensembles où le contexte récupéré ne contient pas la réponse ("Non-Answering"), ARAM évite de dégrader les performances, contrairement aux méthodes statiques qui tendent à "sur-corriger".
• Dynamique de Guidage : L'analyse visuelle montre que ARAM maintient un guidage élevé pour les documents pertinents ("Gold Docs") et le supprime pour les documents bruyants ("Noise Docs"), tout en s'adaptant spatialement aux positions de tokens incertains lors du processus de démasquage.

5. Signification et Limites

Signification :
Ce travail comble un vide important dans la recherche sur les modèles de diffusion linguistique. Il démontre que l'application directe de techniques RAG conçues pour les modèles autoregressifs est sous-optimale pour les MDM. La proposition d'un mécanisme de guidage intrinsèquement adaptatif basé sur le SNR offre une solution élégante au problème du conflit de connaissances, améliorant la fiabilité des modèles de diffusion sans nécessiter de réentraînement coûteux.

Limites :

• Surcoût Computationsnel : Le calcul de l'échelle de guidage pour chaque token et chaque étape de débruitage (impliquant le calcul d'entropie et de divergence sur le vocabulaire entier) introduit une latence d'inférence supplémentaire par rapport au CFG statique.
• Portée : Les résultats sont actuellement limités aux tâches de QA en anglais et aux modèles de diffusion ouverts. L'efficacité sur des tâches de génération ouverte ou dans d'autres langues reste à explorer.

En conclusion, ARAM représente une avancée majeure pour l'intégration robuste de la récupération d'informations dans les modèles de génération de texte basés sur la diffusion, en équilibrant dynamiquement la confiance entre les connaissances internes et externes.