Token Merging (ToMe)

La fusion de tokens (ToMe) est une méthode qui accélère les transformeurs, notamment le transformeur de vision (ViT), en combinant les tokens similaires dans le réseau afin que chaque bloc en traite moins. Les images comportent beaucoup de redondance, par exemple des patches de ciel ou de route presque identiques ; leur fusion par moyenne réduit les calculs tout en préservant leurs informations. ToMe fonctionne sur des modèles déjà entraînés, sans réentraînement, et peut aussi être appliqué pendant l’entraînement. Daniel Bolya et ses collègues de Georgia Tech et Meta AI l’ont présenté dans Token Merging: Your ViT But Faster, un article présenté oralement à ICLR 2023. Sans adaptation particulière, la méthode a doublé le débit de ViT-L à 512 pixels et de ViT-H à 518 pixels, avec une baisse de précision de seulement 0,2 % à 0,3 %, et accéléré 2,2 fois un modèle vidéo ViT-L avec une baisse similaire.

Fonctionnement de la fusion de tokens#

Un ViT divise une image en patches, chacun représenté par un token, et chaque bloc exécute l’auto-attention, dont le coût augmente avec le carré du nombre de tokens, puis un MLP, dont le coût augmente de façon linéaire. ToMe insère une étape de fusion entre l’attention et le MLP de chaque bloc et supprime un nombre fixe de tokens, r, à chaque étape ; un modèle de 24 blocs avec r = 8 termine donc avec 192 tokens de moins qu’au départ. La fusion utilise une mise en correspondance souple bipartite, choisie pour sa rapidité et l’absence de regroupement itératif :

1. Séparation : Répartir les tokens en alternance dans deux ensembles, A et B.
2. Appariement : Associer chaque token de A au token le plus similaire de B, selon la similarité cosinus des clés d’attention moyennées sur les têtes.
3. Conserver les meilleures arêtes : Garder les r paires les plus similaires.
4. Fusion : Faire la moyenne de chaque paire conservée en un seul token, pondérée par le nombre de patches d’origine que chaque token représente déjà, puis concaténer à nouveau les deux ensembles.

Comme les tokens fusionnés représentent plusieurs patches, ToMe ajoute le logarithme de la taille de chaque token aux scores d’attention ; cette méthode, appelée attention proportionnelle, permet à un token construit à partir de nombreux patches de recevoir sa part d’attention. Toute la procédure est différentiable, ce qui permet d’entraîner les modèles avec ToMe pour gagner encore en rapidité.

Cet exemple PyTorch applique la mise en correspondance souple bipartite au ViT DINOv2-small de Meta, depuis Hugging Face Transformers, sur une image de bus de 448 x 448 pixels ; il fusionne r tokens après chacun des 12 blocs et compare l’intégration d’image obtenue à celle du modèle sans fusion :

from urllib.request import urlopen

import torch
import torch.nn.functional as F
from PIL import Image
from transformers import AutoImageProcessor, AutoModel

processor = AutoImageProcessor.from_pretrained("facebook/dinov2-small")
model = AutoModel.from_pretrained("facebook/dinov2-small")
image = Image.open(urlopen("https://ultralytics.com/images/bus.jpg")).resize((448, 448))
pixels = processor(images=image, do_resize=False, do_center_crop=False, return_tensors="pt").pixel_values


def merge(x, r):
    """Bipartite soft matching: merge the r most similar pairs between alternating token sets A and B."""
    cls, a, b = x[:, :1], x[:, 1::2], x[:, 2::2]
    similarity = F.normalize(a[0], dim=-1) @ F.normalize(b[0], dim=-1).T
    best, match = similarity.max(dim=-1)  # each A token's most similar B token
    order = best.argsort(descending=True)
    merged, kept = order[:r], order[r:]
    index = match[merged][None, :, None].expand(1, r, x.shape[-1])
    b = b.scatter_reduce(1, index, a[:, merged], reduce="mean")  # average merged A tokens into their B matches
    return torch.cat([cls, a[:, kept], b], dim=1)


with torch.no_grad():
    for r in (0, 32, 64):
        x, counts = model.embeddings(pixels), []
        for layer in model.encoder.layer:
            counts.append(x.shape[1])
            x = merge(layer(x), r)  # merge r tokens after every transformer block
        embedding = model.layernorm(x)[:, 0]  # CLS token: the image embedding
        reference = embedding if r == 0 else reference
        similarity = F.cosine_similarity(embedding, reference).item()
        print(f"r={r}: mean tokens per block {sum(counts) / 12:.0f}, CLS similarity to unmerged {similarity:.3f}")

Le modèle sans fusion traite 1025 tokens dans chaque bloc, soit 1 024 patches plus le token CLS. La fusion de 32 tokens par bloc ramène la moyenne à 849, tandis que l’intégration d’image conserve une similarité cosinus de 0.983 avec l’originale ; la fusion de 64 tokens par bloc la ramène à 673 avec une similarité de 0.887, ce qui illustre le compromis entre calcul et fidélité contrôlé par r.

Utilisations de la fusion de tokens#

Token Merging for Fast Stable Diffusion applique la même idée aux blocs Transformer du réseau de débruitage de Stable Diffusion, où elle peut supprimer jusqu’à 60 % des tokens, générer des images jusqu’à deux fois plus vite et utiliser jusqu’à 5,6 fois moins de mémoire, sans entraînement supplémentaire. Au-delà des images, l’article ToMe a doublé le débit d’un ViT-B audio entraîné avec fusion, avec une baisse de mAP de 0,4 %, et accéléré jusqu’à deux fois l’ajustement fin d’un autoencodeur masqué sur vidéo. Les chercheurs appliquent des techniques similaires de fusion et d’élagage aux tokens visuels des modèles vision-langage, où LLaVA-1.5 utilise 576 tokens du contexte du modèle de langage pour chaque image.

Dans le package Ultralytics, le Segment Anything Model (SAM) de Meta encode chaque image de 1024 x 1024 pixels en 4 096 tokens de patch ; dans la comparaison avec MobileSAM, son encodeur d’image ViT-H prend 452 ms sur les 456 ms nécessaires à SAM par image sur GPU, soit la part des calculs ciblée par les méthodes de réduction de tokens.

Fusion de tokens et concepts apparentés#

• Élagage des tokens : Des méthodes comme DynamicViT et EViT suppriment les tokens jugés sans importance, ce qui entraîne une perte d’information et nécessite généralement un entraînement. ToMe fusionne plutôt les tokens, en intégrant leurs informations par moyenne, et fonctionne sur des modèles entraînés sans cette méthode.
• Tokenisation : Crée les tokens initiaux à partir de patches ou de texte avant l’exécution du réseau. ToMe réduit leur nombre à l’intérieur du réseau.
• Attention Flash et attention éparse : Réduisent le coût ou accélèrent l’attention pour un nombre donné de tokens. ToMe réduit le nombre de tokens, ce qui accélère aussi les couches MLP, et peut être combiné à ces méthodes.
• Mélange de profondeurs (MoD) : Permet à un routeur appris de décider quels tokens sautent un bloc dans les modèles de langage. ToMe fusionne définitivement les tokens redondants selon une règle fixe de similarité.
• Résolution d’entrée plus faible : Réduit uniformément le nombre de tokens dans toute l’image. ToMe fusionne surtout les régions redondantes, comme l’arrière-plan, et conserve les régions détaillées à pleine résolution.