Tapered Language Models

Resumo Técnico: Modelos de Linguagem Cônicos (Tapered Language Models)

Declaração do Problema
Os modelos de linguagem (LMs) modernos, incluindo Transformers, variantes recorrentes e arquiteturas baseadas em memória, compartilam um padrão estrutural comum: uma pilha de camadas idênticas onde os parâmetros são alocados uniformemente ao longo da profundidade. Apesar das crescentes evidências de que as camadas contribuem de forma não uniforme para o output final — onde as camadas iniciais frequentemente capturam padrões sintáticos superficiais e as camadas posteriores refinam os fluxos residuais em vez de transformá-los — essa alocação uniforme permanece o padrão. O artigo postula que essa abordagem de "tamanho único" falha em refletir a assimetria da importância das camadas, potencialmente desperdiçando a capacidade de parâmetros em camadas posteriores onde ela é menos utilizada.

Metodologia
Os autores introduzem os Modelos de Linguagem Cônicos (TLMs), um princípio arquitetural onde um componente portador de parâmetros é monotonicamente reduzido (tapered) ao longo da profundidade sob um orçamento total fixo.

• Componente Alvo: O estudo instancia este princípio na largura intermediária do MLP ($d_{ff}$). Os MLPs foram escolhidos porque dominam a contagem de parâmetros nos LMs modernos e expõem a largura como um eixo de variação limpo e único, independente do módulo de mistura de tokens.
• Cronogramas de Redução (Tapering Schedules): Em vez de alocações discretas por blocos, os autores propõem cronogramas de decaimento suaves e contínuos para $d_{ff}(l)$ através das camadas $l \in \{0, \dots, L-1\}$. Três cronogramas são avaliados:
  1. Linear: Taxa de decaimento constante.
  2. Cosseno (Cosine): Segue uma curva de meio cosseno, decaindo lentamente nas extremidades e abruptamente no meio.
  3. Sigmoide (Sigmoid): Concentra a transição em uma banda estreita em torno do ponto médio.
• Restrições: A redução é estritamente preservadora de orçamento. A largura média por camada é restringida para ser igual à largura uniforme de base ($\frac{1}{L}\sum d_{ff}(l) = d_{ff}^{baseline}$), garantindo que a contagem total de parâmetros e os FLOPs (tanto de treinamento quanto de inferência) permaneçam idênticos à linha de base uniforme.
• Configuração Experimental: O princípio é avaliado em quatro arquiteturas distintas (Transformer Padrão, Gated Attention, Hope-attention e Titans) e três escalas de modelo (440M, 760M e 1.3B de parâmetros). O estudo utiliza uma configuração de treinamento fixa (tokenizador Llama 3, comprimento de sequência de 4K) e compara os modelos cônicos com as linhas de base uniformes em perplexidade de validação in-distribution, perplexidade out-of-distribution (WikiText, LAMBADA) e oito benchmarks de raciocínio de senso comum.

Principais Resultados

1. A Direção da Alocação Importa: Um experimento motivador em um Transformer de 440M particionou as camadas em grupos: iniciais, médios e tardios. Alocar mais capacidade para as camadas iniciais ("Wider-early") melhorou a perplexidade em 0,32 pontos sobre a linha de base uniforme. Inversamente, alocar mais capacidade para as camadas tardias ("Wider-late") degradou o desempenho em mais de 1,0 ponto, demonstrando que a direção da capacidade é crítica.
2. Cronograma e Razão Ótimos: Uma varredura de razões de largura e cronogramas no modelo de 440M revelou que o cronograma de Cosseno superou consistentemente os cronogramas Linear e Sigmoide. A configuração ideal foi encontrada em uma razão de largura início-fim de 1,50 $\to$ 0,50, alcançando uma perplexidade de validação de 14,44 (uma melhoria de 1,84 pontos sobre a linha de base uniforme de 16,28).
3. Generalização: Aplicar a configuração fixa ótima (Cosseno, 1.5/0.5) aos modelos de 760M e 1.3B em todas as quatro arquiteturas melhorou consistentemente a perplexidade e a precisão de raciocínio de senso comum subsequente. Os ganhos foram observados independentemente de o módulo de mistura de tokens usar atenção softmax, gated attention ou mecanismos recorrentes/baseados em memória.
4. Insight Mecanístico: A análise de transformers pré-treinados revelou que os outputs dos MLPs tornam-se progressivamente mais alinhados com o fluxo residual existente conforme a profundidade aumenta (medido via similaridade de cosseno). Isso sugere que os MLPs posteriores reforçam o conteúdo existente em vez de computar novas características. A redução (tapering) alinha a arquitetura com esse padrão ao reduzir a capacidade onde ela é menos usada (camadas posteriores) e realocá-la para as camadas iniciais, onde ela é mais ortogonal e útil.

Significância e Alegações
O artigo estabelece a alocação de capacidade consciente da profundidade como uma alavanca de design simples e agnóstica à arquitetura para modelos de linguagem. Os autores alegam que o tapering da largura do MLP oferece uma melhoria de desempenho "gratuita" — aumentando a perplexidade e a precisão de tarefas subsequentes sem aumentar a contagem de parâmetros ou o custo computacional.

A significância reside em desafiar a uniformidade herdada do chassi Transformer. Os resultados sugerem que o princípio do tapering diz respeito à alocação geral de parâmetros através da profundidade, e não a propriedades específicas dos mecanismos de atenção. Embora os autores observem que o cronograma específico ideal (Cosseno, 1.5/0.5) foi derivado de um modelo de 440M e pode não ser o ótimo global para todas as escalas, os ganhos consistentes em diversas arquiteturas e tamanhos indicam que o princípio é robusto. O artigo conclui que a alocação consciente da profundidade é uma "alavanca oculta" à vista de todos que pode ser aplicada a modelos de fundação além da modelagem de linguagem, incluindo visão e transformers multimodais.