Résumé Technique : Attention Sinkhorn Différentiable par Blocs

1. Énoncé du Problème

L'article aborde le défi de l'implémentation de mécanismes d'attention différentiables pour contextes longs basés sur le Transport Optimal Entropique (OT) sur du matériel d'accélération (spécifiquement les TPUs). Bien que l'OT fournisse une relaxation lisse et différentiable de l'alignement de séquences discret (contrairement aux méthodes de programmation dynamique classiques comme Smith–Waterman), il se heurte à deux goulots d'étranglement critiques dans les contextes longs :

1. Interdits Mémoires : L'attention OT dense matérialise des matrices de scores ou de plans de transport de taille $L \times L$, ce qui est irréalisable pour de grandes longueurs de séquence $L$.
2. Complexité du Passage Arrière : La différenciation automatique exacte (autodiff) à travers de nombreuses itérations de Sinkhorn est coûteuse en calcul et gourmande en mémoire. À l'inverse, la différenciation implicite nécessite souvent de résoudre de grands systèmes linéaires instables.

L'objectif central est de construire un mécanisme d'attention qui soit explicitement différentiable, exact pour le surrogate spécifique utilisé, et compatible avec les noyaux d'accélérateurs par blocs pour éviter les coûts de mémoire et de calcul en $O(L^2)$.

2. Méthodologie

La méthode proposée utilise un surrogate de raffinement de queue à profondeur fixe, avec base arrêtée. L'approche découple le passage avant en une résolution de base non différentiable et une courte queue de raffinement différentiable.

2.1 Passage Avant : Sinkhorn Équilibré par Blocs

Le modèle opère sur un support actif fixe $\Omega$ avec une largeur de bande $W$. Il intègre la température entropique $\epsilon$ dans la matrice de scores $S$. Le passage avant consiste en :

• Résolution de Base Arrêtée : Une itération Sinkhorn équilibrée de $T$ étapes est effectuée jusqu'à convergence (ou un nombre d'étapes fixe) pour produire les potentiels duaux $(u^{(0)}, v^{(0)})$. Cette étape utilise une opération de stop-gradient.
• Queue de Raffinement : Une courte séquence différentiable de $R$ étapes (où $R$ est petit, par exemple $R=2$) affine les potentiels en partant de l'état de base arrêté.
• Complexité : Pour une largeur de bande fixe $W$, une dimension de tête $d$ et une profondeur de queue $R$, le coût est de $O((T+R)LW)$ en temps, $O(Ld)$ en stockage d'entrée, et $O(L)$ en utilisation supplémentaire de mémoire à large bande (HBM).

2.2 L'Escalier $R=2$ et le Calendrier à Une Tuile de Référence

L'innovation algorithmique clé est le passage arrière exact pour la queue de raffinement.

• Le Défi : Une implémentation directe de l'adjoint exact pour $R=2$ nécessite d'évaluer quatre facteurs de plan "en escalier" distincts : $P^{(2,2)}, P^{(2,1)}, P^{(1,1)}, P^{(1,0)}$.
• La Solution (Calcul de l'Orbite de Transport) : Les auteurs prouvent que ces quatre facteurs résident dans une seule orbite de mise à l'échelle de lignes/colonnes. Spécifiquement, tout plan $P^{(a,b)}$ peut être reconstruit à partir d'un unique plan de référence $P^{(2,2)}$ par multiplication élémentaire avec les différences exponentielles des potentiels duaux :
  $$P^{(a,b)} = P^{(2,2)} \odot \exp(u^{(a)} - u^{(2)}) \odot \exp(v^{(b)} - v^{(2)})$$
• Implication : Le passage arrière n'a pas besoin de matérialiser trois tuiles de plan séparées. Il utilise une seule tuile de référence résidente ($P^{(2,2)}$) et applique des modificateurs vectoriels peu coûteux (re-mise à l'échelle) pour calculer les gradients des autres termes. Cela réduit le passage arrière à une opération par blocs avec une arithmétique de tuiles en $O(LW)$ et un état vectoriel global en $O(L)$.

2.3 Pont de Poubelle Conscient des Espaces Vides

Pour gérer les espaces vides dans l'alignement de séquences (une exigence courante en production), l'article formalise le mécanisme de "poubelle" (dustbin) non pas comme un nouveau solveur OT, mais comme un espace d'état augmenté.

• La méthode ajoute un token poubelle appris et des tokens de remplissage inactifs aux tenseurs requête/clé/valeur.
• Théorème : Le chemin augmenté par la poubelle est mathématiquement équivalent au même surrogate de raffinement de queue à support fixe équilibré appliqué à cet espace d'état plus large.
• Résultat : L'adjoint exact $R=2$ et le calendrier à une tuile de référence se transposent directement au chemin de transport conscient des espaces vides sans nécessiter de nouvelles machineries théoriques.

2.4 Garanties Théoriques

L'article fournit trois certificats théoriques :

1. Limite de Biais de Surrogate Local : Prouve que la différence entre le gradient complet (à travers la résolution de base) et le surrogate à stop-gradient décroît géométriquement avec la profondeur de queue $R$, à condition que l'application de raffinement soit localement contractante.
2. Certificat de Contraction Projective : Utilise la métrique projective de Hilbert pour montrer que les blocs actifs strictement positifs de l'application de Sinkhorn sont des contractions, assurant la stabilité.
3. Certificat de Biais A Posteriori : Une méthode opérationnelle pour mesurer le reste du gradient ($\eta_R$) et certifier que l'omission du stop-gradient est inférieure à une tolérance définie par l'utilisateur $\tau$. Cela permet la sélection dynamique de la profondeur de queue optimale $R$.

3. Contributions Clés

1. Surrogate Explicite & Adjoint Exact : Un surrogate de raffinement de queue à support fixe équilibré explicite avec un adjoint exact $R=2$ évitant l'autodiff complète à travers la résolution de base.
2. Calendrier à Une Tuile de Référence : Un théorème de planification arrière montrant que le gradient $R=2$ peut être calculé en utilisant une seule tuile de transport de référence plus des modificateurs vectoriels, permettant une implémentation efficace de noyaux TPU.
3. Pont d'Augmentation de Poubelle : Un théorème démontrant que le chemin "poubelle" de production est algébriquement identique à la factorisation principale sur un espace augmenté, unifiant le transport conscient des espaces vides avec la théorie centrale.
4. Validation Empirique : Un pipeline complet allant de la validation d'exactitude sur des données synthétiques à un entraînement de plusieurs heures sur TPU v6e-8.

4. Résultats

4.1 Validation d'Exactitude

Sur des problèmes masqués synthétiques, le noyau optimisé $R=2$ correspond à l'autodiff exacte du surrogate centré avec une erreur relative de $10^{-5}$ à $10^{-10}$.

• Décroissance du Biais : La comparaison des gradients de raffinement de queue avec la rétropropagation complète dans le temps (BPTT) montre que l'écart de gradient diminue significativement de $R=0$ à $R=1$ et $R=2$. À $R=2$, l'erreur est pratiquement équivalente à celle de $R=4$ pour les configurations testées.
• Performance du Certificat : Le certificat de biais a posteriori reconstruit correctement le reste du gradient complet avec une précision numérique et sélectionne avec succès $R=2$ comme profondeur minimale satisfaisant une tolérance de $10^{-5}$.

4.2 Entraînement TPU v6e-8

La méthode a été validée sur une tâche de modélisation de séquences protéiques Pfam :

• Dépistage : Quatre configurations (équilibré $R=2$, équilibré $R=4$, et deux variantes poubelle) ont terminé un entraînement de bout en bout. La configuration $R=2$ a atteint un débit de ~8,5 exemples/seconde (après mise en chauffe) sans différence de perte significative par rapport à $R=4$.
• Exécution à Budget Long : Un entraînement promu équilibré $R=2$ a maintenu l'entraînement pendant trois heures, atteignant l'étape 1437.
• Métriques de Performance : Par rapport à l'initialisation (Étape 0), le point de contrôle final a montré :
  • Erreur de Reconstruction : Améliorée de 3,17 à 0,99.
  • Entropie Croisée Sparsifiée : Améliorée de 5,86 à 5,69.
  • Stabilité : L'exécution est restée stable sans signe de divergence, démontrant l'entraînabilité du chemin soutenu par la factorisation.

5. Signification et Revendications

L'article revendique d'établir un chemin pratique, exact et compatible avec le matériel pour l'attention OT à contexte long.

• Signification Théorique : Il fournit un "pont conscient des espaces vides correspondant au théorème", prouvant que des fonctionnalités de production complexes (poubelles) ne nécessitent pas d'abandonner la factorisation équilibrée centrale. Il offre une justification rigoureuse de la profondeur $R=2$ via des certificats de contraction.
• Signification Pratique : Le travail démontre qu'une théorie de passage arrière à profondeur fixe exacte peut être mappée sur des noyaux TPU, atteignant un débit élevé et une stabilité sur des données protéiques réelles.
• Modestie des Revendications : Les auteurs déclarent explicitement qu'il s'agit d'une contribution méthodologique et systémique, et non d'une revendication de domination en temps réel sur tous les autres mécanismes d'attention (par exemple, FlashSinkhorn ou les transformeurs épars) sans benchmarks supplémentaires. Les résultats sont présentés comme une preuve de l'entraînabilité et de la stabilité systémique pour le chemin de production spécifique, plutôt que comme une preuve définitive de supériorité dans toutes les tâches de prise de décision en aval. Le travail limite délibérément son champ d'application à l'OT équilibré à support fixe, laissant le transport déséquilibré et la géométrie de support différentiable pour des travaux futurs.