In-Place Tokenizer Expansion for Pre-trained LLMs

Resumo Técnico: Expansão de Tokenizador In-Place para LLMs Pré-treinados

Declaração do Problema
Os tokenizadores para Modelos de Linguagem de Grande Escala (LLMs) são tipicamente fixos no início do pré-treinamento, com a alocação do vocabulário refletindo a distribuição do corpus inicial. À medida que as prioridades de implantação mudam para incluir idiomas sub-representados, esses idiomas sofrem de alta "fertilidade de tokens" — eles são divididos em significativamente mais tokens por palavra em comparação com os idiomas para os quais o tokenizador foi otimizado. Essa fragmentação aumenta a latência de decodificação, o computo e o consumo de energia, particularmente para modelos de borda (on-device), onde as matrizes de embedding e da cabeça de linguagem (LM-head) constituem uma parcela material da largura de banda por token na decodificação. Enquanto modelos em nuvem podem suportar vocabulários amplos, modelos de borda devem ser enviados com vocabulários compactos, aceitando a fragmentação fora de um conjunto de idiomas fixo. Remédios padrão, como o pré-treinamento do zero com um novo tokenizador ou o uso de métodos de transferência zero-shot para tokenizadores de terceiros arbitrários, são ou computacionalmente dispendiosos ou resolvem um problema mais difícil do que o necessário.

Metodologia
Os autores propõem a Expansão de Tokenizador, uma receita in-place para atualizar o tokenizador de um modelo pré-treinado quando o produtor do modelo controla o design. A abordagem consiste em três componentes principais:

1. Construção do Tokenizador (Continued-BPE):
   Em vez de treinar um novo tokenizador independentemente ou realizar a união de vocabulários, os autores continuam o treinamento de fusão (merge training) de Byte-Pair Encoding (BPE) do tokenizador de origem em um corpus multilíngue ponderado em direção a idiomas sub-tokenizados.

   • Propriedades: As fusões existentes são mantidas fixas; novas fusões são anexadas. Isso garante um mapeamento determinístico onde a maioria dos tokens de origem é transportada como tokens únicos, e cada novo token possui uma decomposição exata em um ou mais tokens de origem.
   • Escala: Aplicado a um modelo Mixture-of-Experts (MoE) de 8 bilhões de parâmetros, expandindo o vocabulário de 65K para 128K entradas ativas.

2. Inicialização de Embedding:
   A matriz de embedding é reinicializada para corresponder ao novo tamanho do vocabulário ($|V_{new}|$).

   • Tokens transportados: As linhas para tokens que mapeiam 1-para-1 da origem são copiadas sem alterações.
   • Novos tokens: As linhas para novos tokens são inicializadas como a média das linhas de embedding dos seus sub-tokens de origem constituintes. Isso garante que as novas linhas tenham uma norma comparável às linhas existentes sem exigir inicialização aleatória.

3. Adaptação em Dois Estágios:
   Para recuperar a qualidade do checkpoint de origem, o modelo passa por um pipeline de treinamento específico:

   • Estágio 1 (Treinamento Apenas de Embedding): O corpo do modelo é congelado. Apenas as novas linhas de embedding (e seus correspondentes acoplados na LM-head) são treinadas por 600B de tokens em uma mistura multilíngue focada em idiomas sub-tokenizados. Crucialmente, as linhas copiadas da origem são congeladas para evitar que se desviem do corpo congelado.
   • Estágio 2 (Pré-treinamento Continuado Total): Todos os parâmetros são descongelados. O modelo continua o pré-treinamento por 400B de tokens em um corpus multilíngue equilibrado. Este estágio é crítico para re-equilibrar o modelo e integrar o novo vocabulário sem regredir o desempenho nos idiomas previamente suportados.
   • Mid-Training/Post-Training: O checkpoint adaptado prossegue através das fases padrão de mid-training (extensão de contexto para 128K) e pós-treinamento.

Resultados Principais
O método foi aplicado ao checkpoint LFM2-8B-A1B para produzir o LFM2.5-8B-A1B.

• Eficiência de Tokens: O tokenizador expandido reduz significativamente a contagem de tokens para idiomas sub-tokenizados, mantendo a paridade para inglês e código.
  • Reduções: Hindi (2,4× menos tokens), Vietnamita (2,6×), Tailandês (~4,0×).
• Recuperação de Qualidade:
  • Troca Zero-shot: Simplesmente trocar o tokenizador causou uma queda de ~5,8 pontos agregados em benchmarks.
  • Estágio 1: Recuperou ~4,8 pontos, restaurando amplamente o desempenho em inglês e código.
  • Estágio 2: Fechou a lacuna restante e superou o modelo de origem em 3,6 pontos agregados, igualando ou excedendo a origem em categorias de Conhecimento, Matemática, Código e Multilíngue.
  • Benchmarks Multilíngues: Idiomas anteriormente suportados (ex: Árabe, Alemão, Japonês) mantiveram ou melhoraram levemente o desempenho, enquanto idiomas sub-tokenizados viram ganhos substanciais (ex: Vietnamita +11,6 pontos, Hindi +7,6 pontos no Global-MMLU).
• Latência em Dispositivos (On-Device):
  • Trade-off: Um vocabulário maior aumenta o custo de decodificação por token (medido como uma queda de ~7–10% em tokens/seg em dispositivos de referência devido a restrições de largura de banda de memória).
  • Ganho Líquido de Velocidade: Ao combinar a compressão de tokens com o custo por token, o tokenizador expandido gera um aumento de velocidade de decodificação de 2,2× a 3,7× por caractere para Tailandês, Bengali, Vietnamita e Hindi em dispositivos de referência (M4 Max, Snapdragon 8 Elite).
  • Regressão: Para idiomas com compressão de paridade próxima (Inglês, código, principais línguas europeias), o tempo de decodificação por caractere regride em até ~9%.

Achados Negativos e Ablações
O artigo relata dois achados negativos críticos que moldaram a receita final:

1. Treinar o Embedding Total no Estágio 1: Treinar todas as linhas de embedding (incluindo as linhas copiadas da origem) enquanto o corpo estava congelado causou uma degradação nas capacidades generativas (ex: colapso em MATH500, HumanEval+) e repetição severa, embora a precisão de múltipla escolha permanecesse intacta. Restringir o treinamento apenas às novas linhas foi necessário.
2. Mistura de Dados do Estágio 2: Usar uma mistura predominantemente em inglês para o pré-treinamento continuado do Estágio 2 causou um "colapso generativo" (perda de raciocínio e geração multilíngue), apesar de manter ou melhorar as pontuações de múltipla escolha. Uma mistura multilíngue equilibrada era necessária para preservar a qualidade generativa.

Significância e Alegações
Os autores afirmam que este é o primeiro relato de uma receita do lado do modelo para expansão de tokenizador via continued-BPE in-place em escala de produção (especificamente para um MoE de 8B para dispositivos).

• Eficiência: Demonstra que a expansão de vocabulário pode ser alcançada sem descartar o computo acumulado do pré-treinamento, desde que o tokenizador seja construído como uma continuação das fusões de origem.
• Implantação: Oferece uma solução prática para modelos de borda suportarem idiomas adicionais de forma eficiente, trocando um pequeno custo uniforme por token por ganhos massivos de eficiência em scripts específicos.
• Insight de Avaliação: O trabalho destaca que benchmarks de múltipla escolha padrão (como MMLU) podem falhar em detectar a degradação nas capacidades generativas durante a adaptação de vocabulário, necessitando da inclusão de benchmarks generativos na avaliação.

Os autores disponibilizam os pesos do modelo (LFM2.5-8B-A1B) e o tokenizador expandido, observando que a receita é especificamente desenhada para cenários onde o produtor do modelo controla a atualização do tokenizador, em vez de adaptar-se a tokenizadores de terceiros arbitrários.