Resumo Técnico: Alinhamento Distribucional como Critério para o Projeto de Vetores de Tarefa em Aprendizado em Contexto

Declaração do Problema

O Aprendizado em Contexto (ICL) permite que Modelos de Linguagem de Grande Escala (LLMs) se adaptem a novas tarefas ao prependir demonstrações rotuladas à consulta de entrada. Embora eficaz, o ICL incorre em custos de inferência crescentes e sobrecargas de memória à medida que o comprimento do contexto aumenta, dificultando sua utilidade prática sob restrições de recursos. Para abordar isso, pesquisas recentes propuseram Vetores de Tarefa (TVs): representações compactas de estados ocultos extraídas das ativações internas de um LLM durante o ICL. Esses vetores encapsulam informações de tarefa inferidas das demonstrações, permitindo que o modelo realize adaptação de tarefa durante a inferência zero-shot sem reprocessar as demonstrações.

No entanto, os métodos existentes para extração de vetores de tarefa carecem de um critério de projeto principiado. Atualmente, a qualidade de um método de extração de vetores de tarefa é avaliada exclusivamente por meio da precisão em tarefas downstream. Essa métrica indireta oferece insights limitados sobre por que um método supera outro ou como projetar mecanismos de extração mais eficazes. Há uma necessidade de um critério direto e fundamentado teoricamente para avaliar e otimizar a qualidade dos vetores de tarefa.

Metodologia

1. O Critério Proposto: Alinhamento Distribucional

Os autores postulam que o objetivo principal de um vetor de tarefa é condensar o efeito das demonstrações de modo que a distribuição preditiva resultante se alinhe estreitamente com a do ICL padrão. Para quantificar esse alinhamento, eles introduzem o dNTP (discrepância na Probabilidade do Próximo Token).

• Definição: O dNTP mede a divergência de Kullback-Leibler (KL) entre a distribuição preditiva do próximo token sob ICL ($P_{icl}$) e a distribuição sob inferência baseada em vetor de tarefa ($P_{tv}$).
  $$dNTP(f; Z) = \mathbb{E}_{x \sim D} \left[ D_{KL} \left( P_{icl}(\cdot | x, Z) \parallel P_{tv}(\cdot | x, f(Z)) \right) \right]$$
  onde $f$ é a função de extração e $Z$ representa as demonstrações.
• Validação: A análise empírica em múltiplos modelos e benchmarks revela uma forte correlação negativa entre dNTP e precisão downstream. Um dNTP mais baixo prevê consistentemente maior precisão, estabelecendo o dNTP como um proxy válido para a qualidade do vetor de tarefa que não requer rótulos do conjunto de teste.

2. Método Proposto: Vetor de Tarefa Linear (LTV)

Motivados pelo objetivo de minimizar o dNTP, os autores propõem o Vetor de Tarefa Linear (LTV). O método é derivado de um problema de otimização proxy.

• Racional: Minimizar o dNTP diretamente é intratável, pois carece de uma solução de forma fechada e requer atualizações iterativas. Os autores provam que minimizar o Erro Quadrático Médio (MSE) entre os estados ocultos do ICL ($h_{icl}$) e o estado oculto zero-shot mais o vetor de tarefa ($h_{zs} + v$) serve como um limite superior para o dNTP.
  $$LMSE(f) = \mathbb{E} \left[ \| h_{icl} - (h_{zs} + f(Z)) \|_2^2 \right]$$
• Mapeamento Linear: Para preservar a vantagem "sem treinamento" do ICL (evitando atualizações de parâmetros do modelo), o LTV restringe a função de extração $f$ a um mapeamento linear. O vetor de tarefa $v$ é estimado como uma transformação linear do estado oculto zero-shot $h_{zs}$:
  $$v = W h_{zs}(x)$$
• Solução de Forma Fechada: O mapeamento linear ótimo $W^*$ é computado via regressão ridge em um conjunto de consultas de treinamento não rotuladas. A regressão visa a diferença entre o estado oculto do ICL e o estado oculto zero-shot ($h_{icl} - h_{zs}$).
  $$W^* = \arg\min_W \left( \| Y - W H \|_F^2 + \lambda \| W \|_F^2 \right)$$
  onde $H$ contém estados ocultos zero-shot e $Y$ contém as diferenças alvo. Isso produz uma solução de forma fechada, garantindo eficiência.

Contribuições Principais

1. Métrica dNTP: Introdução de uma métrica que quantifica a discrepância entre as distribuições de inferência baseadas em TV e baseadas em ICL. O artigo demonstra que o dNTP é um forte proxy de desempenho, correlacionando-se negativamente com a precisão downstream.
2. Vetor de Tarefa Linear (LTV): Desenvolvimento de um método de extração de vetores de tarefa que minimiza o dNTP por meio de um mapeamento linear de forma fechada. Esta abordagem evita otimização iterativa e ajuste fino do modelo.
3. Superioridade Empírica: Extensos experimentos mostrando que o LTV supera as bases existentes sem treinamento (Function Vector, Task Vector, State Vector, I2CL) em oito benchmarks de classificação e cinco LLMs.
4. Transferência Cross-Modelo: Demonstração de que vetores de tarefa extraídos de um modelo maior (ex: 72B) podem ser transferidos para um modelo menor (ex: 7B) para melhorar o desempenho, abordando cenários onde o modelo de inferência tem capacidade ou comprimento de contexto limitados.

Resultados Experimentais

• Desempenho em Classificação: Em cinco LLMs (incluindo variantes de LLaMA-2/3 e Qwen) e oito benchmarks (ex: SST-2, AGNews, DBPedia), o LTV alcançou a maior precisão média. Especificamente, no LLaMA-2-13B, o LTV melhorou a precisão média em 9,2% sobre a melhor base (State Vector), atingindo 79,46%.
• Eficiência: O LTV mantém latência de inferência comparável à inferência zero-shot (aprox. 0,0226s por amostra), significativamente mais rápida que o ICL. Embora a extração exija um custo de pré-cálculo único, isso é amortizado sobre o conjunto de teste.
• Tarefas de Regressão: O LTV também foi avaliado em tarefas de regressão linear e ReLU, alcançando o menor Erro Quadrático Médio (MSE) entre todas as bases de vetores de tarefa.
• Robustez: O método permanece robusto a variações em hiperparâmetros (número de consultas de treinamento $N$ e regularização $\lambda$) e ao número de demonstrações ($k$).
• Transferibilidade: Transferir o LTV de um modelo Qwen-2.5-72B para um modelo Qwen-2.5-7B melhorou a precisão média de classificação do modelo menor em 6,43%, igualando em alguns casos o desempenho do próprio LTV do modelo maior.

Significado e Alegações

O artigo alega que o alinhamento distribucional é um critério fundamental e útil para projetar métodos de extração de vetores de tarefa. Ao deslocar o foco da avaliação da precisão downstream indireta para a discrepância distribucional direta (dNTP), os autores fornecem um framework principiado para entender e melhorar vetores de tarefa.

O significado do LTV reside em sua capacidade de:

1. Preencher a lacuna entre o objetivo teórico do ICL (inferência bayesiana implícita) e a implementação prática (vetores compactos) sem exigir atualizações do modelo.
2. Oferecer uma alternativa computacionalmente eficiente ao ICL, reduzindo a latência de inferência enquanto mantém alta precisão.
3. Permitir transferência cross-modelo, sugerindo que representações de tarefa extraídas de modelos maiores e mais capazes podem servir como recurso para modelos menores com contexto ou capacidade limitados.

Os autores concluem que o LTV oferece um framework mais eficiente e eficaz para compressão de tarefa em LLMs, validando a hipótese de que minimizar a discrepância entre as distribuições TV e ICL leva a uma adaptação de tarefa superior.