Résumé Technique : Tokenisation Protéique Adaptative (APT)

Énoncé du Problème

Les approches actuelles de la tokenisation de la structure des protéines reposent sur la localité, où les tokens sont créés en agrégeant des informations provenant de résidus spatialement voisins le long d'une séquence protéique. Bien que ces méthodes atteignent une haute fidélité de reconstruction, les auteurs identifient deux limitations critiques :

1. Accumulation d'erreurs dans la génération : Les modèles génératifs basés sur une tokenisation locale souffrent de la propagation d'erreurs. Un seul échantillonnage erroné d'un token peut entraîner des déviations significatives dans la structure prédite, ce qui fait que les approches discrètes sont moins performantes que les modèles de diffusion continus sur les métriques génératives.
2. Faible scalabilité : Le nombre de tokens augmente linéairement avec la taille de la protéine. Cela rend la modélisation de complexes protéiques de grande taille coûteuse en termes de calcul et limite la capacité de compression efficace des protéines.
3. Limitations de la représentation : Les tokeniseurs locaux nécessitent souvent des opérations de mean-pooling (moyennage) pour générer des représentations de taille fixe pour les tâches en aval, ce qui peut perdre des informations globales critiques et s'adapter difficilement à l'échelle de la taille des protéines.

Méthodologie

Les auteurs proposent la Tokenisation Protéique Adaptative (APT), une méthode de tokenisation globale où les tokens successifs contribuent à des niveaux de détail croissants vers une représentation globale, plutôt que de correspondre à des voisinages locaux spécifiques.

Architecture

L'APT est implémentée comme un autoencodeur de diffusion avec un goulot d'étranglement discret :

• Entrée : Une séquence de coordonnées $C_\alpha$ brutes (normalisées à masse nulle) de longueur $L$.
• Encodeur : Un transformeur à attention bidirectionnelle qui mappe les positions brutes vers une séquence latente $c \in \mathbb{R}^{L \times d}$.
• Discrétisation : La séquence latente est discrétisée à l'aide de la Quantification Scalaire Finie (FSQ) avec des niveaux $(8, 5, 5, 5)$, résultant en une taille de codebook effective de 1 000.
• Mécanisme d'adaptativité : Pour imposer l'adaptativité, le modèle utilise un dropout imbriqué. Pendant l'entraînement, un seuil supérieur $U$ est échantillonné uniformément dans $[1, \min(L, k_{max})]$. Les tokens au-delà de ce seuil sont supprimés. Cela encourage le modèle à placer les informations globales critiques (composantes de basse fréquence) dans les premiers tokens et les détails de haute fréquence dans les suivants.
• Décodeur : Un décodeur de diffusion, entraîné à l'aide d'un objectif de flow-matching, reconstruit les coordonnées atomiques conditionnées par les tokens discrets. Le décodeur utilise l'équivariance stochastique, apprenant les symétries plutôt que de les imposer explicitement.
• Prédiction de la taille : La taille de la protéine est régressée à partir du premier token via une perte d'entropie croisée, permettant au modèle de découpler la taille de la protéine de la longueur de la séquence de conditionnement.

Entraînement et Inférence

• Entraînement : Le modèle est entraîné sur environ 473 000 prédictions synthétiques AlphaFold2 issues de la base de données AFDB classée par Foldseek. L'entraînement comporte deux étapes : (1) l'entraînement de l'autoencodeur avec des rotations aléatoires pour apprendre l'équivariance stochastique, et (2) l'entraînement d'un modèle autorégressif (de type GPT) sur les tokens APT.
• Stratégies d'inférence :
  • Tail Dropout (Dropout de queue) : Lors de la génération, le modèle peut arrêter la génération prématurément. Tout préfixe de la séquence de tokens constitue un signal de conditionnement valide. La queue (détails de haute fréquence) est déléguée au décodeur de diffusion.
  • Arrêt basé sur l'entropie : Les auteurs introduisent des critères d'arrêt basés sur l'entropie des tokens (seuil fini, entropie minimale basée sur les splines) pour déterminer le point optimal pour interrompre la génération, équilibrant la fidélité de la représentation et la capacité générative.
  • Recuit du classificateur (Classifier Annealing) : Pour atténuer les artefacts causés par l'éloignement de la variété de données réelle, les auteurs utilisent le recuit du classificateur, interpolant entre les champs de guidage conditionnels et inconditionnels lors du décodage.

Contributions Clés

1. Tokeniseur Global Adaptatif : Introduction d'un tokeniseur où les tokens représentent des descripteurs globaux dans une hiérarchie du grossier au fin, plutôt que des voisinages locaux.
2. Validation de la Performance : Démonstration que le modèle égale ou surpasse les tokeniseurs locaux de pointe (ex: DPLM2, ESM3, Kanzi) sur les tâches génératives, d'apprentissage de représentation et de reconstruction.
3. Techniques d'Inférence : Développement de l'échantillonnage basé sur l'entropie et du recuit du classificateur pour gérer la complexité de l'échantillon et atténuer l'exposition aux erreurs, améliorant ainsi la designabilité.
4. Applications Zero-Shot : Démonstration du découplage de la taille des protéines du séquençage de conditionnement, permettant des applications zero-shot telles que le rétrécissement de protéines (réduction du nombre de résidus tout en préservant la structure) et la maturation d'affinité via le passage à l'échelle lors de l'inférence.

Résultats

• Reconstruction : Sur les ensembles de test CATH, CAMEO et AFDB, l'APT atteint des RMSD et des scores TM comparables aux modèles de diffusion continus de pointe. Notamment, l'utilisation de seulement 32 à 64 tokens produit des RMSD $< 2$ Å, affirmant la compressibilité des structures protéiques.
• Génération : Un modèle autorégressif entraîné sur les tokens APT atteint une designabilité de 0,871 (fraction d'échantillons avec scRMSD $< 2$ Å), surpassant de manière significative les modèles de diffusion discrète comme DPLM2 (0,486) et ESM3. Ceci est réalisé sans techniques de "best-of-N" sampling.
• Apprentissage de Représentation : Sur la tâche de classification CATH, le sondage non linéaire (MLP) sur les tokens APT surpasse le sondage des représentations d'ESM3 et DPLM2. Crucialement, l'APT fournit des représentations globales de taille fixe sans mean-pooling, et même des représentations hautement compressées (16 tokens) surpassent des modèles plus larges.
• Applications :
  • Rétrécissement de Protéines : Génération réussie de versions plus petites de protéines (ex: hémoglobine, tonneau bêta) en préservant la structure globale et locale en conditionnant le processus de diffusion sur moins de résidus.
  • Maturation d'Affinité : Capacité démontrée à faire évoluer des liants faibles en liants forts via une recherche par faisceau (beam search) guidée par des fonctions de récompense (ex: iPAE) dans l'espace latent.

Signification

L'article affirme que la Tokenisation Protéique Adaptative offre une approche sensible aux tâches pour passer à l'échelle des tâches biologiques. En passant d'une tokenisation locale à une tokenisation globale, la méthode résout le compromis entre la fidélité de reconstruction et la capacité générative. La capacité de compresser les protéines en vecteurs de taille fixe sans perte d'information via le mean-pooling offre une voie pratique pour modéliser de grands complexes protéiques et effectuer une mise à l'échelle lors de l'inférence. Les auteurs positionnent ce travail comme une étape vers des modèles génératifs capables de raisonner à travers plusieurs échelles de longueur, répondant à un défi clé de la bio-ingénierie de pointe.

Note : Les auteurs reconnaissent des limitations, spécifiquement que les tokens APT sont efficaces pour les tâches globales mais moins adaptés aux tâches locales comme le scaffolding de motifs, et que le tokeniseur actuel ne prend pas en compte les chaînes latérales ou la composition en acides aminés, qui sont critiques pour certaines fonctions.