What DINO saw: ALiBi positional encoding reduces positional bias in Vision Transformers

Este trabajo demuestra que el ajuste fino de Vision Transformers con codificación posicional ALiBi reduce significativamente los sesgos posicionales, preservando al mismo tiempo las semánticas generales y permitiendo una segmentación efectiva en imágenes de microscopía complejas.

Lo esencial

¡Hola! Imagina que los Transformers de Visión (como el famoso modelo DINOv2) son como super-inteligentes detectives entrenados para mirar fotos y entender qué hay en ellas (un perro, un coche, una flor). Estos detectives son increíbles: pueden aprender cosas nuevas sin que nadie les diga la respuesta correcta, simplemente mirando millones de fotos.

Sin embargo, los autores de este paper descubrieron un defecto secreto en estos detectives, especialmente cuando miran fotos de materiales científicos (como microscopios de baterías o metales).

Aquí te explico qué pasó, usando analogías sencillas:

1. El Problema: El Detective "Prejuicioso"

Imagina que entrenaste a un detective para que reconozca perros. Lo hiciste mirando fotos de perros en parques, playas y salones. El detective aprendió muy bien.

Pero, sin darse cuenta, el detective desarrolló un prejuicio por la posición.

• Si ve una foto, piensa: "¡Ah! Como el perro está en la esquina superior izquierda, ¡seguro es un perro!".
• Si el perro está en el centro, el detective se confunde.

En el mundo de la ciencia de materiales, esto es un desastre. Las fotos de microscopio suelen ser homogéneas (como un trozo de pan o una batería cortada por la mitad). No hay un "arriba" o un "abajo" especial; todo es igual en todas partes.

El modelo DINOv2 original, al mirar estas fotos, seguía pensando: "¡El borde izquierdo es diferente al derecho!", aunque no lo fuera. Esto hacía que cuando intentaban usar el modelo para segmentar (dibujar líneas alrededor de las partes importantes de la imagen), el modelo fallaba y dibujaba líneas en lugares incorrectos, solo porque la imagen estaba en una posición específica.

La analogía: Es como si tuvieras un mapa del tesoro, pero el mapa siempre te dijera que el tesoro está en la esquina superior derecha, sin importar dónde estés realmente.

2. La Solución: "ALiBi" (El Detective Justo)

Los autores decidieron arreglar este defecto. ¿Cómo? Cambiando la "brújula" interna del modelo.

• Lo viejo (DINOv2): Usaba una "posición aprendida" que era como un mapa rígido. El modelo memorizaba: "El token 1 es arriba-izquierda, el token 2 es arriba-derecha".
• Lo nuevo (ALiBi): Usan una técnica llamada ALiBi (Atención con Sesgos Lineales). Imagina que en lugar de dar al detective un mapa con coordenadas fijas, le das una regla de distancia relativa.
  • En lugar de decir "Estás en la posición 5", le dicen: "Estás a 3 pasos del vecino de la izquierda y a 5 del de la derecha".

Esto hace que el modelo deje de obsesionarse con dónde está el objeto en la foto y se centre en qué es el objeto y cómo se relaciona con sus vecinos.

3. El Experimento: ¿Funciona?

Los científicos hicieron algo muy inteligente:

1. Tomaron un modelo DINOv2 ya entrenado (el "detective prejuicioso").
2. Le quitaron su vieja brújula y le pusieron la nueva regla de distancia (ALiBi).
3. Le dijeron: "Mira, sigue aprendiendo, pero ahora usa esta nueva regla. Intenta que tus respuestas sean iguales a las que daban antes".

El resultado:

• El nuevo modelo (ALiBi-Dv2) aprendió a ser justo. Ya no le importaba si el objeto estaba en la esquina o en el centro.
• Cuando lo usaron para analizar fotos de baterías y metales, ¡funcionó de maravilla! Logró separar los poros, las grietas y los materiales con mucha más precisión que el modelo original.
• Además, siguió siendo bueno en fotos normales (perros, gatos, paisajes), así que no perdió sus habilidades generales.

4. ¿Por qué es importante?

En la ciencia de materiales, los investigadores a menudo tienen que analizar miles de imágenes de microscopio para ver si una batería es segura o si un metal tiene grietas.

• Antes: Tenían que limpiar las imágenes o usar trucos complicados para engañar al modelo y que no se confundiera con la posición.
• Ahora: Con este nuevo modelo, pueden simplemente mirar la imagen y confiar en que el modelo verá la estructura real del material, no un "fantasma" creado por la posición de la foto.

En resumen

Este paper nos dice: "Los modelos de IA son geniales, pero a veces tienen prejuicios ocultos sobre dónde están las cosas. Si cambiamos su forma de entender la posición (usando ALiBi), podemos hacerlos más justos y precisos, especialmente para tareas científicas donde todo debe ser tratado por igual, sin importar la esquina de la foto."

Es como quitarle los anteojos de color rosa a un detective para que vea la realidad tal como es, sin distorsiones.

Resumen técnico

Resumen Técnico: Reducción del Sesgo Posicional en Vision Transformers mediante Codificación ALiBi

1. El Problema: Sesgo Posicional en Modelos Fundacionales

Los Vision Transformers (ViT), especialmente los modelos fundacionales de características como DINOv2, han demostrado un rendimiento excepcional en tareas de visión por computadora. Sin embargo, el artículo identifica un defecto crítico: estos modelos desarrollan un sesgo posicional (positional bias) en sus representaciones de características.

• Naturaleza del problema: Las características aprendidas no dependen únicamente del contenido semántico, sino que están fuertemente correlacionadas con la posición del parche en la imagen. Esto se manifiesta como gradientes lineales (rampas) en los canales de salida, independientemente del contenido de la imagen.
• Impacto en Ciencias de los Materiales: Este sesgo es particularmente perjudicial en campos como la ciencia de materiales, donde las imágenes (microscopía electrónica SEM/TEM) suelen ser secciones transversales homogéneas sin una dirección preferida. Al aplicar segmentación "zero-shot" o entrenada débilmente, el modelo confunde la posición del píxel con la clase del material, generando artefactos de segmentación (ej. clasificar erróneamente el centro de la imagen de manera diferente a los bordes).
• Limitaciones de soluciones previas: Métodos como el uso de tokens "register" o redes desruidoras (Denoising Vision Transformers) reducen parcialmente el problema, pero no eliminan el sesgo fundamental inherente a la codificación posicional aprendida o a la arquitectura del modelo.

2. Metodología

Los autores proponen una solución que combina el análisis de sesgos con un ajuste fino (finetuning) estratégico utilizando ALiBi (Attention with Linear Biases).

• Caracterización del Sesgo (Linear Probing):

  • Entrenaron "sondas lineales" (linear probes) para predecir funciones de rampa 1D (izquierda-derecha, arriba-abajo, diagonales, radiales) a partir de las características de salida de varios ViT (DINO, DINOv2, DINOv3, MAE, etc.).
  • Hallazgo clave: Se descubrió que ciertos canales de salida son casi puramente funciones de rampa posicional. Este fenómeno está presente en modelos auto-supervisados (DINO, MAE) pero es mucho menor o inexistente en modelos supervisados (ViT entrenados con ImageNet, CLIP). Curiosamente, incluso DINOv3, que usa codificación posicional rotacional (RoPE), mostró un alto sesgo posicional.

• Propuesta de Solución: ALiBi-Dv2:

  • Reemplazo de PE: Se eliminó la codificación posicional (PE) aprendida de un checkpoint de DINOv2 pre-entrenado.
  • Integración de ALiBi: Se introdujeron codificaciones posicionales relativas basadas en ALiBi. En lugar de añadir un vector posicional a la entrada, ALiBi añade un sesgo lineal a las puntuaciones de atención en función de la distancia relativa entre tokens.
  • Configuración 2D: Se adaptó ALiBi para ser consciente de la bidimensionalidad (2D) de las imágenes, utilizando distancias euclidianas con condiciones de frontera cilíndricas (wrap-around) para evitar asimetrías.
  • Estrategia de Entrenamiento:
    • El modelo se ajustó finamente (finetuning) para recuperar las características semánticas originales de DINOv2, utilizando los embeddings sesgados del modelo original como objetivo (teacher).
    • Se congelaron los canales más "posicionales" identificados durante el entrenamiento para forzar al modelo a aprender representaciones más homogéneas.
    • No se requirió un nuevo entrenamiento desde cero; el objetivo era puramente desviar el modelo de sus sesgos aprendidos.

3. Contribuciones Clave

1. Caracterización Exhaustiva del Sesgo: Demostraron mediante linear probing que el sesgo posicional es un problema generalizado en modelos fundacionales auto-supervisados (incluyendo DINOv3 con RoPE), afectando tanto a imágenes naturales como a datos fuera de distribución (micrografías).
2. Método de Mitigación Eficaz: Presentan un método para transformar un modelo DINOv2 existente en una versión "ALiBi-Dv2" que elimina casi por completo el sesgo posicional sin perder la riqueza semántica.
3. Validación en Dominios Críticos: Demuestran que las características resultantes son superiores para la segmentación de imágenes de microscopía electrónica, donde la homogeneidad es crucial.

4. Resultados Principales

• Reducción del Sesgo Posicional:

  • Las puntuaciones $R^2$ de las sondas lineales (que miden cuánto se puede predecir la posición a partir de las características) cayeron drásticamente. Mientras que DINOv2 tenía puntuaciones altas (ej. 0.83 en micrografías), ALiBi-Dv2 obtuvo puntuaciones negativas o cercanas a cero (-0.23), indicando que las características ya no codifican la posición de manera lineal.
  • Las visualizaciones de PCA mostraron una eliminación de los gradientes de borde y rampas radiales presentes en DINOv2 y DINOv3.

• Rendimiento en Segmentación Semántica (Benchmarks Estándar):

  • En conjuntos de datos estándar como PASCAL VOC y ADE20K, el modelo ALiBi-Dv2 mantuvo o mejoró ligeramente el rendimiento (mIoU) en comparación con DINOv2 original, demostrando que la eliminación del sesgo no degrada la capacidad de reconocimiento de objetos.

• Segmentación Entrenable en Ciencias de Materiales:

  • En tareas de segmentación interactiva (entrenando clasificadores como XGBoost sobre características ViT para imágenes de baterías de iones de litio y aleaciones), el modelo ALiBi-Dv2 superó significativamente a DINOv2 y DVT (Denoising ViT).
  • Caso de éxito: En la segmentación de efectos de "pore-back" (material fuera de plano en baterías), DINOv2 falló al confundir la posición con la clase, mientras que ALiBi-Dv2 logró una segmentación homogénea y precisa, crucial para el análisis cuantitativo de porosidad y grietas.

5. Significado e Impacto

Este trabajo es fundamental para la aplicación de la IA en ciencias de los materiales y microscopía avanzada.

• Resolución de una brecha de dominio: Los modelos fundacionales actuales están sesgados hacia imágenes naturales (con horizontes, cielos, etc.). Este método permite adaptar estos modelos potentes a imágenes homogéneas y sin dirección preferida, que son comunes en la investigación de materiales.
• Eficiencia: Ofrece una vía para "limpiar" modelos fundacionales existentes sin necesidad de entrenar desde cero, lo cual es computacionalmente costoso.
• Generalización: Sugiere que el sesgo posicional es una propiedad inherente a la auto-supervisión en ViTs y que el uso de codificaciones posicionales relativas como ALiBi es una arquitectura más robusta para tareas que requieren invariancia posicional estricta.

En conclusión, "What DINO Saw" demuestra que al reemplazar la codificación posicional aprendida por una basada en sesgos lineales (ALiBi), se pueden obtener características visuales ricas en semántica pero libres de artefactos posicionales, habilitando una segmentación de alta precisión en imágenes científicas complejas.