LLMs Eficientes: Arquitetura e Leis de Escala

Introdução: Navegando o Trade-off de Eficiência em LLMs

Como Engenheiro de Software Sênior e Arquiteto de Soluções na AITY, acompanho de perto a evolução dos Large Language Models (LLMs). É inegável que eles se tornam cada vez maiores e melhores. No entanto, o custo de operá-los — seja gerando texto, respondendo perguntas ou impulsionando aplicações em tempo real — também escala. Para aplicações web baseadas em IA em tempo real, a precisão do modelo é vital, mas não pode vir à custa da eficiência.

Em um artigo que apresentamos na International Conference on Learning Representations (ICLR), oferecemos um framework para navegar nesse trade-off entre precisão e eficiência. Conectamos diretamente as leis de escala às decisões de design arquitetônico.

A Lacuna nas Leis de Escala Atuais

Em 2022, o Google DeepMind publicou um estudo com o LLM experimental Chinchilla, demonstrando uma lei de escala que permitia a otimização conjunta do tamanho do modelo e dos dados de treinamento para atingir um nível de perda desejado, dado um orçamento computacional específico. Esta "lei de Chinchilla" relaciona a perda do modelo (L) ao número de parâmetros do modelo (N) e ao número de tokens no conjunto de dados de treinamento.

Apesar de sua importância, essa lei não especifica escolhas arquitetônicas, como o "tamanho oculto" (hidden size) das representações internas do modelo, ou a alocação relativa de parâmetros entre camadas de atenção e camadas Multilayer Perceptron (MLP). Dois modelos com o mesmo número de parâmetros (por exemplo, um bilhão), treinados nos mesmos dados e com a mesma precisão, podem diferir em até 40% no throughput em tempo de inferência, dependendo dessas escolhas arquitetônicas adicionais. Nosso trabalho visou deduzir leis de escala que pudessem prever essas escolhas.

A Arquitetura Transformer

A arquitetura Transformer, que está no cerne de todos os LLMs, consiste principalmente em blocos empilhados de atenção e MLP.

• Blocos de Atenção: Determinam o peso a ser dado a cada token anterior ao atualizar a representação do token atual.
• Blocos MLP: Transformam essa representação, armazenando grande parte do conhecimento aprendido pelo modelo.
• Mecanismo de Atenção: Utiliza três matrizes — Query, Key e Value — que codificam o que cada token procura, o que cada token tem a oferecer e o conteúdo que cada token pode contribuir. A comparação de queries com keys informa ao modelo a relevância de cada token para os outros.
• Atenção Multi-cabeça (Multihead Attention): A maioria dos LLMs utiliza múltiplas computações de atenção rodando em paralelo, cada uma com suas próprias projeções de query, key e value. Isso permite que diferentes cabeças se especializem em diversos aspectos da entrada.

É fundamental entender que a arquitetura não é um detalhe. As configurações corretas podem desbloquear grandes ganhos de eficiência sem custo na precisão.

Nossa Abordagem: Arquitetura como Variável de Primeira Classe

Em nosso artigo da ICLR, introduzimos uma lei de escala que amplia o framework Chinchilla com três fatores arquitetônicos críticos:

• Tamanho Oculto (d_model): A dimensão dos vetores que fluem pelos blocos de embedding, atenção e MLP.
• Razão MLP-para-Atenção (r_mlp/attn): A proporção do número de parâmetros MLP em relação ao número de parâmetros de atenção.
• Atenção de Consulta Agrupada (GQA - Grouped-Query Attention): Grupos de cabeças de atenção compartilham matrizes de key e value, enquanto preservam matrizes de query distintas.

Cada fator impacta diretamente o throughput de inferência:

• Tamanho Oculto (d_model): Sob um orçamento de parâmetros fixo, tamanhos ocultos maiores reduzem as operações FLOPs totais de inferência e diminuem o cache de key-value, melhorando o throughput.
• Razão MLP-para-Atenção (r_mlp/attn): Uma razão maior aloca mais parâmetros para o MLP e menos para a atenção, encolhendo o cache de key-value e reduzindo gargalos de largura de banda de memória.
• Grouped-Query Attention (GQA): Comprimir cabeças de key-value reduz ainda mais os custos de entrada/saída durante a geração.

É importante notar que ajustar esses fatores puramente para maior throughput pode ter um custo na precisão. Tanto o tamanho oculto quanto a razão MLP-para-atenção exibem curvas de perda em forma de "U", indicando um ponto ótimo. O GQA, por ter um efeito mais errático na perda, é tratado como um hiperparâmetro discreto ajustado por busca local.

Nossa lei de escala é deduzida em duas etapas: 1. Ajustamos a lei de Chinchilla padrão ao modelo em investigação, calculando os coeficientes (E, A, B, α, β) para estabelecer uma perda de referência ótima. 2. Calibramos como cada escolha arquitetônica (diferenças nos três fatores) afeta essa perda, aprendendo uma superfície de correção sobre o espaço de design. Os efeitos de tamanho oculto e razão MLP-para-atenção na perda se mostraram separáveis, permitindo a otimização independente de cada fator.

Duas Famílias de Modelos: Panda e Surefire

Essa lei de escala nos permitiu desenvolver um framework de busca que identifica arquiteturas Pareto-ótimas para qualquer alvo de precisão. O resultado dessa busca foram duas famílias de modelos:

• Panda: Maximiza a precisão.
• Surefire: É Pareto-ótimas na fronteira de precisão-eficiência.

Para validar o framework, treinamos mais de 200 modelos com arquiteturas variadas (de 80 milhões a três bilhões de parâmetros, oito bilhões a 100 bilhões de tokens). Os resultados dos experimentos com modelos de um e três bilhões de parâmetros são notáveis:

• O modelo Panda de um bilhão de parâmetros ganha 2,1% de precisão sobre o LLaMA-3.2-1B, e o modelo de três bilhões de parâmetros ganha 0,6% sobre o LLaMA-3.2-3B, porém com menor throughput.
• Os modelos Surefire igualam ou excedem a precisão do LLaMA-3.2 enquanto melhoram o throughput em 12-47%. Ganhos de até 42% foram observados em A100 (vLLM) e 47% em H200 (SGLang), sob diferentes configurações de tamanho de modelo e tamanho de lote.

Impacto Prático na AITY

Nossa pesquisa na AITY revela insights cruciais com impacto direto no desenvolvimento e implantação de LLMs eficientes:

• Arquitetura não é um detalhe: A proporção MLP-para-atenção ideal para modelos LLaMA-3.2 é em torno de 1.0, significativamente menor do que as versões de peso aberto existentes (por exemplo, 4.8 para LLaMA-3.2-1B). Isso sugere que os modelos atuais superalocam recursos para camadas MLP. As configurações corretas de tamanho oculto, proporção MLP-para-atenção e GQA podem desbloquear grandes ganhos de eficiência sem custo na precisão.
• Pequenas experiências preveem grandes resultados: Nossa lei de escala condicional, calibrada em modelos com apenas 80 milhões a 297 milhões de parâmetros, prevê de forma confiável a melhor arquitetura em um e três bilhões de parâmetros. Isso permite uma exploração de baixo custo antes do treinamento completo e dispendioso.
• Generalização em hardware e sistemas de serving: Os ganhos de eficiência que identificamos são consistentes em GPUs A100/H200 e sistemas de serving vLLM/SGLang. Essa generalização torna os resultados diretamente acionáveis para as diversas infraestruturas de produção da AITY, garantindo que nossas aplicações de IA em tempo real sejam não apenas precisas, mas também economicamente viáveis e performáticas.

Esta abordagem estratégica na AITY nos capacita a construir e implantar LLMs de ponta que são simultaneamente poderosos e eficientes, garantindo que nossos clientes recebam o melhor desempenho ao menor custo operacional possível.