Resumen Técnico: KLAS - Uso de la Similitud para Unir Redes Neuronales y Mejorar los Compromisos entre Precisión y Eficiencia

1. Enunciado del Problema

El despliegue de grandes transformadores de visión preentrenados (ViT) y redes neuronales convolucionales (CNN) enfrenta un desafío fundamental: entornos diversos que van desde servidores en la nube potentes hasta dispositivos periféricos con recursos limitados requieren una selección flexible de modelos dentro de presupuestos de cómputo específicos. Aunque la Búsqueda de Arquitectura Neuronal (NAS) intenta equilibrar la precisión y la eficiencia, la NAS tradicional uno-a-uno es computacionalmente prohibitiva, y los enfoques uno-a-muchos o NAS de disparo cero a menudo se limitan a un único espacio de diseño de modelos, fallando en aprovechar el rico conjunto de modelos preentrenados existentes.

La unión de modelos (model stitching) ofrece una alternativa prometedora al combinar bloques de diferentes modelos "ancla" preentrenados para crear redes interpoladas. Sin embargo, los enfoques existentes de unión muchos-a-muchos, como SN-Net, dependen de heurísticas ingenuas (por ejemplo, "unión más cercana" de anclas adyacentes o emparejamiento fijo de bloques) para seleccionar configuraciones de unión. Estas heurísticas a menudo producen compromisos subóptimos entre precisión y eficiencia y carecen de generalización en diversas familias de modelos. Además, el espacio de búsqueda para las configuraciones de unión es vasto ($O(k^2n^2)$ para $k$ modelos de profundidad $n$), lo que hace inviable la búsqueda exhaustiva. El problema central es la falta de un método principiado y automatizado para identificar configuraciones de unión compatibles que maximicen el rendimiento sin incurrir en altos costos de entrenamiento.

2. Metodología: El Marco KLAS

Los autores proponen KLAS (Unión de Anclas basada en Divergencia KL), un marco que automatiza la selección de uniones aprovechando la similitud entre las representaciones intermedias de modelos preentrenados.

Concepto Central: Divergencia KL como Métrica de Similitud

KLAS postula que una unión exitosa depende tanto de la similitud representacional (alineación geométrica de las activaciones) como de la similitud funcional (preservación de la información relevante para la tarea y los límites de decisión). Los autores argumentan que la divergencia de Kullback-Leibler (KL) es única para capturar ambos aspectos:

• Alineación Representacional: Cuantifica la distancia estadística entre las distribuciones de activación, indicando si una transformación ligera puede mapear las activaciones de origen a las activaciones de destino.
• Alineación Funcional: Evalúa qué tan bien las salidas de un bloque de origen preservan los límites de decisión del modelo objetivo, asegurando que la red unida mantenga el rendimiento de la tarea aguas abajo con un ajuste fino mínimo.

A diferencia de métricas como la Alineación de Núcleo Centrado (CKA), el Error Cuadrático Medio (MSE) o la Entropía Cruzada (CE), que a menudo fallan al predecir el éxito de la unión, la divergencia KL mide directamente la compatibilidad de las distribuciones de probabilidad generadas por sondas lineales.

El Algoritmo KLAS

El marco opera en dos etapas principales:

1. Selección de Anclas:

   • KLAS calcula la divergencia KL entre las distribuciones de probabilidad softmax de los últimos bloques de los modelos preentrenados candidatos.
   • Los modelos con baja divergencia KL en el último bloque se consideran anclas compatibles, ya que comparten límites de decisión y distribuciones de confianza de clase similares. Esto reemplaza la heurística de "unión más cercana" utilizada en trabajos anteriores.

2. Selección de Bloques:

   • Una vez seleccionadas las anclas, KLAS evalúa los pares de bloques utilizando una puntuación de unión $\Gamma(i, j)$ definida como:
     $$\Gamma(i, j) = \frac{\Theta(P^f_i, P^g_j)}{\Theta(P^g_j, P^g_{j+1})}$$
     Donde:
     • $\Theta(P^f_i, P^g_j)$ es la distancia de activación entre anclas cruzadas (divergencia KL entre el bloque de origen $i$ y el bloque de destino $j$).
     • $\Theta(P^g_j, P^g_{j+1})$ es la capacidad intra-ancla de bloques (divergencia KL entre bloques consecutivos en la ancla de destino).
   • Esta puntuación equilibra la carga de transformación en la capa de unión hipotética contra la capacidad del bloque de destino para absorber diferencias distribucionales.
   • ProbeNet: Para calcular estas puntuaciones eficientemente sin entrenar capas de unión completas, KLAS emplea una arquitectura unificada llamada ProbeNet. Adjunta sondas de clasificador lineal a cada bloque de los modelos ancla y las entrena conjuntamente en un solo paso hacia adelante y hacia atrás. Esto incurre en un costo único insignificante (por ejemplo, 0.25 días de GPU para Swin-B) en comparación con el entrenamiento de sondas individuales.

3. Selección de Configuración:

   • Las configuraciones de unión candidatas se filtran basándose en un umbral ajustable $\tau$ relativo a la puntuación mínima dentro de los grupos de FLOPs. Esto asegura una frontera de Pareto densa de modelos desplegable a lo largo del espectro de precisión-eficiencia.

3. Contribuciones Clave

1. Identificación de Limitaciones Heurísticas: El artículo demuestra que las heurísticas ingenuas (como las de SN-Net) no son universalmente óptimas, a menudo perdiendo configuraciones de "unión lejana" superiores (por ejemplo, conectar Swin-Ti directamente con Swin-B) que las heurísticas pasan por alto.
2. Métrica de Similitud Principiada: La introducción de la divergencia KL como métrica para la selección de uniones que captura tanto la alineación representacional como la funcional sin requerir la instanciación o el entrenamiento de capas de unión.
3. Marco KLAS: Un marco generalizable y de bajo costo que automatiza la selección de anclas y bloques en diversas familias de modelos (ViT, CNN y LLM).
4. Validación Empírica: Experimentos exhaustivos que muestran mejoras consistentes sobre las líneas base en múltiples conjuntos de datos (ImageNet-1K, CIFAR-100, ADE20K, TruthfulQA) y arquitecturas.

4. Resultados Experimentales

Los autores evaluaron KLAS frente a SN-Net y otras líneas base bajo condiciones idénticas de ajuste fino (50 épocas para visión, 10 épocas para LLM).

• Compromisos Precisión-Eficiencia:

  • ImageNet-1K: KLAS logra hasta un 1.21% más de precisión Top-1 al mismo costo computacional o una reducción de 1.33× en FLOPs a precisión equivalente en comparación con SN-Net.
  • Puntuaciones AUC: En la familia Swin, KLAS alcanzó un Área bajo la Curva (AUC) de 0.8950, superando significativamente a SN-Net (0.8345) y a otras métricas de similitud como CKA (0.8124).
  • Entre Arquitecturas: KLAS unió exitosamente ResNet-18 con Swin-B, superando las líneas base de inferencia en cascada en un 2-3% de precisión Top-1 con FLOPs similares.

• Predicción Densa (Segmentación Semántica):

  • En ADE20K usando Mask2Former, las uniones seleccionadas por KLAS lograron hasta +0.9% mIoU sobre las líneas base de SN-Net con presupuestos de FLOPs comparables.

• Modelos de Lenguaje Grande (LLM):

  • Usando anclas Llama 1B y 3B en la prueba TruthfulQA, los modelos unidos por KLAS superaron a ESTA (un método de unión de LLM basado en heurísticas) en hasta 0.017 ROUGE-1 y 0.033 ROUGE-2, a pesar de tener menos parámetros en algunas configuraciones.

• Efectividad de la Métrica:

  • La divergencia KL recuperó exitosamente el 88.9% de las configuraciones de unión viables Swin Ti-B, mientras que CKA recuperó solo el 5.5% y las heurísticas de SN-Net el 61.1%.

5. Significado y Afirmaciones

El artículo afirma que KLAS aborda la brecha crítica en la unión de modelos al reemplazar heurísticas ad-hoc con un enfoque principiado y guiado por la similitud. Su significado radica en:

• Generalización: Funciona a través de diversas familias arquitectónicas (Transformers, CNN, Híbridos) y modalidades (Visión, Lenguaje), superando las limitaciones de métodos restringidos a familias individuales.
• Eficiencia de Costos: Genera una frontera de Pareto continua de modelos desplegable con sobrecarga computacional adicional insignificante (solo entrenamiento de sondas y ajuste fino de la capa de unión), evitando los costos prohibitivos de la NAS tradicional.
• Optimización Sistemática: Al capturar explícitamente las diferencias distribucionales mediante la divergencia KL, KLAS permite el descubrimiento de configuraciones de unión no intuitivas y de alto rendimiento (por ejemplo, "unión lejana") que los métodos heurísticos pasan por alto.

Los autores concluyen que KLAS proporciona una solución robusta y automatizada para optimizar el compromiso entre precisión y eficiencia en escenarios de despliegue, aprovechando efectivamente el ecosistema existente de modelos preentrenados sin requerir entrenamiento desde cero.