Résumé Technique : LeapAlign

Énoncé du Problème

L'article aborde le défi de l'alignement des modèles de matching de flux avec les préférences humaines par le biais d'un post-entraînement. Si les méthodes à gradient direct offrent une alternative prometteuse aux approches à gradient de politique (comme GRPO) en exploitant la nature différentiable de l'échantillonnage par matching de flux, elles rencontrent des obstacles majeurs lorsqu'elles sont appliquées à des trajectoires de génération longues. Plus précisément, la rétropropagation des gradients de récompense à travers toute la séquence d'étapes d'échantillonnage d'EDP (équations différentielles ordinaires) engendre des coûts mémoire prohibitifs et risque une explosion des gradients. Par conséquent, les méthodes à gradient direct existantes limitent généralement les mises à jour à un seul pas de temps proche de l'image finale ou éliminent les gradients imbriqués pour éviter l'instabilité. Cette limitation empêche l'optimisation des étapes de génération précoces, qui sont critiques pour déterminer la disposition globale et la composition structurelle de l'image finale.

Méthodologie : LeapAlign

Les auteurs proposent LeapAlign, une méthode de fine-tuning conçue pour permettre une propagation directe et stable des gradients vers n'importe quelle étape de génération, y compris les premiers pas de temps, sans engendrer de surcharge mémoire excessive. La méthodologie centrale repose sur trois composants clés :

1. Construction de Trajectoire à Saut en Deux Étapes

Au lieu de rétropropager à travers la trajectoire de génération complète, LeapAlign construit une trajectoire à saut raccourcie composée de seulement deux étapes dérivées d'une trajectoire complète standard.

• Processus : À chaque itération, une trajectoire complète est échantillonnée à partir du bruit ($x_1$) jusqu'à une image propre ($x_0$). Deux pas de temps, $k$ et $j$ (où $k > j$), sont sélectionnés aléatoirement.
• Sauts : La méthode exploite la propriété de prédiction par saut en une étape des modèles de flux rectifié pour estimer les latents futurs en une seule étape :
  • Saut 1 : Prédit le latent $x_j$ à partir de $x_k$ en utilisant le champ de vitesse en $x_k$ : $\hat{x}_{j|k} = x_k - (k-j)v_\theta(x_k)$.
  • Saut 2 : Prédit le latent final $x_0$ à partir de $x_j$ en utilisant le champ de vitesse en $x_j$ : $\hat{x}_{0|j} = x_j - j v_\theta(x_j)$.
• Connecteurs Latents : Pour maintenir la différentiabilité tout en alignant les états prédits avec les états réels échantillonnés depuis la trajectoire complète, des « connecteurs latents » sont introduits en utilisant des opérations stop_gradient (par exemple, $x_j = \hat{x}_{j|k} + \text{stop\_gradient}(x_j - \hat{x}_{j|k})$).
• Avantage : Cela réduit le chemin de rétropropagation à deux étapes, maintenant les coûts mémoire constants quelle que soit la longueur de la trajectoire originale, tout en permettant la sélection aléatoire de $(k, j)$ pour mettre à jour n'importe quelle partie du processus de génération.

2. Découplage des Gradients

Pour résoudre le problème de l'explosion des gradients tout en préservant les signaux d'apprentissage utiles, LeapAlign introduit un mécanisme de découplage des gradients.

• Le Problème : La rétropropagation à travers deux étapes introduit un terme de « gradient imbriqué » ($\frac{\partial v_\theta(x_j)}{\partial x_j} \frac{\partial v_\theta(x_k)}{\partial \theta}$) qui capture les interactions entre les pas de temps. Les méthodes précédentes comme DRTune éliminent ce terme entièrement en arrêtant les gradients à l'entrée du modèle, ce qui rejette des informations de dépendance précieuses.
• La Solution : LeapAlign met à l'échelle le terme de gradient imbriqué par un facteur $\alpha \in [0, 1]$ plutôt que de le supprimer. Cela est réalisé en modifiant l'entrée de la prédiction de vitesse dans le deuxième saut : $\hat{x}_{0|j} = x_j - j v_\theta(\alpha x_j + (1-\alpha)\text{stop\_gradient}(x_j))$.
• Effet : Cela réduit l'amplitude des termes de gradient importants pour stabiliser l'entraînement tout en conservant la structure complète du gradient, y compris le gradient imbriqué.

3. Pondération par Similarité de Trajectoire

Pour garantir que les trajectoires à saut raccourcies fournissent une supervision fiable, la méthode attribue des poids plus élevés aux sauts qui ressemblent étroitement au véritable chemin de génération.

• Mécanisme : La perte est pondérée par l'inverse de la différence absolue moyenne entre les états prédits ($\hat{x}$) et les états réels ($x$) aux points de connexion ($x_j$ et $x_0$).
• Formule : $w_{sim} = \frac{1}{\max(d_j, \tau) + \max(d_0, \tau)}$, où $d$ représente la distance et $\tau$ est un seuil minimum pour éviter une sur-pondération des paires quasi identiques.
• Objectif : Cela assure que les sauts s'écartant significativement de la dynamique réelle contribuent moins à la mise à jour du gradient, améliorant ainsi la stabilité et la fidélité de la supervision.

4. Objectif de Fine-Tuning

La méthode emploie une perte de type « hinge » pour maximiser les récompenses tout en empêchant le piratage de récompense : $L_{raw} = \max(0, \lambda - r(x_0))$. Crucialement, contrairement aux méthodes qui estiment les récompenses à partir des prédictions de saut, LeapAlign évalue la récompense sur l'image finale réelle ($x_0$) générée par la trajectoire complète, garantissant que le modèle de récompense évalue la véritable qualité visuelle et sémantique.

Contributions Clés

1. Cadre LeapAlign : Une méthode de post-entraînement novatrice qui construit des trajectoires à saut en deux étapes pour permettre la rétropropagation directe des gradients vers les étapes de génération précoces avec un coût mémoire constant.
2. Techniques de Stabilité des Gradients : L'introduction du découplage des gradients pour mettre à l'échelle (plutôt que supprimer) les gradients imbriqués, et de la pondération par similarité de trajectoire pour privilégier les sauts cohérents, répondant aux problèmes d'instabilité et de perte de signal des méthodes à gradient direct antérieures.
3. Évaluation Complète : Démonstration que LeapAlign surpasse les méthodes à gradient de politique de l'état de l'art (basées sur GRPO) et les méthodes à gradient direct sur plusieurs métriques, y compris les scores de préférence humaine (HPSv2.1, HPSv3, PickScore) et les benchmarks d'alignement compositionnel (GenEval).

Résultats

Les auteurs ont effectué un fine-tuning du modèle Flux.1-Dev en utilisant LeapAlign et l'ont évalué par rapport à des bases de référence incluant DanceGRPO, MixGRPO, ReFL, DRaFT-LV et DRTune.

• Alignement des Préférences Générales : LeapAlign a obtenu les scores les plus élevés sur les métriques intra-domaine (HPSv2.1) et extra-domaine (HPSv3, PickScore, UnifiedReward). Notamment, il a surpassé MixGRPO (qui utilise trois modèles de récompense) tout en étant entraîné sur un seul modèle de récompense (HPSv2.1).
• Alignement Compositionnel : Sur le benchmark GenEval, LeapAlign a montré des gains significatifs, en particulier dans les tâches « deux objets », « couleurs », « position » et « liaison d'attributs ». L'amélioration dans ces catégories met en évidence le succès de la méthode à optimiser les étapes de génération précoces qui dictent la disposition globale.
• Analyse Qualitative : Les comparaisons visuelles ont démontré que LeapAlign modifie substantiellement les structures d'images globales pour mieux s'aligner avec les invites textuelles, tandis que les méthodes limitées aux mises à jour tardives (comme ReFL) conservaient souvent la disposition du modèle pré-entraîné.
• Études d'Ablation : Les expériences ont confirmé que la conception de saut en deux étapes offre le meilleur compromis entre performance et mémoire. Le découplage des gradients avec $\alpha=0.3$ a produit des résultats optimaux, et la sélection aléatoire des pas de temps sur toute la plage $[0, 1]$ était supérieure aux plages fixes, soulignant l'importance des mises à jour des étapes précoces.

Importance et Revendications

L'article revendique que LeapAlign résout avec succès le compromis entre l'efficacité mémoire et la couverture des gradients dans le post-entraînement par matching de flux. En permettant des mises à jour aux étapes de génération précoces sans coûts computationnels prohibitifs, il permet l'optimisation de la structure globale de l'image, ce qui est crucial pour les tâches compositionnelles complexes. Les auteurs positionnent LeapAlign comme une alternative stable et efficace aux méthodes stochastiques à gradient de politique et aux approches à gradient direct instables, atteignant un alignement image-texte et une qualité de génération supérieurs. Le travail suggère que le maintien et la mise à l'échelle des gradients imbriqués, plutôt que leur élimination, sont un facteur clé pour améliorer les performances de l'alignement à gradient direct. Des travaux futurs sont notés pour étendre ces principes à la génération vidéo.