Composer: A Search Framework for Hybrid Neural Architecture Design

O artigo apresenta o Composer, um framework de busca arquitetônica modular que descobre novas arquiteturas de modelos de linguagem híbridos, superando o Llama 3.2 em desempenho, eficiência e redução de perda de validação ao explorar automaticamente combinações de primitivas computacionais e extrapolar as melhores configurações para escalas maiores.

O essencial

Imagine que você é um chef de cozinha de renome mundial tentando criar o prato perfeito: um Modelo de Linguagem (como o Llama ou o GPT). Até hoje, a receita padrão era sempre a mesma: uma camada de "atenção" (que olha para o contexto da frase) seguida de uma camada de "perceptron" (que processa a informação), repetida 32 vezes. É como fazer um bolo onde você sempre coloca uma camada de chocolate e uma de baunilha, sempre na mesma ordem. Funciona bem, mas será que não existe uma combinação melhor?

O problema é que tentar descobrir essa combinação perfeita manualmente é como tentar encontrar uma agulha num palheiro, mas o palheiro é gigante e cada tentativa de assar um bolo novo custa milhões de dólares em energia e tempo.

É aqui que entra o Composer, o "astrofísico da culinária" criado pelos pesquisadores do Meta e da Universidade do Texas.

O Que é o Composer?

O Composer é um sistema de busca inteligente que não tenta cozinhar o bolo gigante de uma vez. Em vez disso, ele faz o seguinte:

1. Cozinha em Miniatura (Escala Pequena): Em vez de tentar assar um bolo de 8 bilhões de ingredientes (o tamanho de um modelo grande), o Composer assa mini-bolos de apenas alguns milhões de ingredientes. É como testar a receita em uma panela pequena antes de encher a assadeira industrial.
2. Experimenta Combinações Loucas: Ele mistura os ingredientes de formas que os humanos nunca pensariam. Em vez de "Chocolate-Baunilha-Chocolate-Baunilha", ele testa "Chocolate-Chocolate-Baunilha-Chocolate-Baunilha-Baunilha".
3. Adivinha o Futuro (Extrapolação): Aqui está a mágica. O Composer descobre que, se você encontrar o mini-bolo perfeito, você pode "esticar" ou "empilhar" essa receita para fazer o bolo gigante sem perder o sabor. É como descobrir que a proporção perfeita de farinha e açúcar em uma xícara de chá funciona perfeitamente se você multiplicar os ingredientes por 1.000 para fazer um balde gigante.

Como Funciona a Busca? (As 4 Etapas)

O Composer tem quatro "cozinheiros" principais trabalhando juntos:

• O Motor de Busca (Search Engine): Ele é o explorador. Ele usa uma técnica chamada "Otimização Bayesiana" (basicamente, um palpite matemático muito esperto) para testar milhares de combinações de camadas rapidamente. Ele descobre que, às vezes, colocar mais camadas de "processamento" (MLP) do que de "atenção" (Attention) funciona melhor.
• O Avaliador (Evaluator): Ele prova o mini-bolo. Mas provar em um palito gigante é caro e demorado. Então, o Composer usa "sabores sintéticos" (dados pequenos e artificiais, como histórias de crianças ou jogos de palavras) que são baratos e rápidos de testar, mas que dizem se a receita vai funcionar no mundo real.
• O Agregador (Aggregator): Imagine que você testou 100 mini-bolos e 10 deles ficaram ótimos. O Agregador pega os melhores pedaços de cada um e cria um "Super-Bolo" definitivo. Ele não escolhe apenas o melhor de cada vez, mas mistura as melhores partes para criar uma receita que não tenha "ruído" ou defeitos de uma única tentativa.
• O Extrapolador (Extrapolator): Este é o mágico que transforma o mini-bolo em um gigante. Ele usa duas técnicas:
  • Empilhamento (Stacking): Repetir o bloco de receita várias vezes.
  • Esticamento (Stretching): Aumentar o tamanho de cada camada mantendo a proporção perfeita descoberta.

O Que Eles Descobriram?

Ao usar o Composer, eles criaram novos modelos chamados "Composite". A descoberta principal foi uma mudança na proporção dos ingredientes:

• A Velha Regra: 1 camada de Atenção para 1 camada de Processamento (1:1).
• A Nova Regra do Composer: 1 camada de Atenção para 2 camadas de Processamento (1:2).

Além disso, eles descobriram que a ordem importa. Os melhores modelos começam com camadas de Atenção (para entender o contexto) e terminam com camadas de Processamento (para refinar a resposta), com uma mistura inteligente no meio.

Por Que Isso é Importante?

Os resultados foram impressionantes comparados ao famoso Llama 3.2:

• Mais Inteligentes: Os modelos do Composer erram menos e entendem melhor as tarefas de linguagem.
• Mais Rápidos: Como têm menos camadas de "atenção" (que são pesadas e lentas), eles processam informações mais rápido.
• Mais Baratos: Treinar e usar esses modelos gasta menos energia e memória.

A Analogia Final

Pense no Llama 3.2 como um carro de corrida projetado manualmente por engenheiros brilhantes, mas que usa sempre o mesmo tipo de motor e pneus.

O Composer é como um laboratório de IA que testa milhares de combinações de motores e pneus em uma pista de brinquedo (escala pequena). Ele descobre que um motor V8 com pneus de chuva específicos (a proporção 1:2) é o vencedor na pista de brinquedo. Então, ele usa uma fórmula mágica para escalar essa descoberta para um carro de F1 gigante, e o resultado é um carro que é mais rápido, mais econômico e mais seguro do que qualquer carro feito manualmente antes.

Em resumo, o Composer não apenas melhora os modelos de IA; ele muda a forma como os criamos, passando de "tentativa e erro manual" para "descoberta sistemática e inteligente".

Resumo técnico

Resumo Técnico: Composer

1. O Problema

As arquiteturas de modelos de linguagem grandes (LLMs) baseadas em Transformers (com camadas sequenciais fixas de Atenção e MLP) têm sido o padrão-ouro. No entanto, pesquisas recentes mostram que arquiteturas híbridas, que alteram a proporção e a intercalação de primitivas computacionais (como Atenção, MLP e Modelos de Espaço de Estado - SSM), podem superar os Transformers tradicionais.

O desafio principal é que o espaço de design dessas arquiteturas híbridas é vastíssimo e a busca manual é ineficiente. Além disso, realizar uma busca de arquitetura neural (NAS) diretamente em escala pré-treinada é proibitivamente caro. Métodos existentes de NAS muitas vezes assumem intercalações fixas ou não conseguem extrapolar com sucesso os resultados de modelos pequenos para grandes escalas, falhando em capturar a qualidade do modelo quando escalado.

2. Metodologia: O Framework Composer

Os autores propõem o Composer, um framework de busca de arquitetura neural híbrida (HNAS) que descobre automaticamente arquiteturas híbridas de alto desempenho em pequena escala e as extrapola para grandes escalas. O Composer é composto por quatro módulos principais:

1. Motor de Busca (Search Engine):

   • Explora o espaço de design de modelos pequenos (ex: milhões de parâmetros) para encontrar arranjos ótimos de primitivas (Atenção e MLP).
   • Utiliza Otimização Bayesiana (com modelos substitutos Gaussian Process) para eficiência na amostragem.
   • Propõe três estratégias de busca:
     • One-Shot: Busca em um número fixo de camadas (ex: 6 ou 16) e extrapola.
     • Busca Incremental de Camadas Finais: Constrói o modelo camada por camada, fixando as anteriores.
     • Busca Incremental de Camadas Intermediárias: Fixa as camadas iniciais e finais, buscando apenas o meio.
   • Escalonamento de Largura: Reduz a largura das primitivas durante a busca para diminuir o custo computacional, mantendo a proporção largura/profundidade do modelo alvo.

2. Avaliador (Evaluator):

   • Treina e avalia os candidatos em datasets pequenos para fornecer sinais rápidos de qualidade.
   • Descoberta Crítica: O uso de datasets web-scale amostrados (como DCLM reduzido) ou modelos muito pequenos não reflete bem o desempenho em grande escala.
   • Solução: O Composer utiliza o MAD (um dataset sintético de manipulação de tokens) como proxy. O MAD é barato de treinar, aprendível por modelos pequenos devido ao vocabulário reduzido, mas altamente correlacionado com tarefas de LLM em grande escala.

3. Agregador (Aggregator):

   • Sintetiza os resultados da busca (múltiplos candidatos promissores) em uma única arquitetura final.
   • Utiliza uma técnica de agrupamento (clustering) $N_0$: para cada camada, seleciona a primitiva mais frequente entre os melhores candidatos, independentemente das camadas anteriores. Isso suaviza o ruído e o sobreajuste que podem ocorrer na busca em pequena escala.

4. Extrapolador (Extrapolator):

   • Escala a arquitetura encontrada (ex: 6 ou 16 camadas) para o tamanho desejado (ex: 3B ou 8B parâmetros).
   • Técnicas de Escala:
     • Esticamento (Stretching): Mantém o padrão de intercalação e aumenta o número de camadas de cada grupo proporcionalmente. Funciona melhor para buscas de 16+ camadas.
     • Empilhamento (Stacking): Repete o bloco encontrado sequencialmente. Funciona bem para buscas de 6 camadas.

3. Contribuições Principais

• Framework Automatizado: Primeiro framework sistemático para descobrir arquiteturas híbridas de pré-treinamento de LLMs que superam o estado da arte, sem depender de intuição manual.
• Descoberta de Arquiteturas "Composite": Identificação de novas arquiteturas com proporções de Atenção:MLP de 1:2 (uma camada de Atenção para duas de MLP) e padrões de intercalação não uniformes (ex: blocos de 2A+4M ou arranjos mais complexos como 2A+5M+2A+3M+1A+3M).
• Validação de Proxy Sintético: Demonstração de que datasets sintéticos como o MAD são superiores a datasets web-scale amostrados para guiar a busca de arquitetura em pequena escala.
• Estratégias de Escala: Validação de que a combinação de busca em pequena escala com largura reduzida e técnicas de extrapolação (esticamento/empilhamento) preserva a qualidade do modelo ao aumentar em ~1000x o tamanho.

4. Resultados

Os modelos descobertos pelo Composer (chamados de Composite LLMs) foram comparados ao Llama 3.2 e a outros trabalhos de ponta (como Sandwich Transformer, Striped Attention e STAR) em várias escalas (350M a 8B parâmetros).

• Desempenho:
  • Redução consistente na perda de validação (Validation Loss) em comparação ao Llama 3.2 (redução de 0.03 a 0.05).
  • Melhoria de 2% a 2.1% em média na precisão de tarefas downstream (como ARC, HellaSwag, PIQA, WinoGrande).
  • Superação de todas as arquiteturas híbridas anteriores de ponta (SOTA) quando treinadas com o mesmo número de tokens.
• Eficiência:
  • Devido à menor proporção de camadas de Atenção (1:2 vs 1:1 no Transformer padrão), os modelos Composite têm menos camadas de Atenção totais.
  • Throughput de Treinamento: Aumento de 1.25x.
  • Latência de Inferência: Redução de 1.33x.
  • Tamanho do KV Cache: Redução de 1.69x, permitindo sequências mais longas ou maior batch size com a mesma memória.
• Robustez: A correlação de rank entre a busca em pequena escala e o desempenho em grande escala foi de 0.97, indicando que o framework identifica corretamente as melhores arquiteturas mesmo antes do escalonamento.

5. Significado e Impacto

O trabalho do Composer representa um avanço significativo na automação do design de LLMs. Ele demonstra que:

1. A busca de arquitetura não precisa ser feita apenas em grandes escalas para ser eficaz, desde que se use proxies inteligentes (como o MAD) e técnicas de extrapolação adequadas.
2. As arquiteturas híbridas com proporções de primitivas não padrão (especificamente mais camadas de MLP do que de Atenção) oferecem um melhor equilíbrio entre qualidade, custo computacional e eficiência de inferência.
3. O framework é extensível, permitindo a inclusão de outras primitivas (como Mamba, Delta Net, etc.) no futuro, abrindo caminho para uma nova geração de modelos de linguagem mais eficientes e performáticos.

Em resumo, o Composer fornece um "mapa" metodológico para substituir a intuição humana na criação de arquiteturas de LLMs, resultando em modelos que são simultaneamente mais inteligentes, mais rápidos e mais baratos de treinar e executar.