Pesquisadores da gigante de chips de inteligência artificial Nvidia recentemente publicaram uma inovadora técnica de otimização de arquitetura chamada "FFN Fusion" (Fusão FFN). Essa técnica visa melhorar significativamente a eficiência de inferência de grandes modelos de linguagem (LLMs) ao resolver o gargalo de cálculo serial inerente à arquitetura Transformer, abrindo caminho para a implantação mais ampla de aplicativos de IA de alto desempenho.
Nos últimos anos, os grandes modelos de linguagem demonstraram capacidades poderosas em processamento de linguagem natural, pesquisa científica e agentes de diálogo. No entanto, com o aumento contínuo do tamanho e da complexidade dos modelos, os recursos computacionais necessários para o processo de inferência também cresceram consideravelmente, levando a gargalos de eficiência. A arquitetura Transformer é a base dos LLMs, e seus mecanismos de atenção alternados e camadas de redes feedforward (FFNs) precisam processar as entradas sequencialmente. Essa estrutura serial inerente, quando o tamanho do modelo aumenta, aumenta significativamente o custo computacional e de comunicação entre as GPUs, reduzindo a eficiência e aumentando os custos de implantação. Esse problema é ainda mais acentuado em cenários que exigem a geração rápida de vários tokens (como assistentes de IA em tempo real).
Para enfrentar esse desafio, os pesquisadores da Nvidia propuseram a técnica de FFN Fusion. A ideia central é combinar camadas FFN consecutivas e com baixa interdependência em um FFN mais amplo. Os pesquisadores observaram que, após a remoção das camadas de atenção, geralmente existem sequências longas de FFNs consecutivas em LLMs. Ao analisar essas sequências, eles descobriram que a dependência entre essas camadas FFN é pequena, portanto, elas podem ser executadas em paralelo.
A base matemática da FFN Fusion reside em concatenar os pesos de vários FFNs em série, criando assim um único módulo equivalente que pode ser calculado em paralelo. Por exemplo, se três FFNs são empilhados sequencialmente, com a saída de cada FFN servindo como entrada para o próximo, a FFN Fusion elimina essa dependência, permitindo que os três FFNs processem simultaneamente a mesma entrada e agreguem suas saídas. A análise teórica mostra que o FFN fundido mantém a mesma capacidade de representação do FFN original.
Ultra-253B-Base: Melhora dupla de desempenho e eficiência
Os pesquisadores da Nvidia aplicaram a técnica de FFN Fusion ao modelo Llama-3.1-405B-Instruct da Meta, criando um novo modelo chamado Ultra-253B-Base por meio de poda e reconstrução. Os resultados experimentais mostraram que o Ultra-253B-Base obteve melhorias significativas na velocidade de inferência e na eficiência de recursos. Especificamente, o modelo, com tamanho de lote 32, reduziu a latência de inferência em 1,71 vezes e o custo computacional por token em 35 vezes.
Ainda mais impressionante é que a melhoria na eficiência não se deu em detrimento da capacidade do modelo. O Ultra-253B-Base obteve excelentes resultados em vários benchmarks de avaliação de autoridade, como: MMLU 85,17%, MMLU-Pro 72,25%, HumanEval 86,58%, Arena Hard 84,92%, MT-Bench 9,19. Esses resultados são geralmente comparáveis ou até superiores aos do modelo original de 405 bilhões de parâmetros, enquanto o Ultra-253B-Base contém apenas 253 bilhões de parâmetros. Além disso, o consumo de memória do modelo também foi reduzido pela metade, graças à otimização do kv-cache.
Os pesquisadores usaram a distância do cosseno para analisar as saídas entre as camadas FFN, a fim de identificar regiões com baixa interdependência, que são as melhores opções para fusão. A técnica de FFN Fusion foi validada em modelos de diferentes escalas (incluindo 49 bilhões, 70 bilhões e 253 bilhões de parâmetros), demonstrando sua boa generalidade.
Esta pesquisa mostra que, por meio de análise aprofundada e design de arquitetura inteligente, é possível melhorar significativamente a eficiência dos LLMs. A FFN Fusion estabelece uma base para o design de LLMs mais paralelizados e adaptados ao hardware. Embora a paralelização de módulos Transformer completos enfrente mais desafios devido à maior dependência entre camadas, o sucesso da FFN Fusion certamente aponta uma direção importante para a otimização da eficiência de LLMs futuros.
Artigo: https://arxiv.org/abs/2503.18908
