Ask the Expert: Collaborative Inference for Vision Transformers with Near-Edge Accelerators

1. Planteamiento del Problema

La implementación de Transformadores de Visión (ViT) en dispositivos de borde (edge) es un desafío debido a su alta complejidad computacional y al gran número de parámetros. Por otro lado, la descarga completa (offloading) a la nube introduce latencias significativas debido al ancho de banda y la distancia de comunicación.

• Limitaciones actuales: Los dispositivos de borde tradicionales (ej. Raspberry Pi) no pueden ejecutar ViTs complejos con tiempos de respuesta aceptables.
• Ineficiencia de la nube: El envío de datos a la nube para inferencia genera una latencia de ida y vuelta (RTT) que es inaceptable para aplicaciones en tiempo real.
• Fallas de enfoques existentes: Las técnicas de inferencia colaborativa diseñadas para Redes Neuronales Convolucionales (CNN), como la inferencia dividida o las salidas tempranas, no son adecuadas para ViTs debido a su estructura columnar constante (que carece de "cuellos de botella" naturales) y sus dependencias globales de autoatención. Además, los modelos de "Mezcla de Expertos" (MoE) tradicionales sufren de inestabilidad en el entrenamiento y colapso de enrutadores.

2. Metodología Propuesta

Los autores proponen un marco de inferencia colaborativa que orquesta un modelo generalista ligero en el dispositivo de borde y múltiples modelos expertos de tamaño medio en un acelerador de "cerca del borde" (near-edge).

A. Arquitectura del Sistema

• Dispositivo de Borde: Ejecuta un ViT ligero (generalista) que procesa todas las entradas.
• Acelerador Near-Edge: Alberga una biblioteca de ViT expertos, cada uno especializado en subconjuntos específicos de clases del dataset.
• Flujo de Trabajo:
  1. El modelo de borde realiza la inferencia local.
  2. Si la confianza (probabilidad máxima softmax) supera un umbral $\tau$, la predicción se acepta localmente.
  3. Si la confianza es baja (muestra incierta), la muestra se envía al acelerador near-edge junto con sus índices Top-k de predicción.

B. Mecanismo de Enrutamiento Top-k (Zero-Overhead)

En lugar de usar un enrutador externo o costoso, el sistema reutiliza las predicciones Top-k del modelo de borde como señal de enrutamiento.

• Si el modelo de borde no está seguro, sus $k$ clases más probables se mapean a los subconjuntos de datos correspondientes.
• Se selecciona dinámicamente el experto especializado en la unión de esos subconjuntos para refinar la predicción.
• Esto elimina la sobrecarga computacional de un módulo de enrutamiento adicional.

C. Estrategia de Entrenamiento de Especialistas Progresivos

Para entrenar a los expertos sin perder la capacidad de generalización, se propone una estrategia de destilación ponderada dinámica:

• Partición de Datos: El dataset se divide en subconjuntos no superpuestos basados en superclases.
• Ponderación Dinámica: Durante el entrenamiento, se utiliza un factor de escala $\omega(t)$ que aumenta linealmente con las épocas.
  • Fases tempranas: El modelo aprende características globales de todo el dataset (generalización).
  • Fases tardías: El peso de los datos del subconjunto objetivo aumenta, forzando la especialización del experto en ese dominio específico.
• Destilación: Los expertos se entrenan utilizando un modelo ViT grande preentrenado como "maestro", minimizando la pérdida de entropía cruzada ponderada.

3. Contribuciones Clave

1. Marco de Inferencia Colaborativa para ViT: Una formulación general que desacopla el enrutador de los expertos, evitando la inestabilidad de optimización acoplada típica de los MoE.
2. Mecanismo de Enrutamiento sin Sobrecarga: Uso inteligente de las predicciones Top-k del modelo de borde para seleccionar expertos, eliminando la necesidad de componentes adicionales.
3. Estrategia de Entrenamiento Progresivo: Un método novedoso que permite a los expertos especializarse en subconjuntos de datos manteniendo una alta precisión general, superando a los expertos estáticos.
4. Validación Experimental Real: Pruebas exhaustivas en un entorno de prueba real (RPi5/Orin Nano como borde y AGX Orin como near-edge) utilizando el dataset CIFAR-100.

4. Resultados Experimentales

Las pruebas se realizaron comparando el enfoque propuesto contra la ejecución solo en borde, solo en near-edge, y co-inferencia con modelos generalistas.

• Precisión:
  • La estrategia de entrenamiento progresivo mejoró la precisión de especialización de los expertos en un 4.12% en los subconjuntos objetivo en comparación con expertos estáticos.
  • La precisión general aumentó un 2.76% sobre los expertos estáticos.
  • Se logró una mejora de precisión de hasta 1.70% frente a la co-inferencia con modelos generalistas de tamaño similar.
• Latencia:
  • Reducción de latencia de hasta un 45% en comparación con la ejecución local en el dispositivo de borde (debido a la descarga selectiva de muestras difíciles).
  • Aunque la latencia es mayor que la descarga completa a near-edge (por el procesamiento local inicial), la reducción en la transmisión de datos compensa la latencia de red.
• Consumo Energético:
  • Reducción del consumo energético de hasta un 46% en comparación con la descarga completa a near-edge, ya que la mayoría de las muestras fáciles se resuelven localmente sin transmitir datos.
• Eficiencia (ROI):
  • El uso de expertos de tamaño medio demostró una mejor relación "precisión-latencia" y "precisión-energía" que el uso de modelos generalistas grandes en el acelerador near-edge.

5. Significado e Impacto

Este trabajo presenta una solución viable para el despliegue de modelos de visión avanzados en entornos de recursos limitados.

• Optimización de Recursos: Demuestra que es más eficiente desplegar múltiples modelos expertos medianos en aceleradores near-edge compartidos que depender de un único modelo masivo en la nube o en el borde.
• Escalabilidad: El enfoque es particularmente prometedor para tareas más complejas (como ImageNet), donde la brecha entre la precisión Top-1 y Top-k del modelo de borde es mayor, ofreciendo más oportunidades para que los expertos corrijan errores.
• Sostenibilidad: Al reducir la dependencia de la nube y optimizar el uso de energía, el marco contribuye al desarrollo de una IA más sostenible y accesible en entornos remotos o con infraestructura deficiente.
• Complementariedad: El marco es ortogonal a técnicas de compresión (como cuantización), pudiendo compensar la pérdida de precisión que estas técnicas podrían causar en modelos de borde.

En resumen, el artículo valida que una arquitectura colaborativa inteligente, que combina la proximidad de los datos en el borde con la capacidad computacional especializada cerca del borde, ofrece un punto óptimo (Pareto) superior en términos de precisión, latencia y eficiencia energética.