Resumen Técnico: LLMForge

Declaración del Problema

Los modelos de lenguaje Transformer con menos de mil millones de parámetros se despliegan cada vez más en dispositivos de borde para aprovechar las ventajas en privacidad, latencia y costos. Sin embargo, la inferencia en el dispositivo está limitada por un ancho de banda de memoria ajustado, presupuestos de energía y restricciones térmicas. Los enfoques existentes suelen depender de la compresión post-entrenamiento (cuantización, poda), lo que deja sin explorar los grados de libertad arquitectónicos subyacentes. Además, las actuales pipelines de Búsqueda de Arquitectura Neural Consciente del Hardware (NAS) para Transformers adolecen de tres limitaciones:

1. Espacio de Búsqueda Restringido: Las parametrizaciones de atención a menudo fijan las dimensiones por cabeza y acoplan las capacidades de clave/valor, limitando las configuraciones de forma factibles a una única elección de número de cabezas.
2. Cuello de Botella en la Evaluación: Evaluar candidatos mediante entrenamiento desde cero es computacionalmente costoso.
3. Modelado de Costos Grosero: Los costos de hardware se modelan frecuentemente mediante clases únicas de aceleradores o proxies groseros (por ejemplo, recuento de parámetros), enmascarando la re-clasificación dependiente del sustrato de los candidatos.

Metodología

Los autores presentan LLMForge, un marco de NAS consciente del hardware que integra tres contribuciones composables en un único bucle de búsqueda evolutiva NSGA-II.

1. Atención de Cabezas Infinitas (IHA)

IHA desacopla los parámetros de forma de la atención ($n_h$, $n_{kv}$, $d_{qk}$, $d_v$) como variables independientes por capa. A diferencia de la Atención Multi-Cabeza (MHA) estándar o la Atención de Consulta Agrupada (GQA), que imponen restricciones de divisibilidad ($d_{model} = n_h \cdot d_h$) y acoplamiento Q/K–V ($d_h = d_{qk} = d_v$), IHA relaja estas restricciones. Solo preserva la condición de agrupación de GQA ($n_{kv} | n_h$).

• Impacto: Esto expande el espacio de configuración de atención factible por capa en aproximadamente 400× (de 27 configuraciones bajo GQA a 11,250 bajo IHA dentro de los rangos de búsqueda definidos).

2. Forge-Former (Surrogado de Precisión)

Para abordar el cuello de botella en la evaluación, los autores introducen Forge-Former, un modelo surrogado basado en un codificador Transformer.

• Arquitectura: Codifica una arquitectura como una secuencia de tokens de capa, donde cada token concatena los parámetros IHA de la capa. Los campos escalares se elevan mediante capas lineales y se suman con incrustaciones posicionales. Un codificador Transformer pre-LN captura las interacciones dentro de la capa (mediante subcapas feedforward) y las interacciones entre capas (mediante auto-atención).
• Rendimiento: En la división de prueba IHA, Forge-Former alcanza un $\rho$ de Spearman = 0.75 y un $\tau$ de Kendall = 0.58, superando a las líneas base MLP y de bosque aleatorio entre 1.4 y 2.5×. Permite la selección de NSGA-II a una fracción del costo de entrenar desde cero.

3. Forge-DSE (Motor de Exploración del Espacio de Diseño)

Forge-DSE es el motor de búsqueda multiobjetivo que empareja Forge-Former con un modelo de costos de hardware multi-backend.

• Bucle de Búsqueda: Evoluciona poblaciones de arquitecturas parametrizadas por IHA utilizando NSGA-II con ordenamiento no dominado y supervivencia basada en la distancia de congestión. Emplea operadores de mutación personalizados (eliminación, duplicación, rotación, perturbación numérica) y un paso de reparación de factibilidad.
• Co-evolución del Surrogado: El marco refina opcionalmente Forge-Former durante la búsqueda mediante el reentrenamiento de muestras informativas de la población actual y el ajuste fino del surrogado mediante reproducción de experiencia, corrigiendo el desplazamiento de distribución.
• Backends de Hardware: El sistema soporta tres backends distintos:
  1. Backend A: Costos medidos en una GPU NVIDIA A100 utilizando ZEUS.
  2. Backend B: Costos analíticos mediante Timeloop para flujos de datos sístólicos (Gemmini), de fila estacionaria (Eyeriss) y de atención fusionada (FLAT).
  3. Backend C: Un simulador de flujo de datos en anillo multi-chip (rDXE) para aceleradores de borde, realizando una búsqueda conjunta de arquitectura y chip.

Resultados Clave

Validación del Surrogado

En el conjunto de datos IHA, Forge-Former supera significativamente a las líneas base en fidelidad de clasificación y recuperación del top-X%. Reduce el Error Absoluto Medio (MAE) para el 5% superior de arquitecturas en aproximadamente 5× en comparación con las líneas base MLP y Bosque Aleatorio. El análisis t-SNE confirma que las incrustaciones de Forge-Former reflejan mejor la similitud de rendimiento, agrupando estrechamente las colas de alta pérdida.

Rendimiento de Búsqueda Multi-Backend

A través de cuatro sustratos de hardware (A100, Gemmini, Eyeriss, FLAT), las búsquedas convergen hacia arquitecturas de frente de Pareto distintas que siguen los cuellos de botella de costos específicos de cada sustrato:

• A100 (Rico en ancho de banda): Resulta en anchos de MLP planos y las cabezas de atención más anchas.
• Eyeriss/Gemmini (Limitados por caché KV): Colapsan $n_{kv}$ en una atención de consulta agrupada agresiva y compensan con una capacidad de MLP más amplia.
• FLAT (Flexible): Produce el modelo más pequeño con profundidad superficial y cabezas estrechas.
• Mejoras: En comparación con la línea base SmolLM2-135M, la búsqueda produce reducciones de 5–7% en la pérdida de validación, 16–63% en energía por token y 16–59% en métricas de latencia (TTFT/TPOT).

Validación Escalada en rDXE

Los autores realizaron una búsqueda conjunta de arquitectura y chip en el sustrato multi-chip en anillo rDXE, reentrenando candidatos seleccionados en FineWeb-Edu-10BT (13.1B tokens) contra las líneas base SmolLM2 y Qwen:

• Variante de Precisión (347M parámetros): Logró la pérdida de validación más baja (2.798) y un rendimiento competitivo en benchmarks con menos parámetros que SmolLM2-360M y Qwen-0.5B.
• Variante de Energía (294M parámetros): Redujo la energía por token en un 40%.
• Variante de Latencia (365M parámetros): Redujo el Tiempo hasta el Primer Token (TTFT) y el Tiempo por Token de Salida (TPOT) en un 43%.

Significado y Afirmaciones

El artículo afirma que LLMForge demuestra que el NAS consciente del hardware puede producir frentes de Pareto condicionados al sustrato en lugar de una única arquitectura "mejor". Al desacoplar los parámetros de atención mediante IHA y utilizar un surrogado basado en codificador de alta fidelidad, el marco remodela efectivamente los perfiles de ancho y profundidad por capa para coincidir con cuellos de botella de hardware específicos.

Los autores posicionan a LLMForge como una contribución metodológica que permite el descubrimiento de pares de configuración de arquitectura y chip conscientes del despliegue. Señalan que, aunque las reducciones de energía y latencia son modestas, benefician principalmente el costo operativo y el perfil de carbono de la inferencia repetida en el borde. El trabajo se limita a escalas de menos de mil millones de parámetros ($\le$500M) debido al costo computacional de generar el corpus de entrenamiento para el surrogado, identificando la escalabilidad a modelos de mil millones de parámetros como un paso futuro.