Beyond N-gram: Data-Aware X-GRAM Extraction for Efficient Embedding Parameter Scaling

Resumo Técnico: Beyond N-gram (X-GRAM)

1. O Problema

O escalonamento de Grandes Modelos de Linguagem (LLMs) tradicionalmente depende do aumento da capacidade computacional (FLOPs), o que se torna cada vez mais custoso. Uma alternativa é o uso de tabelas de busca indexadas por tokens (memória externa), que permitem expandir a capacidade estática do modelo sem aumentar proporcionalmente o custo computacional.

No entanto, as abordagens anteriores (como MoLE, Engram, Retoken) enfrentam limitações críticas:

• Ineficiência de Parâmetros (Cauda Longa): Devido à distribuição de Zipf dos tokens, a maioria das linhas da tabela (tokens de baixa frequência) recebe poucas atualizações durante o treinamento, permanecendo sub-otimizadas ("under-trained").
• Colapso de Slots (Redundância): Adicionar mais slots de busca paralelos não garante diversidade. Como os slots são treinados em eventos sobrepostos, eles tendem a colapsar em subespaços redundantes, oferecendo pouca informação adicional por byte de memória.
• Custo de Memória: Tabelas grandes exigem tráfego de memória massivo, muitas vezes forçando o uso de memória de host ou SSD, o que cria um gargalo de latência difícil de esconder durante a inferência.

2. Metodologia: X-GRAM

O X-GRAM é um framework de injeção de tokens consciente dos dados, projetado para extrair e injetar informações de forma eficiente. Ele transforma a busca simples de 1-gram (token único) em características locais variáveis de x-gram (n-gram) através de três etapas principais:

A. Mapeamento de Hash Híbrido Consciente da Frequência (Frequency-Aware Hash Mapping)
Para resolver o problema da cauda longa e do desequilíbrio de atualizações:

• Reserva VIP: Tokens de alta frequência (cabeça da distribuição) recebem linhas físicas dedicadas em uma região reservada, garantindo que sejam bem treinados.
• Bucketização Balanceada: Os tokens restantes são agrupados em buckets lógicos baseados em uma massa suavizada da frequência ($p(\omega)^\alpha$), garantindo que cada bucket tenha uma carga de atualização equilibrada.
• Hash Local e Alias: Dentro de buckets densos, utiliza-se hashing local multi-caminho para reduzir colisões, e linhas não utilizadas são recicladas como "alias" para tokens de baixa frequência, melhorando a taxa de utilização da memória.

B. Extração de Características com ShortConv (X-gram Extraction)
Para evitar o colapso de slots e adicionar contexto:

• Em vez de injetar vetores estáticos, o X-GRAM aplica um módulo leve de ShortConv (Convolução de Curta Distância) com portas (gating) e normalização.
• Utiliza múltiplos tamanhos de kernel (ex: $k=2, 3, 5, 7$) para extrair características locais de diferentes escalas (2-gram, 3-gram, etc.) a partir da sequência recuperada.
• Isso transforma a representação estática do token em características contextuais ricas e diversificadas, quebrando a simetria entre slots paralelos.

C. Injeção Contextual Consciente da Profundidade (Contextual Injection)
Para alinhar a memória estática com o fluxo dinâmico do modelo:

• Localização: A injeção ocorre principalmente no fluxo de valores (Value Stream) da atenção, onde o custo de ativação é menor (devido ao GQA - Grouped Query Attention) e o impacto na semelhança é mínimo.
• Resíduos Inter-camadas: Adiciona-se opcionalmente um caminho de resíduo entre camadas para estabilidade.
• Portas Conscientes da Profundidade: O uso de gates ($g^{(m)}_\ell$) que variam conforme a profundidade da camada e o estágio de treinamento (warmup) garante que a injeção seja estável e adaptada à demanda heterogênea do modelo.

3. Contribuições Chave

1. Eixo de Escalonamento Focado em Memória: Propõe um caminho de escalonamento que desacopla a capacidade do modelo dos FLOPs, focando na eficiência da gestão de memória.
2. Análise de Falhas em Abordagens Anteriores: Demonstra empiricamente que o simples aumento de tabelas leva a colapso de representações e sub-treinamento, e que a solução requer diversificação de características e alocação inteligente de recursos.
3. Arquitetura Híbrida Eficiente: Combina alocação de memória baseada em frequência (VIP + Hash) com extração de características não lineares (ShortConv), permitindo que tabelas menores superem tabelas muito maiores de métodos concorrentes.
4. Estabilidade de Treinamento: Introduz técnicas de normalização, escalonamento de taxa de aprendizado esparsa e gates de aquecimento para estabilizar o treinamento de parâmetros indexados por tokens.

4. Resultados Experimentais

Os autores avaliaram o X-GRAM em modelos de 0.73B e 1.15B parâmetros, comparando com baselines fortes (Retoken, MoRT, Engram) e um Transformer vanilla.

• Desempenho Superior: O X-GRAM superou consistentemente os baselines.
  • Na escala de 0.73B, a configuração X-GRAM-50% (com tabela comprimida para 50% do tamanho) alcançou 48.5 de precisão média, superando o MoRT (45.8) e o Engram (47.2) em até 3.2 pontos, usando uma tabela significativamente menor.
  • Na escala de 1.15B, alcançou 50.8 de precisão média na configuração 4x, superando os baselines em 2.3 pontos.
• Eficiência de Memória: Com a configuração de 50% (tabela comprimida), o X-GRAM obteve ganhos maiores do que os métodos concorrentes com tabelas 2x a 4x maiores.
• Escalabilidade: Diferente dos baselines que tendem a estagnar (plateau) à medida que a capacidade da tabela aumenta, o X-GRAM continua a melhorar, indicando que a memória adicional está sendo convertida em características úteis e distintas, não redundantes.
• Estabilidade: O modelo manteve uma vantagem de perda (loss) consistente durante todo o treinamento e alcançou a qualidade da baseline com apenas 57.25% dos dados de treinamento.

5. Significado e Impacto

O X-GRAM representa um avanço significativo na arquitetura de modelos de linguagem aumentados por memória. Ao demonstrar que a eficiência da memória (qualidade da extração e alocação) é mais importante do que a quantidade bruta de parâmetros, o trabalho oferece um paradigma prático para o futuro:

• Permite escalar modelos sem aumentar exponencialmente o custo de inferência (FLOPs).
• Resolve o problema fundamental da "cauda longa" em tabelas de lookup, tornando o treinamento de grandes vocabulários viável e estável.
• Estabelece que a extração de características contextuais locais (n-grams) via convoluções leves é essencial para evitar a redundância em sistemas de memória externa.

Em suma, o X-GRAM transforma a memória indexada por tokens de um gargalo de eficiência em um eixo de escalonamento robusto e prático para a próxima geração de LLMs.