Token Merging (ToMe)

La fusión de tokens (ToMe) es un método para acelerar los transformers, en especial el Transformer de visión (ViT), que combina tokens similares dentro de la red para que cada bloque procese menos. Las imágenes contienen mucha redundancia, como parches de cielo o carretera casi idénticos, y promediar esos tokens reduce los cálculos sin perder su información. ToMe funciona con modelos ya entrenados, sin volver a entrenarlos, y también se puede aplicar durante el entrenamiento. Daniel Bolya y sus colegas de Georgia Tech y Meta AI lo presentaron en Token Merging: Your ViT But Faster, un artículo presentado oralmente en ICLR 2023. Sin modificaciones, duplicó el rendimiento de ViT-L a 512 píxeles y de ViT-H a 518 píxeles en imágenes, con una caída de precisión de solo el 0,2 % al 0,3 %, y aceleró 2,2 veces un modelo de vídeo ViT-L con una caída similar.

Cómo funciona la fusión de tokens#

Un ViT divide una imagen en parches, uno por token, y cada bloque ejecuta autoatención, cuyo coste crece con el cuadrado del número de tokens, seguida de un MLP, cuyo coste crece linealmente. ToMe inserta un paso de fusión entre la atención y el MLP de cada bloque y elimina un número fijo de tokens, r, cada vez; así, un modelo de 24 bloques con r = 8 termina con 192 tokens menos que al principio. La fusión utiliza emparejamiento suave bipartito, elegido por su rapidez y porque no requiere agrupamiento iterativo:

1. Dividir: Separa los tokens alternativamente en dos conjuntos, A y B.
2. Emparejar: Conecta cada token de A con el token más similar de B, medido mediante la similitud coseno de las claves de atención promediadas entre cabezas.
3. Conservar las mejores aristas: Conserva los r pares más similares.
4. Fusionar: Promedia cada par conservado en un solo token, ponderado según el número de parches originales que ya representa cada token, y vuelve a concatenar los dos conjuntos.

Como los tokens fusionados representan varios parches, ToMe añade el logaritmo del tamaño de cada token a las puntuaciones de atención, una técnica llamada atención proporcional, para que los tokens formados por muchos parches reciban la parte de atención que les corresponde. Todo el procedimiento es diferenciable, lo que permite entrenar modelos con ToMe para obtener más velocidad.

Este ejemplo de PyTorch aplica el emparejamiento suave bipartito al ViT pequeño DINOv2 de Meta, disponible en Transformers de Hugging Face, sobre una imagen de 448 x 448 de un autobús; fusiona r tokens después de cada uno de sus 12 bloques y compara la representación resultante de la imagen con la del modelo sin fusiones:

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

El modelo sin fusiones procesa 1025 tokens en cada bloque: 1.024 parches más el token CLS. Fusionar 32 tokens por bloque reduce la media a 849, mientras la representación de la imagen conserva una similitud coseno de 0.983 respecto a la original; fusionar 64 por bloque la reduce a 673, con una similitud de 0.887. Esto muestra cómo r intercambia capacidad de cálculo por fidelidad.

Usos de la fusión de tokens#

Token Merging for Fast Stable Diffusion aplica la misma idea a los bloques Transformer de la red de eliminación de ruido de Stable Diffusion, donde puede eliminar hasta el 60 % de los tokens y generar imágenes hasta dos veces más rápido y con hasta 5,6 veces menos memoria, sin entrenamiento adicional. Más allá de las imágenes, el artículo de ToMe duplicó el rendimiento de un ViT-B de audio entrenado con fusión, con una caída del 0,4 % en mAP, y aceleró hasta dos veces el ajuste fino de un codificador automático enmascarado para vídeo. Los investigadores aplican técnicas similares de fusión y poda a los tokens visuales de los modelos de visión y lenguaje, donde LLaVA-1.5 emplea 576 tokens del contexto del modelo de lenguaje por cada imagen.

En el paquete Ultralytics, Segment Anything Model (SAM) de Meta codifica cada imagen de 1024 x 1024 píxeles como 4.096 tokens de parche; en la comparativa de MobileSAM, su codificador de imágenes ViT-H tarda 452 ms de los 456 ms que SAM necesita por imagen en una GPU, la parte del cálculo a la que se dirigen los métodos de reducción de tokens.

Fusión de tokens frente a conceptos relacionados#

• Poda de tokens: Métodos como DynamicViT y EViT descartan los tokens considerados poco importantes, lo que hace perder su información y suele requerir entrenamiento. ToMe fusiona los tokens en lugar de descartarlos, por lo que incorpora su información mediante promediado y funciona con modelos que no se han entrenado con este método.
• Tokenización: Crea los tokens iniciales a partir de parches o texto antes de que se ejecute la red. ToMe reduce su número dentro de la red.
• Atención Flash y atención dispersa: Reducen el coste o aumentan la velocidad de la atención para un número determinado de tokens. ToMe reduce el número de tokens, lo que también acelera las capas MLP, y se puede combinar con estos métodos.
• Mixture of Depths (MoD): Permite que un router aprendido decida qué tokens se saltan un bloque en los modelos de lenguaje. ToMe fusiona permanentemente los tokens redundantes mediante una regla fija de similitud.
• Menor resolución de entrada: Reduce uniformemente el número de tokens en toda la imagen. ToMe fusiona principalmente las regiones redundantes, como el fondo, y conserva las regiones detalladas a resolución completa.