Token Merging (ToMe)

Il merging dei token (ToMe) è un metodo per accelerare i Transformer, in particolare il Vision Transformer (ViT), combinando token simili all'interno della rete, così che ogni blocco ne elabori un numero minore. Le immagini contengono molte ridondanze, come patch di cielo o strada quasi identiche; fare la media di questi token elimina calcoli preservandone le informazioni. ToMe funziona su modelli già addestrati senza riaddestrarli e può essere applicato anche durante l'addestramento. Daniel Bolya e colleghi di Georgia Tech e Meta AI lo hanno presentato in Token Merging: Your ViT But Faster, un articolo presentato oralmente a ICLR 2023. Pronto all'uso, ha raddoppiato il throughput di ViT-L a 512 pixel e di ViT-H a 518 pixel, con immagini e un calo di accuratezza di appena lo 0,2%-0,3%; inoltre, ha accelerato di 2,2 volte un modello video ViT-L con un calo simile.

Come funziona il merging dei token#

Un ViT divide un'immagine in patch, assegnando un token a ciascuna, e ogni blocco esegue la self-attention, il cui costo cresce con il quadrato del numero di token, seguita da una MLP, il cui costo cresce linearmente. ToMe inserisce un passaggio di merging tra l'attenzione e la MLP di ogni blocco e rimuove ogni volta un numero fisso di token, r; di conseguenza, un modello con 24 blocchi e r = 8 termina con 192 token in meno rispetto al numero iniziale. Il merging usa un abbinamento soft bipartito, scelto perché è veloce e non richiede clustering iterativo:

1. Divisione: Dividi i token alternandoli tra due insiemi, A e B.
2. Abbinamento: Collega ogni token di A al token più simile di B, misurando la similarità coseno delle chiavi di attenzione mediate tra le teste.
3. Selezione dei collegamenti migliori: Mantieni le r coppie più simili.
4. Merging: Calcola la media di ogni coppia selezionata in un unico token, ponderandola in base al numero di patch originali già rappresentate da ciascun token, quindi concatena nuovamente i due insiemi.

Poiché i token uniti rappresentano più patch, ToMe aggiunge ai punteggi di attenzione il logaritmo della dimensione di ciascun token, un metodo chiamato attenzione proporzionale, così che un token composto da molte patch riceva la quota di attenzione che gli spetta. L'intera procedura è differenziabile, consentendo di addestrare i modelli con ToMe per ottenere ulteriore velocità.

Questo esempio PyTorch applica l'abbinamento soft bipartito al ViT DINOv2-small di Meta, disponibile in Transformers di Hugging Face, su un'immagine di un autobus da 448 x 448 pixel; esegue il merging di r token dopo ciascuno dei suoi 12 blocchi e confronta l'embedding dell'immagine risultante con quello del modello senza merging:

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}")

Il modello senza merging elabora 1025 token in ogni blocco: 1.024 patch più il token CLS. Unire 32 token per blocco riduce la media a 849, mentre l'embedding dell'immagine mantiene una similarità coseno di 0.983 rispetto all'originale; unirne 64 per blocco la riduce a 673 con una similarità di 0.887, mostrando come r bilanci il costo computazionale e la fedeltà.

Dove si usa il merging dei token#

Token Merging for Fast Stable Diffusion applica la stessa idea ai blocchi Transformer della rete di denoising di Stable Diffusion, dove può rimuovere fino al 60% dei token, generando immagini fino a due volte più velocemente e con un consumo di memoria fino a 5,6 volte inferiore, senza ulteriore addestramento. Oltre alle immagini, l'articolo su ToMe ha raddoppiato il throughput di un ViT-B audio addestrato con merging, con un calo di mAP dello 0,4%, e ha accelerato fino a 2 volte il fine-tuning di un autoencoder mascherato su video. I ricercatori applicano tecniche analoghe di merging e pruning ai token visivi dei modelli visione-linguaggio, nei quali LLaVA-1.5 occupa 576 token del contesto del modello linguistico per ogni immagine.

Nel pacchetto Ultralytics, il Segment Anything Model (SAM) di Meta codifica ogni immagine da 1024 x 1024 pixel come 4.096 token di patch; nel confronto con MobileSAM, il suo encoder d'immagine ViT-H impiega 452 ms dei 456 ms necessari a SAM per immagine su GPU: è proprio questa quota di calcolo che i metodi di riduzione dei token mirano a ridurre.

Merging dei token e concetti correlati#

• Pruning dei token: Metodi come DynamicViT ed EViT eliminano i token giudicati poco importanti, perdendone le informazioni e richiedendo di solito l'addestramento. ToMe invece unisce i token, integrandone le informazioni tramite la media, e funziona con modelli addestrati senza ToMe.
• Tokenizzazione: Crea i token iniziali da patch o testo prima dell'esecuzione della rete. ToMe ne riduce il numero all'interno della rete.
• Attenzione Flash e attenzione sparsa: Rendono l'attenzione meno costosa o più veloce a parità di token. ToMe riduce il numero di token, accelerando anche i livelli MLP, e può essere combinato con questi metodi.
• Mixture of Depths (MoD): Permette a un router appreso di decidere quali token saltano un blocco nei modelli linguistici. ToMe fonde in modo permanente i token ridondanti usando una regola fissa di similarità.
• Risoluzione di input inferiore: Riduce uniformemente il numero di token nell'immagine. ToMe esegue il merging soprattutto nelle aree ridondanti, come lo sfondo, mantenendo la risoluzione massima nelle aree dettagliate.