Token Merging (ToMe)

La fusión de tokens (ToMe) es una técnica de vanguardia diseñada para optimizar el rendimiento y la eficiencia de las arquitecturas Transformer reduciendo el número de tokens procesados durante las pasadas hacia delante. Desarrollada originalmente para acelerar los modelos Transformer de visión (ViT), ToMe funciona identificando y combinando sistemáticamente tokens redundantes dentro de la red sin requerir entrenamiento adicional. Dado que la complejidad computacional del mecanismo de autoatención aumenta cuadráticamente con el número de tokens, fusionar tokens similares reduce drásticamente el total de operaciones de coma flotante (FLOPs), lo que permite una inferencia en tiempo real considerablemente más rápida.

Comprender el proceso de fusión de tokens#

ToMe es fundamentalmente diferente de la tokenización, que es el paso inicial de preprocesamiento consistente en dividir una imagen o un texto en tokens individuales. Mientras que la tokenización crea los elementos discretos, la fusión de tokens actúa como un mecanismo dinámico de submuestreo durante la ejecución hacia delante del modelo.

El algoritmo suele utilizar emparejamiento bipartito para evaluar la similitud entre tokens, calculando a menudo la similitud coseno entre las claves de los tokens en las capas de atención. Los tokens que comparten información visual o semántica muy similar se fusionan, normalmente promediando sus características. Esto garantiza que se conserve la información espacial o contextual esencial y, al mismo tiempo, se descarte la carga computacional innecesaria, lo que permite a frameworks como PyTorch procesar modelos de visión complejos mucho más rápido.

Aplicaciones reales de la fusión de tokens#

La fusión de tokens se ha convertido en una estrategia de optimización fundamental para desplegar arquitecturas pesadas basadas en atención en entornos con recursos computacionales limitados.

1. IA generativa y síntesis de imágenes: En los populares modelos de difusión de texto a imagen, ToMe se utiliza con frecuencia para acelerar la generación de imágenes. Al fusionar tokens del fondo o con poco detalle, el proceso de generación requiere menos pasos, ahorra una enorme cantidad de recursos de la GPU y reduce la latencia para los usuarios finales que dependen de modelos generativos. Puedes obtener más información sobre los procesos de difusión en la investigación fundamental disponible en arXiv.

2. Despliegues de IA en el borde: Desplegar modelos masivos como el Modelo Segment Anything (SAM) en dispositivos móviles es especialmente difícil debido a las limitaciones de memoria. ToMe ayuda a reducir dinámicamente la huella de memoria, permitiendo ejecutar tareas complejas de segmentación de imágenes en hardware periférico. En situaciones en las que la velocidad pura es fundamental, los ingenieros suelen decantarse por arquitecturas optimizadas de forma nativa y sin atención, como Ultralytics YOLO26, para lograr una inferencia en el borde más rápida y de extremo a extremo.

Ejemplo en Python: cálculo de la similitud entre tokens#

Aunque integrar ToMe en una arquitectura completa requiere modificar los bloques de atención, el concepto central consiste en encontrar tokens similares. El siguiente fragmento de PyTorch muestra cómo se podría calcular la similitud coseno entre un conjunto de tokens para identificar cuáles son candidatos a fusionarse.

import torch
import torch.nn.functional as F

# Simulate a batch of 4 image patches (tokens) with 64-dimensional features
tokens = torch.randn(1, 4, 64)

# Normalize the tokens to easily compute cosine similarity via dot product
normalized_tokens = F.normalize(tokens, p=2, dim=-1)

# Compute the similarity matrix between all tokens (1 x 4 x 4)
similarity_matrix = torch.matmul(normalized_tokens, normalized_tokens.transpose(1, 2))

# Tokens with high similarity scores (close to 1.0) off the diagonal
# are prime candidates for Token Merging.
print("Similarity Matrix:", similarity_matrix)

Las canalizaciones modernas de aprendizaje automático requieren equilibrar cuidadosamente la precisión y la velocidad. Tanto si utilizas la fusión de tokens para optimizar un ViT personalizado como si aprovechas la eficiencia de vanguardia de YOLO26, la gestión de estos complejos flujos de trabajo de datos se simplifica enormemente con la Ultralytics Platform. La Platform proporciona un ecosistema intuitivo para la anotación de datos automatizada, el entrenamiento fluido en la nube y el despliegue sólido de modelos en diversos entornos de hardware de computación en el borde. Las organizaciones que amplían sus iniciativas de visión artificial confían en estas herramientas para poner en producción modelos de vanguardia de forma fiable y eficiente.