Resumo Técnico: Ajuste de Contexto para Otimização em Contexto

1. Declaração do Problema

Os Grandes Modelos de Linguagem (LLMs) demonstraram capacidades notáveis de adaptação de poucos disparos (few-shot) por meio do Aprendizado em Contexto (ICL), onde os modelos aproveitam informações relevantes para a tarefa a partir de exemplos de demonstração fornecidos no contexto de entrada sem atualizar os pesos. No entanto, o ICL depende de uma única passagem direta (forward pass) para formar uma representação de memória, frequentemente enfrentando dificuldades com raciocínio complexo ou mudanças de domínio quando o contexto inicial é insuficiente.

Por outrooso lado, métodos de adaptação baseados em prompts (ex: Prompt Tuning, Prefix Tuning) otimizam vetores treináveis para guiar o comportamento do modelo, mas tipicamente inicializam esses vetores aleatoriamente ou via instruções de alto nível, independentemente dos exemplos de demonstração específicos de poucos disparos. Embora o Treinamento em Tempo de Teste (TTT) tenha mostrado que uma adaptação eficaz é possível através do ajuste fino dos parâmetros do modelo (ex: via LoRA) em permutações de demonstração, isso acarreta custos computacionais significativos e atualiza os pesos do modelo em vez do próprio contexto.

O artigo identifica uma lacuna em unir eficientemente as forças complementares do ICL (extração de informação do contexto) e da adaptação baseada em prompts (otimização de representações via gradientes). O objetivo é desenvolver um método que inicialize uma representação treinável diretamente a partir de exemplos de poucos disparos e a refine através de otimização baseada em gradientes, alcançando alta precisão com maior eficiência do que abordagens de atualização de pesos.

2. Metodologia: Otimização em Contexto (ICO)

Os autores propõem um framework denominado Otimização em Contexto (ICO), que unifica abordagens que adaptam ou os pesos do modelo ou a representação do contexto no tempo de inferência usando pares de demonstração. Dentro deste framework, eles introduzem o Ajuste de Contexto (Context Tuning), um método que congela os parâmetros do modelo e, em vez disso, otimiza uma representação de contexto leve derivada das demonstrações de poucos disparos.

O Ajuste de Contexto é instanciado em duas variantes:

2.1 CT-Prompt

Esta variante aplica os princípios do Prompt Tuning. Ela inicializa um conjunto de tokens de prompt suaves treináveis ($P_{CT}$) utilizando os embeddings de prompt do modelo derivados da concatenação de pares de demonstração ($\mathcal{C}$). Esses tokens são então otimizados via gradiente descendente para minimizar a perda nos pares de demonstração.

2.2 CT-KV

Esta variante aplica os princípios do Prefix Tuning, mas com uma estratégia de inicialização específica. Ela inicializa um prefixo de chave-valor (KV) ($\Theta_{CT} = \{K_j, V_j\}_{j=1}^L$) para cada camada do transformer. Esses prefixos são derivados das ativações por camada do modelo no contexto de demonstração $\mathcal{C}$. Os prefixos KV são então otimizados para minimizar a perda.

2.3 Escolhas Críticas de Design

Para garantir um desempenho forte e evitar que o modelo simplesmente recupere respostas do contexto, os autores introduzem duas escolhas de design críticas para ambas as variantes:

1. Mascaramento Leave-One-Out: Para evitar que o modelo recupere diretamente a saída alvo $y_i$ do $i$-ésimo par de demonstração do contexto ao prever $y_i$, a parte correspondente do vetor de contexto é mascarada durante a otimização.

   • Para CT-Prompt, os tokens de prompt suaves correspondentes ao par $i$ são mascarados.
   • Para CT-KV, os prefixos KV correspondentes ao par $i$ são mascarados.
   • Nota: O artigo observa que isso é crucial para a maioria dos benchmarks, mas foi considerado menos eficaz (ou prejudicial) para o benchmark ARC devido ao seu número extremamente baixo de pares de demonstração.

2. Token Dropout: Para regularizar o treinamento e evitar o overfitting a tokens específicos (especialmente dado o maior número de tokens introduzidos em comparação ao prompt tuning padrão), os tokens na representação do contexto são descartados aleatoriamente com uma probabilidade fixa durante a otimização. A perda é computada em expectativa sobre essas máscaras estocásticas.

2.4 Abordagem Híbrida: TTT + CT-KV

Os autores propõem um fluxo de trabalho complementar onde o Treinamento em Tempo de Teste (TTT) primeiro atualiza os pesos do modelo usando LoRA, e então o CT-KV é aplicado para refinar a representação de contexto do modelo. Isso combina adaptação baseada em peso e baseada em contexto.

3. Principais Contribuições

• Framework de Ajuste de Contexto: Introdução de um método que inicializa prompts treináveis ou prefixos KV diretamente de pares de demonstração de poucos disparos, aproveitando a capacidade inerente de ICL do modelo antes de refiná-la via gradientes.
• Eficiência do CT-KV: Demonstração de que a otimização de caches de KV (CT-KV) alcança tempo de treinamento de complexidade linear em relação ao número de pares de demonstração, contrastando com a complexidade quadrática do CT-Prompt e TTT.
• Complementaridade: Mostra que a otimização de contexto (CT-KV) e a otimização de peso (TTT) são complementares; aplicar CT-KV após o TTT gera ganhos de desempenho adicionais.
• Framework de Otimização em Contexto (ICO): Formalização de uma categoria mais ampla de métodos de adaptação em tempo de inferência que otimizam ou parâmetros do modelo ou representações de contexto derivadas de demonstrações.

4. Resultados Experimentais

Os autores avaliaram o Ajuste de Contexto em um conjunto diversificado de benchmarks: NLP-LR (26 tarefas do CrossFit e UnifiedQA), MMLU (57 assuntos), BIG-Bench Hard (BBH) (23 tarefas de raciocínio complexo) e o Abstraction and Reasoning Corpus (ARC). Os modelos testados incluíram variações do GPT-2, Llama 3 (1B, 3B, 8B) e modelos maiores como Mistral-12B e DeepSeek-32B.

Destaques de Desempenho:

• Precisão: O CT-KV supera significativamente o ICL padrão, o Prompt Tuning tradicional, o Prefix Tuning e as variantes de LoRA em todos os benchmarks. Por exemplo, no NLP-LR, o CT-KV alcança 44,2% de precisão, superando o ICL (35,6%) e o MetaICL (43,3%).
• Eficiência: O CT-KV alcança desempenho competitivo com o TTT, mas requer significativamente menos tempo de treinamento. Em todos os benchmarks, o CT-KV requer no máximo metade do tempo de treinamento por tarefa em comparação ao TTT.
• Escalabilidade: O CT-KV demonstra melhorias consistentes sobre o ICL e o Prefix Tuning em modelos variando de 1B a 32B de parâmetros.
• Robustez:
  • Contagem de Demonstrações: O CT-KV escala efetivamente conforme o número de pares de demonstração ($k$) aumenta, mostrando maiores ganhos do que o ICL ou Prefix Tuning.
  • Ruído: O CT-KV mantém um desempenho robusto mesmo com até 75% de ruído de rótulo nos exemplos de demonstração.
• Ablação: Remover o Mascaramento Leave-One-Out ou o Token Dropout leva a quedas significativas de desempenho na maioria dos benchmarks, confirmando sua necessidade para generalização e prevenção de aprendizado de atalhos (shortcut learning).

5. Significância e Alegações

O artigo afirma que o Ajuste de Contexto representa uma direção poderosa e escalável para o aprendizado de poucos disparos. Ao otimizar a representação de contexto leve em vez dos pesos do modelo, o método alcança:

1. Alta Eficiência: Oferece um compromisso favorável entre desempenho e tempo de treinamento, superando o ICL e métodos baseados em prompt, sendo mais eficiente que o TTT.
2. Inicialização Eficaz: Inicializar representações treináveis a partir de pares de demonstração (em vez de tokens aleatórios) reduz a variância e aproveita o conhecimento pré-treinado do modelo de forma mais eficaz.
3. Complementaridade: A capacidade de combinar o ajuste de contexto com a adaptação baseada em peso (TTT) sugere uma abordagem modular para a otimização em tempo de inferência.

Os autores posicionam seu trabalho como um passo em direção a estratégias de adaptação mais eficientes que não exigem recursos computacionais pesados ou retreinamento total do modelo, tornando-os adequados para cenários dinâmicos de poucos disparos. Eles reconhecem limitações relativas ao potencial de overfitting em tarefas específicas e sugerem trabalhos futuros em regularização mais forte e compressão de cache KV.