TiTok: Transfer Token-level Knowledge via Contrastive Excess to Transplant LoRA

1. Problema

Os Grandes Modelos de Linguagem (LLMs) são amplamente utilizados, mas o fine-tuning completo é computacionalmente caro e exige muito armazenamento. Métodos de Fine-Tuning Eficiente em Parâmetros (PEFT), como o LoRA (Low-Rank Adaptation), mitigam esses custos atualizando apenas um pequeno subconjunto de parâmetros. No entanto, uma limitação crítica do LoRA é que seus adaptadores são dependentes do modelo base (backbone) em que foram treinados e não podem ser transferidos diretamente para diferentes arquiteturas ou versões de modelos.

Abordagens existentes para contornar isso incluem:

• Distilação de Conhecimento (KD): Geralmente depende de dados de treinamento reais do domínio alvo, que muitas vezes não estão disponíveis.
• TransLoRA: Utiliza dados sintéticos gerados por LLMs para transferir conhecimento, mas exige o treinamento de um modelo discriminador adicional para filtrar dados de baixa qualidade, o que aumenta a complexidade e o custo computacional.

O objetivo do TITOK é permitir a transplantação eficiente de LoRA entre modelos diferentes (mesma família, famílias distintas, tamanhos diferentes ou versões diferentes) sem dados reais e sem modelos auxiliares complexos.

2. Metodologia: TITOK

O TITOK (Transfer Token-Level Knowledge via Contrastive Excess) propõe um framework que transfere conhecimento através de sinais ao nível de token, em vez de tratar sequências inteiras como unidades indivisíveis. O processo consiste em três etapas principais:

A. Geração de Dados Sintéticos

O modelo "especialista" (o modelo base fonte + o LoRA treinado) gera pares de query-resposta (pergunta e rótulo) para a tarefa alvo, utilizando poucos exemplos (few-shot) como semente. Diferente do TransLoRA, que usa o modelo alvo para gerar consultas, o TITOK usa o modelo especialista para gerar tanto a consulta quanto o rótulo, garantindo maior coerência na distribuição dos dados.

B. Cálculo de "Excesso Contrastivo" (Contrastive Excess)

Esta é a inovação central do método. O framework compara as previsões de dois modelos:

1. Amador: O modelo base fonte (sem LoRA).
2. Especialista: O modelo base fonte + LoRA.

Para cada token $y_i$ na resposta gerada, calcula-se uma pontuação de excesso ($S(y_i)$):
$$S(y_i) = \log P_{\text{Especialista}}(y_i | \text{contexto}) - \log P_{\text{Amador}}(y_i | \text{contexto})$$

• Interpretação: Tokens com alto excesso indicam onde o LoRA injetou conhecimento específico da tarefa que o modelo base sozinho não conseguia capturar com confiança. Esses tokens são os mais informativos para a transferência.
• Fundamento Teórico: A pontuação é análoga a uma Razão de Verossimilhança Logarítmica (LLR), um estatístico robusto para identificar diferenças entre distribuições de probabilidade.

C. Filtragem e Seleção de Treinamento

Com as pontuações de excesso calculadas, o TITOK aplica um filtro de dois níveis:

1. Filtragem de Amostra (Sample Filtering): Calcula a pontuação média de excesso para cada amostra sintética. Apenas as $M$ amostras com as maiores pontuações médias são mantidas, descartando exemplos pouco informativos.
2. Seleção de Token (Token Selection): Dentro das amostras mantidas, apenas os tokens com as pontuações de excesso mais altas (top-$k\%$) são utilizados para treinar o novo LoRA no modelo alvo. Isso foca o aprendizado nas regiões onde o conhecimento transferível está concentrado.

D. Alinhamento de Tokenizadores

Para lidar com transferências entre modelos com tokenizadores diferentes (ex: Mistral para Llama), o TITOK utiliza um algoritmo de alinhamento de ponteiros duplos. Ele mapeia os tokens do modelo fonte para o modelo alvo e propaga as máscaras de seleção (baseadas no excesso) através de regras de cópia, replicação e média, permitindo o treinamento mesmo com vocabulários distintos.

3. Principais Contribuições

• Eliminação de Modelos Auxiliares: Ao contrário do TransLoRA, o TITOK não requer o treinamento de um discriminador, reduzindo a complexidade e o custo computacional.
• Transferência Baseada em Token: Em vez de transferir sequências inteiras, o método identifica e transfere apenas os tokens críticos onde o LoRA adiciona valor, tornando o processo mais eficiente e preciso.
• Robustez a Dados Externos: O método demonstra eficácia mesmo quando utiliza dados externos (não sintéticos) ou dados de domínios ligeiramente diferentes, provando que não depende exclusivamente da qualidade perfeita dos dados sintéticos.
• Alinhamento de Tokenizadores: Uma solução prática para permitir a transferência entre modelos com arquiteturas de tokenização incompatíveis.

4. Resultados Experimentais

Os autores avaliaram o TITOK em três benchmarks principais: Big-Bench Hard (BBH), MMLU (tarefas de raciocínio e conhecimento geral) e LaMP (tarefas de personalização e geração de texto).

• Desempenho Geral: O TITOK superou consistentemente as linhas de base (baselines), incluindo modelos Vanilla (sem ajuste), Distilação de Conhecimento (KD) e TransLoRA.
  • Ganho médio de +4% a +10% em relação às linhas de base.
  • Especificamente, superou o TransLoRA em +4,4% e o KD em +8,5% em média.
• Cenários de Transferência:
  • Mesmo Modelo: Transferência dentro da mesma família (ex: Mistral v0.3 para Mistral v0.3) mostrou ganhos significativos.
  • Famílias Diferentes: Transferência entre Mistral e Llama demonstrou a capacidade de generalização arquitetural.
  • Tamanhos Diferentes: Transferência de modelos menores para maiores (ex: Llama 3B para 8B) e vice-versa foi eficaz.
  • Versões Diferentes: Transferência entre versões de modelos (ex: Llama 2 para Llama 3) manteve a eficácia.
• Análise de Componentes: Estudos de ablação confirmaram que tanto a filtragem de amostras quanto a seleção de tokens são essenciais para o desempenho final.
• Eficiência Computacional: O TITOK reduziu o tempo total de computação em 1,5x a 2,5x comparado ao TransLoRA, principalmente devido à eliminação do treinamento do discriminador.

5. Significado e Conclusão

O TITOK representa um avanço paradigmático na transferência de conhecimento para PEFT. Ao focar em sinais ao nível de token derivados da diferença entre um modelo base e sua versão adaptada, o método consegue "transplantar" o conhecimento do LoRA de forma seletiva e eficiente.

Sua principal vantagem é a simplicidade e eficiência: remove a necessidade de dados reais de treinamento e de modelos auxiliares complexos, tornando a adaptação de LLMs mais acessível e escalável para cenários do mundo real onde os dados são escassos ou privados, e onde a diversidade de modelos base exige flexibilidade na transferência de conhecimento.