Résumé Technique : DiNa-LRM – Modélisation de Récompense Latente Native par Diffusion

1. Énoncé du Problème

L'alignement des modèles de diffusion et d'appariement de flux avec les préférences humaines repose fortement sur des fonctions de récompense. Les approches actuelles les plus avancées utilisent principalement des Modèles Vision-Langage (VLM) comme backbones de récompense. Bien que les VLM offrent des capacités discriminatives robustes, ils présentent deux limitations critiques pour l'alignement des modèles de diffusion :

1. Surcharge Computatoire et Mémoire : Les récompenses basées sur les VLM sont coûteuses à interroger de manière répétée lors de l'optimisation, créant un goulot d'étranglement pour l'efficacité de l'entraînement.
2. Inadéquation de Domaine : Les VLM opèrent généralement dans l'espace des pixels, tandis que les générateurs de diffusion latente sont entraînés et optimisés dans l'espace latent. Cette inadéquation pixel-latente complique l'alignement, en particulier pour les méthodes de gradient de récompense, et introduit une surcharge système inutile.

À l'inverse, bien que les modèles de diffusion aient montré des promesses en tant que backbones discriminatifs, les tentatives antérieures de les utiliser comme récompenses ont souvent été liées à des algorithmes d'optimisation spécifiques à un niveau d'étape ou se sont concentrées sur des états intermédiaires bruyants sans établir de cadre général comparable aux VLM pour le scoring d'échantillons propres.

2. Méthodologie : DiNa-LRM

Les auteurs proposent DiNa-LRM (Modèle de Récompense Latente Native par Diffusion), un cadre qui formule l'apprentissage des préférences directement sur les états de diffusion bruyants au sein de l'espace latent. La méthode se compose de trois composants principaux :

A. Formulation des Préférences Natives par Diffusion

Au lieu d'évaluer directement les échantillons propres ($x_0$), DiNa-LRM étend le modèle de Thurstone aux états de diffusion bruyants ($x_t$).

• Incertitude Calibrée par le Bruit : Les auteurs reconnaissent que les jugements de préférence sur des états bruyants sont intrinsèquement moins certains à mesure que le niveau de bruit augmente. Ils introduisent une incertitude de comparaison calibrée par le bruit $\sigma_u^2(t)$ qui évolue avec le niveau de bruit de diffusion $\sigma^2(t)$ :
  $$\sigma_u^2(t) = k\sigma^2(t) + \sigma_u^2$$
  Cette formulation garantit que l'objectif de récompense reste cohérent avec la variété bruyante rencontrée lors du pré-entraînement par diffusion. Elle agit comme un régularisateur conscient de l'incertitude, empêchant les gradients à forte variance de déstabiliser l'entraînement dans les régimes à fort bruit.
• Objectif d'Entraînement : Le modèle est optimisé en utilisant une perte de fidélité sur des paires de préférences $(x_t^+, x_t^-)$, où la probabilité de comparaison est modulée par l'incertitude dépendante du pas de temps.

B. Architecture de Récompense Latente

DiNa-LRM est instancié sur un backbone de diffusion latente pré-entraîné (par exemple, SD3.5-Medium) avec une tête de scoring légère basée sur des requêtes.

• Extraction de Caractéristiques Multi-couches : Des caractéristiques visuelles et textuelles sont extraites de multiples couches du backbone.
• Adaptation Conditionnée par le Pas de Temps : Les caractéristiques sont modulées par des embeddings de pas de temps utilisant FiLM (Modulation Linéaire Spécifique aux Caractéristiques) pour rendre la tête de récompense explicitement consciente du niveau de bruit.
• Scoring Q-Former à Porte : Un Q-Former à porte agrège ces caractéristiques multi-couches conditionnées par le pas de temps pour produire un score de récompense scalaire. Cette architecture opère entièrement dans l'espace latent du VAE, évitant le décodage dans l'espace des pixels.

C. Mise à l'Échelle au Moment de l'Inférence par Ensembling de Bruit

Pour améliorer la robustesse et réduire la sensibilité à un pas de temps d'évaluation spécifique, DiNa-LRM prend en charge l'ensembling de bruit.

• Au lieu de scorer un échantillon propre à un seul niveau de bruit fixe, le modèle évalue l'échantillon à plusieurs pas de temps $\{t_k\}$.
• Les caractéristiques extraites à ces différents pas de temps sont concaténées et traitées par la même tête Q-Former pour produire un score agrégé. Cela sert de mécanisme de mise à l'échelle au moment du test "natif par diffusion".

3. Contributions Clés

1. Formulation des Préférences Natives par Diffusion : L'article étend le modèle de préférence de Thurstone aux états de diffusion bruyants en introduisant une incertitude calibrée par le bruit qui évolue avec le niveau de bruit de diffusion, comblant ainsi l'écart distributionnel entre le pré-entraînement et l'apprentissage de la récompense.
2. Mise à l'Échelle au Moment de l'Inférence : Les auteurs proposent une stratégie d'ensembling de bruit qui agrège des caractéristiques multi-pas de temps, fournissant un mécanisme de mise à l'échelle au moment du test qui exploite la trajectoire globale de diffusion plutôt qu'une estimation ponctuelle unique.
3. Validation Empirique : Le travail démontre que les récompenses natives par diffusion peuvent atteindre des performances compétitives avec des juges VLM puissants tout en opérant nativement dans l'espace latent.

4. Résultats Expérimentaux

Les auteurs ont évalué DiNa-LRM sur plusieurs benchmarks et scénarios d'alignement :

• Précision de la Modélisation de Récompense :

  • DiNa-LRM surpasse considérablement les baselines de récompense basées sur la diffusion existantes (par exemple, LRM-SD1.5, LRM-SDXL) sur ImageReward, HPDv2, HPDv3 et GenAI Bench.
  • Il atteint des performances compétitives avec les récompenses basées sur les VLM les plus avancées (par exemple, HPSv3, UnifiedReward) en moyenne, avec un écart notable qui se réduit sur des jeux de données spécifiques comme HPDv2 et HPDv3.
  • Ensembling de Bruit : L'application de l'ensembling de bruit améliore davantage la précision moyenne (par exemple, passant de 71,49 % à 72,48 % en moyenne), validant l'avantage de l'agrégation des preuves à travers les niveaux de bruit.

• Alignement des Préférences (ReFL) :

  • Dans les expériences d'alignement post-entraînement utilisant ReFL (Apprentissage par Feedback de Récompense) sur SD3.5-M, DiNa-LRM a démontré une convergence plus rapide des scores proxy et une amélioration des métriques "golden" (PickScore retenu) par rapport à HPSv3.
  • Efficacité : DiNa-LRM réduit considérablement les coûts de ressources par rapport aux récompenses basées sur les VLM. Il réduit l'utilisation maximale de VRAM de 51,4 %, les FLOPs de calcul de récompense de 71,1 % et les FLOPs de la phase d'optimisation de 46,4 %.

• Généralisation :

  • La formulation se généralise à travers différents backbones de diffusion (SDXL, FLUX.1-Dev, Z-Image-Turbo), surpassant constamment les baselines basées sur la diffusion antérieures.
  • Les expériences de transférabilité montrent qu'un modèle de récompense entraîné sur SD3.5-M peut aligner efficacement SD3.5-L au sein du même espace latent.

5. Importance et Revendications

L'article revendique que DiNa-LRM offre une alternative pratique et efficace aux récompenses VLM dans l'espace des pixels pour l'alignement des préférences. Son importance principale réside dans :

• Avantage Structurel : En opérant nativement dans l'espace latent, il élimine l'inadéquation latent-pixel et la surcharge computationnelle du décodage, le rendant hautement compatible avec les méthodes d'alignement par gradient de récompense.
• Facilité d'Optimisation : Il fournit un signal différentiable et efficace qui améliore la dynamique d'entraînement et réduit les besoins en ressources sans sacrifier la qualité de l'alignement.
• Évolutivité : Le mécanisme d'ensembling de bruit fournit une manière nouvelle et native par diffusion de mettre à l'échelle les performances au moment de l'inférence.

Les auteurs reconnaissent des limites, notant que la récompense est liée à des espaces latents spécifiques (limitant le transfert inter-backbone sans réentraînement) et que la modélisation dans l'espace latent peut sous-estimer certains artefacts au niveau des pixels, conduisant potentiellement à des comportements de piratage de récompense (par exemple, injection humaine spuriaire ou dérive stylistique) lors d'optimisations à long horizon. Ils suggèrent que les travaux futurs devraient se concentrer sur des backbones unifiés plus puissants et l'intégration d'une régularisation légère dans l'espace des pixels.