Distilling Vision Transformers for Distortion-Robust Representation Learning

Resumen Técnico: Destilación de Vision Transformers para el Aprendizaje de Representaciones Robustas ante Distorsiones

1. El Problema

El aprendizaje de representaciones visuales robustas es un desafío crítico en visión por computadora, especialmente cuando las imágenes presentan ruido, desenfoque (blur), oclusiones o pérdida de píxeles. En aplicaciones reales como la imagen médica, el teledetección o los sistemas autónomos, los datos de alta calidad suelen ser escasos o inexistentes, obligando a los modelos a operar sobre observaciones corruptas.

Los modelos actuales, incluso aquellos preentrenados con éxito mediante aprendizaje autosupervisado (SSL), sufren una degradación significativa en su capacidad semántica cuando se enfrentan a distorsiones severas, ya que estas interrumpen tanto las señales visuales de bajo nivel como la información estructural de alto nivel. El problema central que aborda este estudio es: ¿Es posible recuperar las representaciones de un codificador limpio observando únicamente versiones distorsionadas de las imágenes?

2. Metodología

Los autores proponen un marco de destilación de conocimiento asimétrica basado en Vision Transformers (ViT). El enfoque no busca reconstruir los píxeles de la imagen (como en los métodos de restauración de imágenes), sino recuperar el espacio de características semánticas en el espacio latente.

• Configuración Maestro-Estudiante: Se utilizan dos modelos con la misma arquitectura ViT, ambos inicializados con un modelo preentrenado (DINO-ViT). El Maestro permanece congelado y procesa imágenes limpias ($x$), mientras que el Estudiante es el único que se entrena y procesa las versiones distorsionadas ($\tilde{x}$).
• Objetivo de Destilación Multinivel: Para guiar al estudiante de manera efectiva, proponen una función de pérdida compuesta por tres niveles de alineación:
  1. Alineación Semántica Global ($L_{cls}$): Utiliza el error cuadrático medio (MSE) para alinear el class token del maestro y el estudiante, asegurando que la información categórica se preserve.
  2. Alineación Espacial Local ($L_{patch}$): Alinea los patch tokens mediante MSE para evitar que la distorsión colapse la estructura espacial y el diseño de la imagen.
  3. Alineación de la Estructura de Atención ($L_{attn}$): Utiliza la divergencia KL para alinear los mapas de atención de las cabezas del Transformer. Esto obliga al estudiante a "mirar" las mismas regiones de interés que el maestro.

La función de pérdida total es una combinación ponderada de estos tres componentes: $\mathcal{L} = \lambda_{cls} \mathcal{L}_{cls} + \lambda_{patch} \mathcal{L}_{patch} + \lambda_{attn} \mathcal{L}_{attn}$.

3. Contribuciones Clave

• Marco de Destilación Asimétrica: Un método para entrenar codificadores robustos sin necesidad de acceder nunca a datos limpios durante el entrenamiento.
• Objetivo de Destilación Multinivel: La integración de la alineación de embeddings globales, características de parches y mapas de atención, proporcionando una supervisión mucho más rica que los métodos de contraste tradicionales.
• Demostración de Robustez y Transferibilidad: El método no solo mejora el rendimiento en el dominio de entrenamiento, sino que demuestra una capacidad excepcional para generalizar a dominios completamente distintos (médico y satelital).

4. Resultados Principales

• Rendimiento In-Domain (ImageNet-100): El método supera consistentemente tanto al aprendizaje supervisado convencional como al método de Contrastive Inversion (CI), especialmente bajo ruido gaussiano y desenfoque.
• Transferencia de Dominio: En conjuntos de datos como RESISC45 (imágenes aéreas) y CAMELYON17 (histopatología médica), el modelo propuesto muestra mejoras sustanciales, demostrando que las representaciones aprendidas son semánticamente profundas y no solo adaptaciones superficiales al ruido.
• Eficiencia de Etiquetas: Los experimentos muestran que el modelo es altamente eficiente en el uso de etiquetas; incluso con solo un 5% de las etiquetas disponibles, el rendimiento se aproxima al de un modelo entrenado con el 100% de los datos.
• Análisis de Atención: Las visualizaciones confirman que, a diferencia de los modelos supervisados que pierden el foco ante la distorsión, los modelos destilados mantienen mapas de atención centrados en los objetos de interés, similares a los de una imagen limpia.

5. Significado y Conclusión

Este trabajo cambia el paradigma de la robustez ante la distorsión: en lugar de intentar "limpiar" la imagen o entrenar con aumentos de datos masivos, propone aprender a inferir la semántica limpia desde la observación corrupta.

La importancia de este estudio radica en su aplicabilidad en entornos donde el etiquetado es costoso y la calidad de la imagen es inherentemente baja. Al demostrar que la destilación multinivel puede recuperar la estructura de atención y la coherencia espacial, los autores ofrecen una vía robusta para el despliegue de modelos de visión en condiciones del mundo real imperfectas.