Investigadores de la empresa líder en chips de inteligencia artificial, NVIDIA, han publicado recientemente una innovadora técnica de optimización de arquitectura denominada "Fusión FFN" (FFN Fusion). Esta técnica tiene como objetivo mejorar significativamente la eficiencia de la inferencia de los modelos lingüísticos grandes (LLM) al resolver el cuello de botella de cálculo en serie inherente a la arquitectura Transformer, allanando el camino para una implementación más amplia de aplicaciones de IA de alto rendimiento.
En los últimos años, los modelos lingüísticos grandes han demostrado una gran capacidad en áreas como el procesamiento del lenguaje natural, la investigación científica y los agentes conversacionales. Sin embargo, a medida que aumenta el tamaño y la complejidad de los modelos, también lo hace la cantidad de recursos computacionales necesarios para su proceso de inferencia, lo que genera un cuello de botella en la eficiencia. La arquitectura Transformer, base de los LLM, requiere el procesamiento secuencial de la entrada debido a sus mecanismos de atención y capas de redes feed-forward (FFN) alternas. Esta estructura en serie inherente aumenta significativamente los costes de cálculo y de comunicación entre las GPU a medida que aumenta el tamaño del modelo, reduciendo la eficiencia y aumentando los costes de implementación. Este problema es especialmente evidente en escenarios que requieren la generación rápida de varios tokens (como en asistentes de IA en tiempo real).
Para abordar este desafío, los investigadores de NVIDIA propusieron la técnica de Fusión FFN. La idea principal de este método es fusionar las capas FFN consecutivas con baja dependencia mutua en un FFN más ancho. Los investigadores observaron que, después de eliminar las capas de atención, suelen existir secuencias largas de FFN consecutivas en los LLM. Al analizar estas secuencias, descubrieron que la dependencia entre estas capas FFN es mínima, por lo que se pueden ejecutar en paralelo.
La base matemática de la Fusión FFN radica en concatenar los pesos de múltiples FFN en serie para crear un único módulo equivalente que se puede calcular en paralelo. Por ejemplo, si tres FFN están apilados secuencialmente, con la salida de cada FFN como entrada del siguiente, la Fusión FFN elimina esta dependencia, permitiendo que los tres FFN procesen simultáneamente la misma entrada y luego agreguen sus salidas. El análisis teórico demuestra que el FFN fusionado mantiene la misma capacidad de representación que el FFN original.
Ultra-253B-Base: Mejora del rendimiento y la eficiencia
Los investigadores de NVIDIA aplicaron la técnica de Fusión FFN al modelo Llama-3.1-405B-Instruct de Meta, creando un nuevo modelo llamado Ultra-253B-Base mediante poda y reconstrucción. Los resultados experimentales muestran una mejora significativa en la velocidad de inferencia y la eficiencia de los recursos. Específicamente, el modelo redujo la latencia de inferencia en 1,71 veces y el coste computacional por token en 35 veces con un tamaño de lote de 32.
Más impresionante aún, esta mejora en la eficiencia no se logró a expensas de la capacidad del modelo. Ultra-253B-Base obtuvo excelentes resultados en varios puntos de referencia de evaluación autorizados, como: MMLU 85,17 %, MMLU-Pro 72,25 %, HumanEval 86,58 %, Arena Hard 84,92 %, MT-Bench 9,19 %. Estos resultados son comparables o incluso superiores a los del modelo original de 405 mil millones de parámetros, mientras que Ultra-253B-Base solo contiene 253 mil millones de parámetros. Además, el consumo de memoria se redujo a la mitad gracias a la optimización de la caché kv.
Los investigadores utilizaron el análisis de la distancia del coseno entre las salidas de las capas FFN para identificar las regiones con baja interdependencia, que son las mejores opciones para la fusión. La técnica de Fusión FFN se ha probado en modelos de diferentes tamaños (incluidos modelos de 49 mil millones, 70 mil millones y 253 mil millones de parámetros), lo que demuestra su buena generalización.
Esta investigación muestra que, mediante un análisis profundo y un diseño de arquitectura inteligente, se puede mejorar significativamente la eficiencia de los LLM. La Fusión FFN sienta las bases para el diseño de LLM más paralelizados y adaptados al hardware. Aunque la paralelización de los módulos Transformer completos presenta mayores desafíos debido a la mayor dependencia entre capas, el éxito de la Fusión FFN sin duda indica una dirección importante para la optimización de la eficiencia de los LLM en el futuro.
Artículo:https://arxiv.org/abs/2503.18908
