Token Merging (ToMe)

Token birleştirme (ToMe), özellikle Vision Transformer'ı (ViT) hızlandırmak için Transformer'ların ağ içindeki benzer tokenları birleştirerek her bloğun daha az token işlemesini sağlayan bir yöntemdir. Görüntüler, neredeyse aynı görünen gökyüzü veya yol parçaları gibi çok fazla tekrar içerir; bu tokenların ortalamasını almak, bilgilerini korurken hesaplama yükünü azaltır. ToMe, yeniden eğitim gerektirmeden önceden eğitilmiş modellerde çalışır ve eğitim sırasında da uygulanabilir. Georgia Tech ve Meta AI'dan Daniel Bolya ve çalışma arkadaşları yöntemi, ICLR 2023'te sözlü sunulan Token Merging: Your ViT But Faster makalesinde tanıttı. Hazır kullanıldığında, doğrulukta yalnızca %0,2 ila %0,3 düşüşle 512 pikselde ViT-L'nin ve 518 pikselde ViT-H'nin verimini ikiye katladı; benzer bir düşüşle ViT-L video modelini 2,2 kat hızlandırdı.

Token Birleştirme Nasıl Çalışır#

ViT, bir görüntüyü her biri tek bir token olan parçalara böler ve her blokta, maliyeti token sayısının karesiyle artan öz dikkat ile maliyeti doğrusal olarak artan bir MLP çalıştırır. ToMe, her bloğun dikkat ve MLP katmanları arasına bir birleştirme adımı ekler ve her seferinde sabit sayıda (r) tokenı kaldırır; böylece r = 8 olan 24 bloklu model, başladığından 192 daha az tokenla sona erer. Birleştirme, hızlı olduğu ve yinelemeli kümeleme gerektirmediği için seçilen iki parçalı yumuşak eşleştirmeyi kullanır:

1. Böl: Tokenları dönüşümlü olarak A ve B adlı iki kümeye ayır.
2. Eşleştir: A kümesindeki her tokenı, başlar üzerinden ortalaması alınan dikkat anahtarlarının kosinüs benzerliğine göre B kümesindeki en benzer tokene bağla.
3. En iyi kenarları koru: Benzerliği en yüksek r çifti tut.
4. Birleştir: Korunan her çifti, her tokenın temsil ettiği özgün parça sayısına göre ağırlıklandırıp tek bir tokenda ortala ve iki kümeyi yeniden birleştir.

Birleştirilen tokenlar birden fazla parçayı temsil ettiğinden ToMe, orantılı dikkat adı verilen yöntemde her tokenın boyutunun logaritmasını dikkat puanlarına ekler; böylece çok sayıda parçadan oluşturulan token, dikkat payını alır. Tüm süreç türevlenebilir olduğundan modeller ek hız kazanmak için ToMe ile eğitilebilir.

Bu PyTorch örneği, Meta'nın Hugging Face Transformers kütüphanesindeki DINOv2-small ViT modeline ait iki parçalı yumuşak eşleştirmeyi 448 x 448 boyutundaki bir otobüs görüntüsüne uygular; 12 bloğun her birinden sonra r tokenı birleştirir ve ortaya çıkan görüntü gömmesini, birleştirme yapılmamış modelin gömmesiyle karşılaştırır:

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

Birleştirme yapılmayan model her blokta 1025 token işler; bunlar 1.024 parça ve CLS tokenıdır. Her blokta 32 tokenı birleştirmek ortalama sayıyı 849 düzeyine indirirken görüntü gömmesinin özgün gömmeyle kosinüs benzerliği 0.983 olur; her blokta 64 tokenı birleştirmek ise bu sayıyı 673 düzeyine, benzerliği 0.887 düzeyine indirir ve r değerinin hesaplama yükü ile aslına uygunluk arasında nasıl bir ödünleşim yarattığını gösterir.

Token Birleştirme Nerelerde Kullanılır#

Hızlı Stable Diffusion için Token Birleştirme, aynı fikri Stable Diffusion gürültü giderme ağının Transformer bloklarına uygular; ek eğitim gerektirmeden tokenların %60'a kadarını kaldırabilir, görüntüleri iki kata kadar hızlı ve 5,6 kata kadar daha az bellek kullanarak üretebilir. ToMe makalesi, görüntülerin yanı sıra birleştirmeyle eğitilmiş bir ses ViT-B modelinin verimini mAP'de %0,4 düşüşle ikiye katladı ve videoda maskeli otomatik kodlayıcı ince ayarını iki kata kadar hızlandırdı. Araştırmacılar, LLaVA-1.5'in her görüntü için dil modelinin bağlamından 576 token harcadığı görsel-dil modellerinin görsel tokenlarına da benzer birleştirme ve budama yöntemleri uygular.

Ultralytics paketinde Meta'nın Segment Anything Model (SAM) modeli her görüntüyü 1024 x 1024 piksel çözünürlükte 4.096 parça tokenı olarak kodlar; MobileSAM karşılaştırmasında ise ViT-H görüntü kodlayıcısı, SAM'in GPU'da görüntü başına ihtiyaç duyduğu 456 ms'nin 452 ms'sini alır. Token azaltma yöntemleri hesaplamanın bu bölümünü hedefler.

Token Birleştirme ve İlgili Kavramlar#

• Token budama: DynamicViT ve EViT gibi yöntemler önemsiz olduğuna karar verilen tokenları atar; bu da bilgiyi kaybettirir ve genellikle eğitim gerektirir. ToMe bunun yerine tokenları birleştirerek bilgiyi ortalamaya dahil eder ve bu yöntem olmadan eğitilmiş modellerde de çalışır.
• Tokenleştirme: Ağ çalışmadan önce parçalardan veya metinden başlangıç tokenlarını oluşturur. ToMe, ağ içindeki token sayısını azaltır.
• Flash attention ve seyrek dikkat: Belirli sayıda token için dikkati daha ucuz veya hızlı hâle getirir. ToMe token sayısını azaltarak MLP katmanlarını da hızlandırır ve bu yöntemlerle birlikte kullanılabilir.
• Derinlik Karışımı (MoD): Dil modellerinde öğrenilmiş bir yönlendiricinin hangi tokenların bir bloğu atlayacağına karar vermesini sağlar. ToMe ise gereksiz tokenları sabit bir benzerlik kuralıyla kalıcı olarak birleştirir.
• Daha düşük giriş çözünürlüğü: Görüntünün tamamındaki token sayısını aynı oranda azaltır. ToMe, arka plan gibi çoğunlukla tekrar içeren bölgeleri birleştirir ve ayrıntılı bölgeleri tam çözünürlükte tutar.