Dense2MoE: Pushing the Pareto Frontier of On-Device LLMs via Unified Pruning and Upcycling

Resumen Técnico: Dense2MoE

Enunciado del Problema
El despliegue de Modelos de Lenguaje Grandes (LLM) en dispositivos periféricos con recursos limitados (por ejemplo, conducción autónoma, IA en el borde) se ve severamente obstaculizado por el "muro de memoria de inferencia". Según la teoría Roofline del hardware, los LLM densos están limitados por la memoria durante la decodificación autoregresiva, lo que significa que simplemente reducir los FLOPs teóricos es insuficiente para mejorar la latencia. Las soluciones existentes enfrentan un compromiso fundamental:

1. Poda Estructurada: Métodos como UIDL reducen la latencia descartando capas redundantes, pero inevitablemente sacrifican la precisión del modelo al eliminar Perceptrones Multicapa (MLP), que sirven como los principales repositorios de conocimiento de los LLM.
2. Reciclaje a MoE: Convertir modelos densos a arquitecturas de Mezcla de Expertos (MoE) aumenta la capacidad sin incrementar linealmente la computación dinámica. Sin embargo, el reciclaje estándar a menudo sufre de "colapso de representación" debido a una inicialización homogénea de expertos (duplicando MLP nativos), lo que requiere un preentrenamiento continuo masivo para inducir especialización. Además, los métodos recientes de reaprovechamiento disperso a menudo se centran en la redundancia intra-capas, sin abordar las superposiciones estructurales inter-capas a nivel macro que causan redundancia en la atención y sobrecargas en la carga de la memoria caché KV.

Metodología: Dense2MoE
Los autores proponen Dense2MoE, un marco unificado que integra la poda estructurada de capas con el reciclaje a MoE mediante una técnica novedosa llamada Reciclaje por Fusión de Capas (LF-UC). El método opera en tres etapas:

1. Análisis de Similitud Inter-Capas: El marco identifica la redundancia estructural analizando capas adyacentes en busca de alta similitud coseno tanto en las salidas del decodificador como en las entradas de los MLP. Crucialmente, impone una restricción estricta de coincidencia de norma de características para asegurar que la reutilización de MLP redundantes no introduzca anomalías de activación o cambios de escala.
2. Reciclaje por Fusión de Capas (LF-UC): En lugar de descartar completamente las capas redundantes, Dense2MoE poda físicamente solo los módulos de atención (MHA) de estas capas para eliminar operaciones intensivas en ancho de banda de memoria. Los MLP preentrenados de estas capas podadas se preservan y reaprovechan como expertos heterogéneos dentro de las capas base retenidas.
   • Esto crea un grupo de expertos donde los "expertos base" son MLP nativos duplicados y los "expertos suplementarios" son los MLP preservados de las capas podadas.
   • Un enrutador dinámico a nivel de token despacha selectivamente tokens a estos expertos o los omite mediante rutas residuales.
   • Este diseño aprovecha la heterogeneidad estructural inter-capas para mitigar el colapso de enrutamiento "El Ganador se lo Lleva Todo" común en el reciclaje homogéneo.
3. Preentrenamiento Continuo: La arquitectura modificada pasa por una fase ligera de preentrenamiento continuo (225 mil millones de tokens) para optimizar el mecanismo de enrutamiento y alinear los expertos fusionados. Este presupuesto consume menos del 1% de la computación requerida para entrenar un modelo MoE desde cero.

Para determinar la profundidad de poda óptima, los autores emplean una métrica de recompensa consciente del hardware inspirada en MnasNet, mapeando las restricciones de poda a la frontera de Pareto de latencia versus precisión para hardware objetivo específico (por ejemplo, NVIDIA Jetson Thor-U).

Contribuciones Clave

• Marco Unificado: Dense2MoE integra perfectamente la poda estructurada de capas y el reciclaje a MoE, resolviendo el compromiso entre la eficiencia de inferencia y la precisión del modelo.
• Arquitectura Novedosa (LF-UC): Al podar los cuellos de botella de atención mientras se reciclan los MLP como expertos heterogéneos, el método reduce el consumo de ancho de banda de memoria sin destruir la capacidad representacional.
• Optimalidad de Pareto de Vanguardia: El marco avanza la frontera de Pareto para la latencia de inferencia en dispositivos versus la precisión del modelo, superando a las líneas base densas, a los métodos de compresión más avanzados y a los enfoques estándar de reciclaje.
• Escalabilidad Robusta: El método se valida en diversas escalas de modelos (de 0.5B a 7B) y familias arquitectónicas distintas (Qwen2.5 y LLaMA2).

Resultados Experimentales
Evaluado en 9 benchmarks diversos (incluyendo MMLU, ARC, GSM8K y HumanEval) utilizando modelos Qwen2.5 y LLaMA2:

• Rendimiento vs. Latencia: La configuración óptima de Dense2MoE (reteniendo 19 capas de una semilla Qwen2.5-0.5B de 24 capas) logró una latencia física de inferencia de 271.0 ms en el SoC Thor-U (bajando de 307.5 ms para el modelo semilla) mientras alcanzaba una puntuación promedio en benchmarks de 47.31, superando significativamente a la línea base densa (37.53).
• Comparación con Líneas Base: A diferencia de los métodos de poda tradicionales (por ejemplo, UIDL, LLM-Streamline) y enfoques de reaprovechamiento disperso (por ejemplo, ToMoE, LLaMA-MoE), que sufrieron un colapso severo de capacidad en tareas de razonamiento y codificación, Dense2MoE mantuvo una retención representacional robusta. Superó a la línea base densa en más de 9.78 puntos porcentuales en puntuación promedio mientras activaba solo 0.42B de parámetros.
• Escalabilidad: Al aplicarse a modelos más grandes (Qwen2.5-1.5B y LLaMA2-7B), Dense2MoE superó consistentemente a sus contrapartes densas en 3.04 y 3.54 puntos porcentuales, respectivamente, después del ajuste fino en el mismo conjunto de datos de 225 mil millones de tokens.

Significado
El artículo afirma que Dense2MoE establece una receta práctica y de estándar industrial para reaprovechar LLMs densos de código abierto en arquitecturas altamente eficientes para dispositivos periféricos. Al explotar sistemáticamente la redundancia estructural inter-capas, permite la optimización conjunta de la latencia de inferencia y la precisión del modelo. El trabajo demuestra que es posible romper el muro de memoria de inferencia y preservar capacidades complejas de razonamiento sin los costos prohibitivos de entrenar modelos MoE desde cero o la degradación de precisión asociada con la poda agresiva. Esto proporciona un continuo flexible de modelos óptimos para implementaciones en el borde críticas para la latencia.