Token Merging (ToMe)

La fusion de tokens (ToMe) est une technique de pointe conçue pour optimiser les performances et l'efficacité des architectures Transformer en réduisant le nombre de tokens traités pendant les passes avant. Initialement développée pour accélérer les modèles Transformer de vision (ViT), ToMe fonctionne en identifiant et en combinant systématiquement les tokens redondants au sein du réseau, sans nécessiter d'entraînement supplémentaire. Comme la complexité de calcul du mécanisme d'auto-attention évolue de façon quadratique avec le nombre de tokens, la fusion de tokens similaires réduit considérablement le nombre total d'opérations en virgule flottante (FLOPs), ce qui permet une inférence en temps réel nettement plus rapide.

Comprendre le processus de fusion de tokens#

ToMe se distingue fondamentalement de la tokenisation, qui est l'étape initiale de prétraitement consistant à décomposer une image ou un texte en tokens individuels. Alors que la tokenisation crée les éléments discrets, la fusion de tokens agit comme un mécanisme de sous-échantillonnage dynamique pendant l'exécution de la passe avant du modèle.

L'algorithme utilise généralement un appariement biparti pour évaluer la similarité entre les tokens, en calculant souvent la similarité cosinus entre les clés des tokens dans les couches d'attention. Les tokens qui partagent des informations visuelles ou sémantiques très similaires sont fusionnés, souvent en faisant la moyenne de leurs caractéristiques. Cela garantit la préservation des informations spatiales ou contextuelles essentielles tout en éliminant la charge de calcul superflue, ce qui permet à des frameworks comme PyTorch de traiter beaucoup plus rapidement des modèles de vision complexes.

Applications concrètes de la fusion de tokens#

La fusion de tokens est devenue une stratégie d'optimisation essentielle pour déployer des architectures lourdes fondées sur l'attention dans des environnements aux ressources de calcul limitées.

1. IA générative et synthèse d'images : dans les modèles de diffusion texte-image populaires, ToMe est fréquemment utilisée pour accélérer la génération d'images. En fusionnant les tokens d'arrière-plan ou de faible niveau de détail, le processus de génération nécessite moins d'étapes, ce qui permet d'économiser d'importantes ressources GPU et de réduire la latence pour les utilisateurs finaux qui s'appuient sur des modèles génératifs. Tu peux en apprendre davantage sur les processus de diffusion dans les travaux de recherche fondamentaux disponibles sur arXiv.

2. Déploiements d'IA en périphérie : le déploiement de modèles massifs comme le modèle Segment Anything (SAM) sur des appareils mobiles est notoirement difficile en raison des contraintes de mémoire. ToMe contribue à réduire dynamiquement l'empreinte mémoire, ce qui permet d'exécuter des tâches complexes de segmentation d'images sur du matériel en périphérie. Lorsque la vitesse pure est essentielle, les ingénieurs se tournent souvent vers des architectures natives optimisées et dépourvues d'attention, comme Ultralytics YOLO26, pour obtenir une inférence en périphérie plus rapide de bout en bout.

Exemple Python : calcul de la similarité entre tokens#

L'intégration de ToMe dans une architecture complète nécessite de modifier les blocs d'attention, mais le concept central repose sur la recherche de tokens similaires. L'extrait PyTorch suivant montre comment calculer la similarité cosinus entre un ensemble de tokens afin d'identifier ceux qui peuvent être fusionnés.

import torch
import torch.nn.functional as F

# Simulate a batch of 4 image patches (tokens) with 64-dimensional features
tokens = torch.randn(1, 4, 64)

# Normalize the tokens to easily compute cosine similarity via dot product
normalized_tokens = F.normalize(tokens, p=2, dim=-1)

# Compute the similarity matrix between all tokens (1 x 4 x 4)
similarity_matrix = torch.matmul(normalized_tokens, normalized_tokens.transpose(1, 2))

# Tokens with high similarity scores (close to 1.0) off the diagonal
# are prime candidates for Token Merging.
print("Similarity Matrix:", similarity_matrix)

Les pipelines modernes de machine learning nécessitent un équilibre soigneux entre précision et vitesse. Que tu utilises la fusion de tokens pour optimiser un ViT personnalisé ou que tu t'appuies sur les performances de pointe de YOLO26, la gestion de ces workflows de données complexes est grandement simplifiée par l'Ultralytics Platform. La Platform fournit un écosystème intuitif pour l'annotation de données automatisée, l'entraînement fluide dans le cloud et le déploiement robuste de modèles sur divers environnements matériels de calcul en périphérie. Les organisations qui développent leurs initiatives de vision par ordinateur s'appuient sur ces outils pour mettre en production de manière fiable et efficace des modèles de pointe.