Resumen Técnico: SAFE-KD

Declaración del Problema

Las arquitecturas de visión modernas logran una alta precisión, pero incurren en costos de inferencia significativos, lo que limita su despliegue en dispositivos de borde (edge devices) y sistemas con sensibilidad a la latencia. Las redes de salida temprana (early-exit networks) abordan esto mediante la incorporación de clasificadores intermedios, permitiendo que las entradas "fáciles" detengan la computación prematuramente. Sin embargo, el despliegue actual depende de reglas de parada heurísticas (por ejemplo, umbrales de confianza o entropía) que carecen de garantías formales. Estos heurísticos suelen ser poco fiables debido a la descalibración del modelo y se degradan rápidamente bajo cambios de distribución, lo que dificulta el intercambio de precisión por latencia de una manera predecible y segura. Existe una necesidad crítica de una política de inferencia adaptativa que proporcione garantías de muestra finita sobre el riesgo de las salidas tempranas.

Metodología

SAFE-KD (Destilación de Salida Temprana Controlada por Riesgo) es un envoltorio (wrapper) multi-salida universal diseñado para permitir salidas calibradas con garantías estadísticas. El método integra tres componentes principales:

1. Arquitectura Multi-Salida Universal

SAFE-KD adjunta cabezales de salida ligeros en profundidades intermedias de arquitecturas modernas (CNN y Vision Transformers) con mínimas suposiciones arquitectónicas:

• Colocación: Para CNN (ej. ResNet, ConvNeXt), las salidas se colocan después de grupos seleccionados de bloques residuales. Para ViT (ej. ViT-S, Swin-T), las salidas se adjuntan después de bloques de transformadores seleccionados, realizando el pooling de los embeddings de los tokens CLS o de la media.
• Diseño del Cabezal: Un pequeño cabezal de predicción (Global Average Pooling o Token Pooling seguido de un MLP/LayerNorm opcional y un clasificador lineal) asegura una sobrecarga mínima mientras produce logits bien formados.
• Señal: La señal de confianza principal es la Probabilidad Máxima de Softmax (MSP).

2. Destilación Jerárquica

Para abordar la debilidad y el ruido de las salidas intermedias, SAFE-KD emplea una estrategia de destilación jerárquica utilizando un profesor de Media Móvil Exponencial (EMA):

• Destilación de Conocimiento Desacoplada (DKD): Cada salida es supervisada por el profesor mediante DKD, que separa los componentes de la clase objetivo y de las clases no objetivo en la pérdida de destilación para aprovechar mejor el "conocimiento oscuro" (dark knowledge).
• Consistencia Profundo-a-Superficial: Una pérdida de consistencia alinea las salidas tempranas con la salida final, asegurando la estabilidad de la predicción a través de la profundidad de la red.
• Objetivo: La pérdida total combina la entropía cruzada estándar, la pérdida DKD contra el profesor y la pérdida de consistencia contra la salida final, ponderadas por los hiperparámetros $\alpha$ y $\beta$.

3. Salida Controlada por Riesgo mediante Control de Riesgo Conformal (CRC)

En lugar de umbrales heurísticos, SAFE-KD utiliza el Control de Riesgo Conformal para calibrar los umbrales de parada por cada salida en un conjunto de calibración retenido ($D_{cal}$):

• Definición de Riesgo Selectivo: La salida temprana se enmarca como clasificación selectiva. La puntuación de incertidumbre se define como $r_j(x) = 1 - \max_c p_j(c|x)$.
• Calibración: Para un nivel de riesgo $\delta$ especificado por el usuario, CRC calcula un umbral $\hat{\tau}_j(\delta)$ para cada salida $j$. Este umbral garantiza que el riesgo de clasificación errónea selectiva (la probabilidad de error entre las muestras aceptadas en la salida $j$) esté acotado por $\delta + O(1/|D_{cal}|)$ bajo la suposición de intercambiabilidad.
• Política de Inferencia: En el tiempo de prueba, el modelo realiza la salida en el primer cabezal $j$ donde la puntuación de incertidumbre $r_j(x) \leq \hat{\tau}_j(\delta)$. Si ninguna salida temprana satisface la condición, el modelo procede a la salida final.

Contribuciones Clave

1. Envoltorio Multi-Salida Universal: Un diseño modular que añade salidas intermedias a CNNs y ViTs con una sobrecarga mínima, preservando las características de la arquitectura base (backbone).
2. Destilación Jerárquica: Un enfoque de entrenamiento que utiliza DKD y consistencia profundo-a-superficial para mejorar la precisión y la estabilidad de las predicciones intermedias.
3. Salida Controlada por Riesgo: Un procedimiento de calibración que convierte las puntuaciones de confianza en umbrales de salida limitados por el riesgo y ajustables por el usuario, reemplazando los heurísticos frágiles con una prueba estadísticamente válida.
4. Validación Empírica: Demostrado en múltiples conjuntos de datos y arquitecturas, mostrando mejoras en el intercambio precisión-computación, una mayor calibración y robustez ante corrupciones.

Resultados Experimentales

Los autores evaluaron SAFE-KD en conjuntos de datos que incluyen CIFAR-10/100, STL-10, Oxford-IIIT Pets, Flowers102 y FGVC Aircraft, utilizando arquitecturas como ResNet-50, MobileNetV3-S, EfficientNet-B0, ConvNeXt-T, ViT-S y Swin-T.

• Precisión y Eficiencia: SAFE-KD iguala o supera la precisión de la salida final de DKD estándar y supera significativamente a los modelos de referencia MultiExit. En ViT-S con CIFAR-100, logró una mejora de precisión de +0.80% sobre ERM.
• Dominancia de Pareto: SAFE-KD logra una frontera de Pareto estrictamente dominante en comparación con el gating basado en confianza fija y entropía. Por ejemplo, en ResNet-50 con un riesgo objetivo de $\delta=0.05$, SAFE-KD redujo la profundidad computacional en un 41% manteniendo un 94.1% de precisión relativa.
• Validez del Riesgo: Las curvas de riesgo selectivo empíricas siguen consistentemente la diagonal $y=x$ o se mantienen ligeramente por debajo de ella, confirmando la validez de muestra finita. En contraste, los baselines heurísticos (ej. confianza fija de 0.95) a menudo violaron las restricciones de riesgo, resultando en riesgos observados (ej. 12% en Flowers102) que excedían con creces el objetivo.
• Robustez: Bajo la corrupción CIFAR-10-C, SAFE-KD se degradó de forma más gradual que los modelos estándar. Los umbrales calibrados rechazaron naturalmente las muestras corruptas en las salidas tempranas, forzando al modelo a utilizar capas más profundas y robustas.
• Ablación: Eliminar DKD causó una caída significativa en la precisión (-2.1%) y forzó umbrales más conservadores (aumentando la profundidad a 0.85). Eliminar la consistencia aumentó la varianza del riesgo. El método completo proporcionó el mejor equilibrio entre baja profundidad y alta precisión.

Significación y Reivindicaciones

El artículo postula que SAFE-KD transforma la inferencia adaptativa de un problema de optimización heurística en un sistema fiable y controlado por el riesgo. Al acotar matemáticamente la tasa de error de las salidas tempranas, proporciona una interfaz fiable entre las salidas del modelo y los requisitos del sistema. Esto permite a los diseñadores de sistemas tratar el parámetro de riesgo $\delta$ como una restricción estricta en lugar de un hiperparámetro ajustable, haciendo que las redes de salida temprana sean viables para despliegues de misión crítica donde la latencia predecible y los límites de error son esenciales.

Los autores señalan limitaciones, específicamente la dependencia de un conjunto de calibración (conjuntos pequeños conducen a umbrales conservadores) y la dependencia de la suposición de intercambiabilidad, que puede romperse bajo cambios severos de distribución. Sin embargo, dentro de estos límites, SAFE-KD ofrece un marco de trabajo con principios sólidos para una inferencia eficiente y segura.