Por qué los Vision Transformer (ViT) necesitan tokens de registro
(openreview.net)- Los tokens de alta norma en los mapas de características de los ViT aparecen principalmente en regiones de fondo con poca información, y se interpretan como un fenómeno en el que los patch tokens se reutilizan como espacio de almacenamiento para cálculos internos
- La solución consiste en agregar tokens de registro al input sequence, separados de los parches de imagen, para delegar ese rol de cálculo interno a tokens dedicados
- En ViT entrenados tanto con aprendizaje supervisado como con auto-supervisión, se reducen los artefactos y los mapas de características y attention map se vuelven más suaves, lo que favorece las tareas visuales downstream
- Se reportaron mejoras en object discovery y dense prediction, pero en OpenCLIP el rendimiento de object localization cayó ligeramente, por lo que se corrigió la afirmación de que “mejora en todos los modelos”
- Al añadir entre 4 y 16 tokens de registro, se genera un 2~6% de cómputo adicional, y aumentar la cantidad de tokens no siempre mejora el rendimiento en tareas densas
Artefactos observados en los mapas de características de ViT
- Vision Transformer se usa como una herramienta potente para el aprendizaje de representaciones visuales, pero se han identificado artefactos en los mapas de características de ViT entrenados con supervisión y auto-supervisión
- Los artefactos se observan como tokens de alta norma durante la inferencia
- Aparecen principalmente en regiones de fondo de la imagen con poca información
- Parece que la red reutiliza esos tokens para cálculos internos
- Estos tokens provienen de parches de fondo similares a los parches vecinos, y se evalúa que contienen poca información local del parche y más información global de la imagen
- En DINOv2 y otros ViT, la degradación en la resolución espacial y la precisión del attention map puede afectar el rendimiento downstream en object discovery y localization
Enfoque con tokens de registro
- La solución propuesta es agregar tokens de registro a la secuencia de entrada de Vision Transformer, separados de los parches de imagen
- Los registros, al igual que el CLS token, están separados de los tokens de parches de imagen, pero se definen como tokens que no se usan directamente en el cálculo de la loss
- El objetivo es evitar que los patch tokens existentes se usen como “espacio de almacenamiento” para cálculos internos, dejando ese papel a tokens dedicados
- Este enfoque está relacionado con trabajos previos como Memory Transformer, pero el análisis de estos artefactos y el uso de nuevos tokens para corregirlos está específicamente enfocado en este paper
Cambios en los mapas de características y en las tareas downstream
- Se reporta que al agregar registros se eliminan estos artefactos en modelos entrenados tanto con supervisión como con auto-supervisión
- Los mapas de características y el attention map se vuelven más suaves, lo que beneficia el procesamiento visual downstream
- Se resume que el modelo alcanzó un nuevo state of the art en dense visual prediction task de modelos visuales auto-supervisados
- También permite usar métodos de object discovery en modelos más grandes
- En DINOv2 y DeiT-III, el rendimiento en object discovery mejora y el rendimiento en tareas a nivel de imagen en general se mantiene
Excepciones por modelo y análisis adicional
- En OpenCLIP, agregar registros no mejoró los resultados de object localization, y en la tabla se observa una ligera caída de rendimiento
- Los autores reconocieron en su respuesta que la afirmación de que “el rendimiento de object discovery mejora en todos los modelos” no estaba respaldada por los datos
- Ese contenido fue corregido en la Section 3.3
- En OpenCLIP, la Figure 7 muestra una reducción en la norma, pero se discute la inconsistencia de que esto no se traduzca en una mejora de object localization en la Table 3
- En la evaluación con LOST, para cada modelo se selecciona el embedding con mejor rendimiento entre keys, queries y values
- En CLIP, values fue el mejor, y los artifacts aparecen en keys o queries, pero no en values, lo que es consistente con la Table 3 y el análisis cualitativo
- En MAE, parece que no aparecen los mismos “outlier patches”
- Esto podría deberse a que MAE se entrena solo con local loss, por lo que no necesita, o necesita menos, reunir información global
- La precisión de clasificación ImageNet con linear probing de MAE-ViT-Large es de 75%, inferior a la de otros métodos SSL
Rendimiento, costo y limitaciones
- Los registros son un método simple, pero aumentan el costo computacional por los tokens adicionales
- Se reporta un aumento de 2~6% en el cómputo al añadir entre 4 y 16 tokens de registro
- Aumentar la cantidad de registros no siempre mejora las tareas densas
- Al pasar de 0 a 1 registro, los artefactos se eliminan y mejoran segmentation y depth prediction
- Después de 1 registro, aunque se agreguen más, el feature map ya está limpio, por lo que segmentation y depth prediction no mejoran mucho más
- El rendimiento en classification sí tendió a mejorar adicionalmente a medida que aumentaba el número de registros, pero todavía no hay una intuición clara sobre la razón
- En el dataset NYU depth, al pasar de 8 a 16 registros, el RMSE aumenta ligeramente
- Es posible que los hyperparameters óptimos hayan cambiado al aumentar la longitud de la secuencia
- La diferencia de 0.03 en RMSE entre 8 y 16 es pequeña comparada con la diferencia de 0.1 al pasar de 0 a 1, por lo que también podría ser simple ruido
Adopción en OpenReview y puntos de debate en la revisión
- La decisión en OpenReview fue Accept (oral)
- En la meta review, los cuatro reviewers dieron una calificación de 8, y se reconoció la contribución técnica del paper
- Los reviewers expresaron preocupaciones sobre inconsistencias de rendimiento, la relación entre dataset bias y los modelos SSL, la elección de configuraciones de optimización y el impacto del gradient clipping
- En el proceso de rebuttal, los autores añadieron experimentos adicionales y análisis en el appendix
- En discusiones posteriores, se reforzó el análisis de mecanismos internos, como si los registros heredan el comportamiento de los outliers, si los image tokens vuelven a tener un carácter más local y cómo cambia la norma del CLS token
1 comentarios
Opiniones en Hacker News
Según el paper, un register es un token adicional aprendible que se agrega a la secuencia de entrada durante el entrenamiento de modelos Vision Transformer.
Se añade después de la capa de embeddings de parches como un valor aprendible, igual que el token [CLS], y al final del Vision Transformer se descartan los tokens register y solo se usan el token [CLS] y los tokens de parches como representación de la imagen.
Este token le da al modelo un espacio para almacenar, procesar y consultar información global durante la propagación hacia adelante, de modo que no tenga que forzar el uso de los tokens de parches para ese fin.
Como resultado, desaparecen los artefactos que aparecen en los mapas de características de los Vision Transformer entrenados y los tokens atípicos de norma alta; los mapas de características se vuelven más suaves, y también mejoran el rendimiento en dense prediction y el descubrimiento no supervisado de objetos.
Dado que el costo de inferencia aumenta solo 2% y mejora bastante el rendimiento de ViT, parece un resultado muy limpio, casi como un almuerzo gratis.
Las expresiones regulares probabilísticas ya son una cosa, pero esto es libertinaje matemático a otro nivel.
Aun así, es realmente sorprendente que funcione así.
Suena parecido a un attention sink o a un pause token.
Me pregunto si alguien investigó pasar la salida cruda de un pase directamente como entrada del siguiente.
Parece relacionado: "Let's Think Dot by Dot: Hidden Computation in Transformer Language Models" https://arxiv.org/abs/2404.15758
Este paper considera que no está claro si la mejora de rendimiento del chain-of-thought en modelos de lenguaje se debe a una descomposición de tareas similar a la humana o a la cantidad adicional de cómputo que aportan los tokens extra.
Mostró que los Transformer pueden resolver tareas algorítmicas difíciles que no podían resolver al responder sin tokens intermedios, incluso usando solo tokens de relleno sin significado, por ejemplo
......; aunque dice que entrenar el uso de esos tokens de relleno requiere supervisión específica y densa.Por eso me pregunto si también habría beneficios en agregar tokens de cálculo/register especiales a los LLM de texto.
Discusión adicional: https://news.ycombinator.com/item?id=40182695, https://www.reddit.com/r/LocalLLaMA/comments/1cf2w5a/transfo...
En los últimos años probé decenas, quizá más de 100 papers e ideas en visión y reconocimiento multimodal, y este fue uno de esos casos raros en los que todo funcionó bien.
Por ejemplo, este modelo se entrena usando 4 tokens register junto con una pérdida estilo Matryoshka, y genera un embedding de 64 dimensiones muy pequeño que puede ser útil para quienes buscan alternativas a CLIP: https://huggingface.co/unum-cloud/uform3-image-text-english-...
Lo vi en ICLR y fue uno de los mejores papers de este año; eso también quedó claro en la sesión de pósters.
Felicitaciones a los autores.
Llevo varios meses usando DinoV2, y también probé un modelo que tiene 4 tokens register junto con [CLS] + tokens de parches.
Como varios tokens de embedding del modelo anterior sin registers ya eran parte de la solución, y el token [CLS] no queda alineado entre el modelo con 0 registers y el de 4 registers, no adoptamos el nuevo modelo con registers.
Ojalá [CLS] y los tokens de parches pudieran alinearse de alguna forma entre ambos modelos.
Parece un attention sink para modelos de visión