Token Merging (ToMe)

A fusão de tokens (ToMe) é um método para acelerar transformers, especialmente o Vision Transformer (ViT), combinando tokens semelhantes dentro da rede, para que cada bloco processe menos tokens. As imagens contêm muita redundância, como patches de céu ou estrada quase iguais, e calcular a média desses tokens remove operações computacionais sem perder as informações que contêm. O ToMe funciona em modelos já treinados, sem novo treinamento, e também pode ser aplicado durante o treinamento. Daniel Bolya e colegas da Georgia Tech e da Meta AI o apresentaram em Token Merging: Your ViT But Faster, artigo apresentado oralmente na ICLR 2023. Sem ajustes adicionais, o método duplicou a taxa de processamento do ViT-L a 512 pixels e do ViT-H a 518 pixels, com uma queda de precisão de apenas 0,2% a 0,3%, e acelerou em 2,2 vezes um modelo de vídeo ViT-L, com queda semelhante.

Como funciona a fusão de tokens#

Um ViT divide uma imagem em patches, com um token por patch, e cada bloco executa autoatenção, cujo custo cresce com o quadrado do número de tokens, seguida por uma MLP, cujo custo cresce linearmente. O ToMe insere uma etapa de fusão entre a atenção e a MLP de cada bloco e remove um número fixo de tokens, r, a cada etapa; assim, um modelo de 24 blocos com r = 8 termina com 192 tokens a menos do que no início. A fusão usa correspondência suave bipartida, escolhida por ser rápida e não exigir agrupamento iterativo:

1. Divisão: Divide os tokens alternadamente em dois conjuntos, A e B.
2. Correspondência: Liga cada token de A ao token mais semelhante de B, medido pela similaridade do cosseno das chaves de atenção, cuja média é calculada entre as cabeças.
3. Manter as melhores arestas: Mantém os r pares mais semelhantes.
4. Fusão: Calcula a média de cada par mantido para formar um token, ponderando pelo número de patches originais que cada token já representa, e concatena novamente os dois conjuntos.

Como os tokens fundidos representam vários patches, o ToMe adiciona o logaritmo do tamanho de cada token às pontuações de atenção. Essa técnica, chamada atenção proporcional, garante que um token formado por muitos patches receba a sua parcela de atenção. Todo o procedimento é diferenciável, o que permite treinar modelos com ToMe para obter velocidade adicional.

Este exemplo de PyTorch aplica correspondência suave bipartida ao ViT DINOv2-small da Meta, disponível no Transformers da Hugging Face, usando uma imagem de ônibus de 448 x 448 pixels. O exemplo funde r tokens após cada um dos 12 blocos e compara a incorporação de imagem resultante com a do modelo sem fusão:

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

O modelo sem fusão processa 1025 tokens em cada bloco: 1.024 patches mais o token CLS. A fusão de 32 tokens por bloco reduz a média para 849, enquanto a incorporação da imagem mantém uma similaridade do cosseno de 0.983 em relação à original; a fusão de 64 tokens por bloco reduz a média para 673, com similaridade de 0.887. Isso mostra como r representa uma troca entre custo computacional e fidelidade.

Onde a fusão de tokens é usada#

Token Merging for Fast Stable Diffusion aplica a mesma ideia aos blocos Transformer da rede de remoção de ruído do Stable Diffusion, onde pode remover até 60% dos tokens e gerar imagens até duas vezes mais rápido, usando até 5,6 vezes menos memória e sem treinamento adicional. Além das imagens, o artigo sobre ToMe duplicou a taxa de processamento de um ViT-B de áudio treinado com fusão, com uma queda de 0,4% em mAP, e acelerou em até 2 vezes o ajuste fino de autoencoders mascarados em vídeo. Pesquisadores aplicam técnicas semelhantes de fusão e poda aos tokens visuais de modelos de visão e linguagem, nos quais o LLaVA-1.5 usa 576 tokens do contexto do modelo de linguagem para cada imagem.

No pacote Ultralytics, o Segment Anything Model (SAM) da Meta codifica cada imagem de 1024 x 1024 pixels como 4.096 tokens de patch. Na comparação com o MobileSAM, o codificador de imagens ViT-H leva 452 ms dos 456 ms que o SAM precisa por imagem em uma GPU — a parcela do processamento que os métodos de redução de tokens visam.

Fusão de tokens e conceitos relacionados#

• Poda de tokens: Métodos como DynamicViT e EViT descartam tokens considerados sem importância, perdendo as informações que contêm e normalmente exigindo treinamento. O ToMe, em vez disso, funde os tokens, incorporando as informações pela média, e funciona em modelos que não foram treinados com ele.
• Tokenização: Cria os tokens iniciais a partir de patches ou texto, antes da execução da rede. O ToMe reduz o número de tokens dentro da rede.
• Atenção Flash e atenção esparsa: Tornam a atenção mais barata ou rápida para um determinado número de tokens. O ToMe reduz o número de tokens, acelerando também as camadas MLP, e pode ser combinado com esses métodos.
• Mixture of Depths (MoD): Permite que um roteador aprendido decida quais tokens ignoram um bloco em modelos de linguagem. O ToMe funde permanentemente tokens redundantes usando uma regra fixa de similaridade.
• Resolução de entrada menor: Reduz uniformemente o número de tokens em toda a imagem. O ToMe funde tokens principalmente em regiões redundantes, como o fundo, e mantém as regiões detalhadas na resolução original.