Resumen Técnico: SnapViT – Vision Transformers Elásticos a partir de Modelos Preentrenados sin Reentrenamiento

Declaración del Problema

Los modelos fundacionales de visión han alcanzado un rendimiento de vanguardia en diversas tareas, pero suelen publicarse en un conjunto limitado de tamaños predeterminados (por ejemplo, conteos de parámetros específicos para las familias DINOv3 o DeIT). Esta rigidez obliga a los profesionales a tomar decisiones de despliegue subóptimas cuando se enfrentan a restricciones reales de cómputo, latencia y costo. Las soluciones existentes para la flexibilidad de modelos enfrentan limitaciones significativas:

• Destilación de Conocimiento: Requiere una arquitectura objetivo predeterminada y a menudo depende de conjuntos de datos de preentrenamiento no públicos, lo que socava la robustez.
• Inferencia Elástica: Los métodos como las representaciones Matryoshka o las Redes Slimmable requieren que los modelos sean diseñados con estructuras anidadas específicas durante el preentrenamiento, lo que los hace inaplicables a modelos preentrenados existentes y listos para usar.
• Poda Estructurada: Aunque es eficaz para reducir la memoria y el cómputo, la mayoría de las técnicas de poda actuales están adaptadas a tareas o restricciones específicas y suelen requerir reentrenamiento, dejando un vacío para modelos universalmente adaptables que no necesiten datos etiquetados ni un cabezal de clasificación.

Metodología: SnapViT

Los autores proponen SnapViT (Single-shot network approximation for pruned Vision Transformers), un método de poda estructurada post-preentrenamiento que genera un continuo de subredes elásticas a partir de un único modelo preentrenado sin necesidad de reentrenamiento ni datos etiquetados. El núcleo del método es una puntuación de capacidad de poda ($P$) que clasifica los parámetros para su eliminación basándose en dos componentes:

1. Aproximación de Hessian Local (Gradientes Autosupervisados)

Para capturar la sensibilidad local de los parámetros sin un cabezal de clasificación, el método aproxima la diagonal de la matriz Hessian utilizando gradientes autosupervisados. Específicamente, emplea el objetivo DINO para calcular gradientes sobre un conjunto de datos de imágenes. La curvatura local $H^{(l)}$ se estima como la norma al cuadrado promedio de los gradientes:
$$H^{(l)} \approx \frac{1}{N_D} \sum_{i=1}^{N_D} \|\nabla_\theta L_i\|^2$$
Esto proporciona una medida base de qué tan críticos son los parámetros individuales dentro de sus bloques específicos (por ejemplo, neuronas de alimentación hacia adelante o cabezales de atención).

2. Estimación de Hessian Global (Algoritmo Evolutivo)

Para abordar la limitación de ignorar las dependencias entre bloques (elementos fuera de la diagonal de la Hessian), SnapViT introduce una novedosa aproximación utilizando la Estrategia de Evolución Natural Exponencial (xNES).

• En lugar de calcular explícitamente la completa e intratable matriz Hessian, xNES aprende un conjunto de factores de escala por estructura ($c$) que reponderan las sensibilidades locales.
• El algoritmo optimiza una distribución Gaussiana multivariante sobre estos factores de escala. La función de aptitud (fitness) carece de etiquetas, midiendo la similitud de coseno entre los embeddings del modelo original y el modelo podado (después de la compresión PCA).
• A través de actualizaciones de gradiente natural, la matriz de covarianza ($\Sigma$) de la estrategia evolutiva evoluciona para aproximar la inversa de la verdadera Hessian inter-bloque ($H^{(g)} \approx \alpha \Sigma^{-1}$). Esto permite al método identificar qué bloques pueden podarse conjuntamente sin degradar la calidad de la representación.

3. Puntuación de Capacidad de Poda Unificada y Generación Elástica

La puntuación de capacidad de poda final es el producto elemento a elemento de la sensibilidad local y los factores de escala globales:
$$P = \text{diag}(H^{(l)}) \odot (M c)$$
donde $M$ es una matriz de membresía que expande los factores a nivel de bloque a parámetros individuales.

• Generación de un solo paso (Single-Shot): Una vez calculada $P$, todos los parámetros se clasifican globalmente. Esta clasificación única permite la extracción de subredes para cualquier nivel de dispersión deseado ($S \in [0, 1]$) sin optimización adicional.
• Alcance: El método poda tanto combinaciones fila-columna dentro de los bloques de alimentación hacia adelante como cabezales de atención completos.

Contribuciones Clave

1. Estrategia de Poda Elástica: Una estrategia de poda eficiente y rápida para Vision Transformers preentrenados que produce modelos elásticos adaptables a cualquier restricción computacional sin reentrenamiento.
2. Aproximación de Hessian Evolutiva: Una estrategia novedosa utilizando un algoritmo genético (xNES) para aproximar los componentes fuera de la diagonal de la Hessian para estructuras de red, capturando correlaciones entre bloques que las aproximaciones diagonales tradicionales omiten.
3. Rendimiento Libre de Etiquetas y Agnóstico al Cabezal: El método alcanza un rendimiento de vanguardia bajo una poda significativa sin requerir datos etiquetados o un cabezal de clasificación, generalizando a través de modelos fundacionales supervisados y autosupervisados.

Resultados Experimentales

Los autores evaluaron SnapViT en ocho conjuntos de datos (incluyendo ImageNet-1k, Pascal VOC y varios benchmarks de grano fino) utilizando clasificación de k-vecinos más cercanos (k-NN), probing lineal y segmentación semántica lineal.

• Rendimiento vs. SOTA: SnapVió de forma consistente iguala o supera a los métodos de vanguardia (LAMP, LLM Surgeon, FPTP, SparseGPT, SNIP Magnitude, NViT) en varios niveles de dispersión.
  • Para DINOv1 ViT-B/16, el método poda hasta un 40% de dispersión con menos del 5% de degradación de precisión, logrando una aceleración de inferencia de 1.58x.
  • Con un 50% de dispersión en modelos AugReg supervisados, supera a SNIP en un 7% y a FPTP en un 12.3% en clasificación lineal.
  • Para SigLIPv2 y DINOv3 (modelos fundacionales de gran escala), el método supera significativamente a las bases, con mejoras de hasta un 34.3% en clasificación lineal respecto a los segundos mejores métodos.
• Eficiencia: Todo el proceso de poda para una familia de modelos (por ejemplo, DeIT-III) toma menos de cinco minutos en una sola GPU A100.
• Post-procesamiento: El método es compatible con la corrección de pesos post-poda (estilo SparseGPT) y el ajuste fino completo (fine-tuning). Con un solo paso de corrección de pesos, el rendimiento en la dispersión extrema (por ejemplo, 50% en SigLIPv2) se preserva en gran medida.
• Escalabilidad: Aunque el preentrenamiento a gran escala (por ejemplo, DINOv3, SigLIPv2) hace que la poda sea más desafiante debido al conocimiento representacional distribuido, SnapViT sigue siendo efectivo, especialmente cuando se combina con corrección de pesos o horizontes de optimización más largos (500 iteraciones de xNES).

Significado y Reivindicaciones

El artículo afirma que SnapViT cierra la brecha entre la rigidez de los despliegues actuales de modelos fundacionales y la necesidad de una inferencia flexible y adaptable a los recursos. Al permitir la extracción de un continuo de subredes a partir de un único modelo preentrenado en un solo paso, permite a los usuarios adaptar con precisión los modelos de vanguardia a sus presupuestos computacionales específicos.

Los autores enfatizan que su enfoque:

• Elimina la necesidad de reentrenamiento, ahorrando recursos computacionales y tiempo significativos.
• Elimina la dependencia de datos etiquetados, haciéndolo aplicable a modelos fundacionales autosupervisados que carecen de cabezales de clasificación.
• Captura correlaciones estructurales globales mediante aproximación evolutiva, superando las limitaciones de las métricas de poda solo locales.
• Reduce el impacto ambiental al disminuir los requisitos de cómputo y el uso de energía de los Vision Transformers, reduciendo potencialmente las emisiones de CO2 de la inferencia.

El trabajo se presenta como una solución práctica para desplegar modelos de visión a gran escala en entornos restringidos manteniendo un alto rendimiento en diversas tareas de aplicación.