💻 computer science

Putting Registers to Work: Task Registers for Token Pruning in Vision Transformers

Este artículo propone la Poda Adaptativa a la Tarea (TAP, por sus siglas en inglés), un método que utiliza registros específicos de la tarea para clasificar dinámicamente los tokens y asignar presupuestos de poda a través de diferentes tareas de Vision Transformer, logrando mejoras significativas en el rendimiento de procesamiento mientras mantiene un desempeño competitivo en clasificación de imágenes, segmentación semántica y detección de objetos.

Autores originales: Hongsen Cao, Mona Jaber, Shanxin Yuan, Ahmed Sayed

Publicado 2026-08-12

📖 5 min de lectura🧠 Análisis profundo

CC BY 4.0

Autores originales: Hongsen Cao, Mona Jaber, Shanxin Yuan, Ahmed Sayed

Artículo original bajo licencia CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta es una explicación generada por IA del artículo a continuación. No ha sido escrita ni avalada por los autores. Para mayor precisión técnica, consulte el artículo original. Leer descargo de responsabilidad completo

Imagina que tienes un cerebro de robot superinteligente que observa imágenes e intenta comprender qué hay dentro de ellas. Este cerebro está construido como un equipo gigante de trabajadores diminutos, cada uno sosteniendo una pieza del rompecabezas (un "token"). Para tomar una decisión, cada trabajador tiene que charlar con todos los demás trabajadores. ¿El problema? Si la imagen es grande, el número de charlas explota y el cerebro se abruma, se vuelve lento y tiene hambre de electricidad. Los científicos han estado intentando solucionar esto diciéndole al cerebro que ignore las piezas aburridas de la imagen y solo conserve las importantes, algo así como un profesor diciéndole a una clase que solo escuche a los estudiantes que tienen algo inteligente que decir. Esto se llama "poda de tokens" (token pruning). Pero aquí está la parte difícil: lo que cuenta como "inteligente" o "importante" cambia dependiendo de lo que el robot esté intentando hacer. Una pieza de una imagen que es crucial para detectar un gato puede ser inútil para averiguar de qué color es el cielo. Durante mucho tiempo, los científicos intentaron usar la misma "lista de ignorados" para cada trabajo, pero resulta que un mismo tamaño rara vez sirve para todo.

Este artículo, titulado "Putting Registers to Work" (Poniendo los registros a trabajar), aborda exactamente ese problema. Los autores, un equipo de la Universidad Queen Mary de Londres, descubrieron que las reglas para decidir qué piezas de la imagen desechar son totalmente diferentes para distintas tareas. Encontraron que una estrategia que funciona de maravilla para detectar objetos (como encontrar un perro en un parque) podría en realidad perjudicar la capacidad del robot para comprender una escena (como averiguar si el perro está sobre hierba o sobre pavimento). Para resolver esto, inventaron un nuevo y astuto sistema llamado Poda Adaptativa a la Tarea (TAP).

Piensa en el cerebro del robot como una cocina con mucho movimiento. Normalmente, el chef (la política de poda) utiliza la misma receta para decidir qué ingredientes desechar, sin importar si está haciendo una sopa o un pastel. Los autores se dieron cuenta de que esto era un error. En su lugar, le dieron a la cocina un conjunto de "Registros de Tarea" especiales —piensa en ellos como tarjetas de recetas inteligentes y mágicas que cambian dependiendo del plato que se esté cocinando—. Cuando el robot necesita identificar un gato, saca la "Tarjeta de Gato", que le dice al cerebro exactamente qué partes de la imagen debe conservar y cuáles debe ignorar. Cuando necesita mapear una habitación entera, cambia a la "Tarjeta de Habitación". Estas tarjetas no solo deciden qué tirar, sino que también determinan cuándo tirarlo mientras el cerebro procesa la imagen, capa por capa.

Los investigadores probaron esto congelando el cerebro de un robot estándar e intentando probar diferentes "reglas de ignorado" sin reentrenarlo. Encontraron tres grandes sorpresas. Primero, la mejor manera de elegir las piezas importantes para encontrar objetos es exactamente lo opuesto a la mejor manera de mapear escenas. Segundo, el robot es muy sensible a cómo maneja las primeras capas de procesamiento, especialmente para tareas sencillas como decir si una imagen es un gato o un perro. Tercero, cuando el robot desecha una pieza de la imagen, necesita recordar cómo era para rellenar el hueco más tarde. Pero aquí está el truco: para encontrar objetos, es mejor simplemente pretender que la pieza nunca estuvo allí y usar un vecino cercano para llenar el hueco. Para mapear escenas, es mejor recordar realmente la diferencia exacta y pegarla de nuevo más tarde.

Para solucionar todo esto, TAP utiliza una única tarjeta de registro "activa" para el trabajo actual. Esta tarjeta viaja a través del cerebro, observando la imagen y tomando tres decisiones sobre la marcha:

Selección: Qué tokens mantener y cuáles descartar.
Presupuesto: Cuántos tokens descartar en cada paso (quizás descartar algunos al principio, más tarde, o viceversa).
Recuperación: Cuánta de la información detallada "olvidada" debe traer de vuelta cuando el robot necesite dibujar un mapa detallado.

Los resultados son impresionantes. Al usar estas tarjetas específicas para cada tarea, el robot se volvió mucho más rápido sin perder su inteligencia. En una prueba estándar para encontrar objetos (COCO), se volvió 1.32 veces más rápido manteniendo casi la misma precisión (cayendo solo 0.3 puntos). En una prueba para mapear escenas (ADE20K), se volvió 1.30 veces más rápido con una mínima caída en la calidad. Incluso para la clasificación de imágenes sencilla, se mantuvo competitivo.

Los autores sugieren que este enfoque es un paso hacia un cerebro de robot verdaderamente unificado: uno que pueda alternar entre detectar un coche, pintar un cuadro de una calle o leer un cartel, todo utilizando el mismo motor central pero con diferentes "estilos de pensamiento" activados sobre la marcha. No solo adivinaron que esto funcionaría; lo demostraron realizando miles de experimentos, congelando el cerebro para probar reglas individuales y mostrando que su nuevo sistema de "registro" supera consistentemente a los métodos antiguos de "talla única". Es un recordatorio de que, en el mundo de la IA, a veces lo más inteligente que puedes hacer es saber exactamente cuándo ignorar el ruido, y que la definición de "ruido" depende enteramente de lo que estés intentando ver.

Resumen Técnico: Poniendo los Registros a Trabajar: Registros de Tarea para la Poda de Tokens en Vision Transformers

Declaración del Problema

Los Vision Transformers (ViTs) exhiben un escalamiento computacional cuadrático con el número de tokens, lo que motiva la poda de tokens para reducir los costos de inferencia. Aunque los criterios de poda (por ejemplo, puntuaciones de atención, cobertura espacial) se desarrollan a menudo para la clasificación de imágenes, las arquitecturas base de ViT preentrenadas se están adaptando cada vez más a tareas de predicción densa como la segmentación semántica y la detección de objetos.

Los autores identifican una brecha crítica: las políticas de poda diseñadas para una tarea no necesariamente se transfieren de manera efectiva a otras. A través de experimentos controlados de "sonda" (probe), el artículo demuestra que:

Divergencia en la Clasificación de Criterios: Las tareas de segmentación y detección clasifican los criterios de poda de manera diferente. Por ejemplo, la poda basada en la atención causa caídas significativas de precisión en las capas iniciales para la clasificación, mientras que la cobertura espacial es más estable. Por el contrario, el criterio que minimiza el error de detección suele funcionar mal para la segmentación.
Sensibilidad a la Profundidad: La efectividad de un criterio de poda varía significamente según la profundidad de la red y difiere entre tareas.
Conflictos de Recuperación: Las tareas de predicción densa requieren la reconstrucción de características en las posiciones eliminadas. El artículo encuentra que la segmentación y la detección prefieren puntos finales de recuperación opuestos: la segmentación se beneficia de restaurar el desplazamiento completo de la característica (la diferencia entre los tokens eliminados y los retenidos), mientras que la detección funciona mejor utilizando solo el token "sustituto" (stand-in) retenido.

Los métodos existentes suelen aplicar una política de poda única y estática o utilizar selectores agnósticos a la tarea, fallando al no tener en cuenta estas variaciones específicas de la tarea y dependientes de la profundidad.

Metodología: Poda Adaptativa a la Tarea (TAP)

El artículo propone la Poda Adaptativa a la Tarea (TAP), un marco que utiliza Registros de Tarea para controlar dinámicamente la poda de tokens, la asignación de presupuesto y la recuperación de características.

Mecanismo Central:

Registros de Tarea: En lugar de usar los registros meramente como almacenamiento agnóstico a la tarea para artefactos de características (como en trabajos previos como Darcet et al., 2024), TAP introduce un registro inicial aprendible por tarea. Durante una pasada hacia adelante (forward pass), solo el registro asociado con la tarea actual está activo.
Estado Evolutivo: El registro activo evoluciona junto con los tokens de la imagen a través del backbone. Su estado en cualquier capa condiciona tres decisiones clave:
1. Selección de Tokens: El registro computa una puntuación para cada parche candidato reutilizando las proyecciones de consulta (query) y clave (key) existentes del bloque. Los tokens se clasifican y un subconjunto se retiene para la auto-atención.
2. Asignación Exacta de Presupuesto: Una lectura lineal compartida, condicionada al estado del registro, determina el número entero exacto de tokens a eliminar en cada capa de poda. Esto asegura que se cumpla exactamente el presupuesto global de eliminación sin requerir una pérdida de tasa (rate loss).
3. Escalamiento de la Recuperación: Para tareas densas, el registro controla un factor de escala ( $\alpha$ ) que determina qué parte del desplazamiento de la característica almacenado (la diferencia entre el token eliminado y su sustituto) se restaura en la lectura.

Variantes de Implementación:

TAP-J (Ajuste Fino Conjunto): Cada tarea tiene un backbone, un registro de tarea y lecturas separadas. Esto permite una especialización total.
TAP-F (Base Congelada): Un único backbone congelado se comparte entre las tareas. Los registros específicos de la tarea, las actualizaciones de bajo rango (LoRA) y los cabezales se entrenan conjuntamente. Esto reduce el almacenamiento de parámetros mientras mantiene la adaptabilidad.

Recuperación de Características:
Los tokens eliminados se emparejan de forma determinista con el token retenido que tenga la mayor similitud de coseno en el espacio de claves (el "sustituto" o stand-in). La diferencia de característica (desplazamiento u offset) se almacena. En los puntos de lectura densa, la reconstrucción se calcula como $\hat{x}_i = x_{\text{stand-in}} + \alpha \cdot \delta_i$ , donde $\alpha$ es aprendido por tarea y profundidad. Las tareas de clasificación omiten este paso de recuperación.

Contribuciones Clave

Marco Controlado: Los autores desarrollan un marco que congela los puntos de control sin poda y aplica criterios de reducción libres de parámetros capa por capa para aislar el impacto de las elecciones de poda a través de la clasificación, segmentación y detección. Esto revela la desalineación en las clasificaciones de los criterios y las preferencias de recuperación.
Registros de Tarea como Controladores: El artículo redefine el papel de los registros de un almacenamiento pasivo a controladores activos y específicos de la tarea del cómputo disperso. El registro coordina la selección de tokens, la distribución exacta del presupuesto de eliminación a través de la profundidad y el escalamiento de la recuperación.
Poda y Recuperación Unificadas: TAP trata la eliminación de tokens y la recuperación de características como una operación unificada. Empareja los tokens descartados con sus sustitutos y escala la restauración de los desplazamientos basándose en el registro de la tarea, manteniendo las características reconstruidas fuera del codificador para mantener la dispersión.

Resultados Experimentales

Los autores evalúan TAP en ImageNet-1K (clasificación), ADE20K (segmentación) y COCO (detección/segmentación de instancias) utilizando un backbone ViT-B con $\rho = 0.5$ (50% de tokens retenidos).

Rendimiento:
- TAP-J logra las caídas de precisión más pequeñas entre los métodos de ajuste fino completo. En ADE20K, alcanza 47.0 mIoU (solo 0.2 de caída respecto a sin poda) con un throughput del codificador de 1.30×. En COCO, alcanza 53.7 box AP (0.3 de caída) con 1.32× de throughput. En ImageNet-1K, alcanza 83.2 top-1 (0.4 de caída).
- TAP-F (backbone compartido) sigue ligeramente por debajo de TAP-J pero supera significativamente a las reglas estáticas y otros baselines, logrando 45.8 mIoU en ADE20K y 52.5 box AP en COCO con una sobrecarga de parámetros mínima (~1.2M de parámetros específicos de la tarea).
Eficiencia: TAP-J proporciona aceleraciones del codificador de 1.26× a 1.32× manteniendo una precisión competitiva. La sobrecarga del controlador (consulta del registro, lectura del presupuesto, contabilidad) representa solo un ~0.7% del tiempo de ejecución del codificador.
Memoria: TAP reduce la memoria de inferencia pico en un 9–18% en comparación con los baselines densos, siendo la segmentación la que observa las mayores ganancias debido a la reducción de los tamaños de activación en los bloques intermedios.
Estudios de Ablación: Eliminar el registro de tarea (usando uno compartido) o fijar la asignación/recuperación a valores estáticos degrada significativamente el rendimiento, confirmando la necesidad de un control adaptativo a la tarea y a la profundidad.

Significado y Reivindicaciones

El artículo afirma que la reducción de tokens no puede tratarse como una regla única transferible. A medida que las representaciones de ViT evolucionan de estructuras locales a semánticas de alto nivel, el valor de los tokens cambia tanto con la profundidad como con la tarea específica que lee el backbone.

La importancia de TAP radica en su capacidad para adaptar el cómputo disperso a las necesidades específicas de diferentes tareas visuales utilizando un backbone compartido. Al poner los registros de tarea en funcionamiento, TAP:

Resuelve los conflictos en las clasificaciones de los criterios de poda entre tareas.
Optimiza la distribución del presupuesto de eliminación a través de la profundidad de la red.
Condiciona la recuperación de características a los requisitos específicos de las lecturas densas frente a las dispersas.

Los autores posicionan a TAP como un paso hacia sistemas de visión unificados que pueden servir múltiples tareas a través de un backbone compartido con un almacenamiento específico de la tarea mínimo, equilibrando el rendimiento predictivo con el throughput del codificador de manera más efectiva que las políticas de poda estáticas. El trabajo se mantiene modesto, señalando que TAP-J queda 0.1 puntos por detrás del mejor método de solo clasificación (DiffRate), probablemente porque la clasificación depende de un token de clase y omite la recuperación densa, dejando uno de los mecanismos de TAP subutilizado para esa tarea específica. Se sugiere como trabajo futuro explorar presupuestos condicionados por la imagen y múltiples sustitutos.