Resumo Técnico: TaskPGM – Aprendendo Misturas de Tarefas por meio de Afinidades de Tarefas

Declaração do Problema

O Ajuste Fino Supervisionado (SFT) para Grandes Modelos de Linguagem (LLMs) depende fortemente da distribuição do orçamento de treinamento através de um conjunto heterogêneo de tarefas de instrução. As práticas atuais frequentemente dependem de heurísticas simples, como amostragem uniforme ou ponderação proporcional ao tamanho do conjunto de dados. Esses métodos falham em considerar as interações entre tarefas, levando a duas ineficiências primárias:

1. Redundância: Sobreponderação de tarefas que induzem atualizações comportamentais quase idênticas, desperdiçando o orçamento computacional.
2. Transferência Perdida: Subutilização de tarefas que oferecem informações complementares ou benefícios de transferência entre domínios.

Abordagens automatizadas existentes (por exemplo, recuperação baseada em influência como LESS, ou seleção submodular como SMART) frequentemente otimizam para subconjuntos discretos em vez de proporções de mistura contínuas, dependem de similaridade semântica que ignora o comportamento do modelo downstream, ou exigem treinamento iterativo dispendioso. O desafio central é determinar automaticamente uma mistura de tarefas contínua que seja simultaneamente representativa (alta utilidade) e diversa (baixa redundância) para maximizar o desempenho downstream sob um orçamento fixo.

Metodologia: TaskPGM

Os autores propõem o TaskPGM (Task Mixtures via Probabilistic Graphical Modelling), um framework que trata a mistura de dados de SFT como um problema de minimização de energia sobre um Campo de Markov Aleatório (MRF) denso.

1. Formulação do Modelo Gráfico Probabilístico

As tarefas são modeladas como nós em um MRF. O objetivo é encontrar uma distribuição de amostragem ideal $p^* \in \Delta_n$ (onde $\Delta_n$ é o simplex de probabilidade) que maximize uma função de energia $E(p)$. A função de energia consiste em dois componentes:

• Potenciais Unários ($\Psi_{un}$): Codificam a utilidade ou representatividade de tarefas individuais. Isso é derivado da massa total de similaridade de uma tarefa em relação à coleção inteira.
• Potenciais Pareados ($\Psi_{pair}$): Codificam redundância e diversidade. Estes termos penalizam a coocorrência de tarefas altamente similares na mistura.

O objetivo de otimização é formulado como:
$$\max_{p \in \Delta_n} E(p) = \Psi_{un}^\top p - \frac{1}{2} p^\top \Psi_{pair} p$$
Isso é equivalente a minimizar uma função quadrática convexa (sob condições específicas) sujeita a restrições de simplex.

2. Métricas de Afinidade Comportamental de Tarefas

Uma inovação fundamental é o método para computar afinidades de tarefas pareadas ($S_{ij}$). Em vez de depender de similaridade semântica (embeddings) ou metadados, o TaskPGM computa afinidades baseadas em divergências comportamentais derivadas das distribuições preditivas de modelos ajustados para tarefa única.

• Informação Mútua de Ponto (PMI): Mede a informação mútua entre as distribuições preditivas de duas tarefas quando avaliadas nos conjuntos de dados uma da outra.
• Divergência de Jensen–Shannon (JSD): Uma medida simétrica da discrepância entre as distribuições preditivas de dois modelos em exemplos de tarefas cruzadas.

Essas métricas capturam como as tarefas moldam o comportamento do modelo em exemplos fora da tarefa, abordando diretamente a transferência e a redundância entre tarefas.

3. Otimização e Solvers

• Mistura Contínua: Os autores derivam uma solução de forma fechada para os pesos da mistura contínua usando as condições de Karush-Kuhn-Tucker (KKT).
  • Se a matriz de similaridade pareada não for Semidefinida Positiva (PSD), eles aplicam deslocamento de espectro (adicionando uma constante à diagonal) para garantir a convexidade, o que introduz um termo de regularização que enviesa a solução em direção à uniformidade se a estrutura não-PSD for severa.
  • Quando a inversa é mal condicionada ou as restrições de não-negatividade são violadas, eles utilizam o Gradiente Projetado (PGD) no simplex de probabilidade.
• Variante Discreta/Orçada: Para cenários que exigem a seleção de um subconjunto de tarefas (top-$k$), os autores elevam a otimização contínua para uma função de conjunto discreto. Eles provam que esta função é monotônica e fracamente submodular sob condições de regularidade padrão (Convexidade Forte Restrita e Suavidade), permitindo a construção gulosa eficiente com garantias de aproximação.

Principais Contribuições

1. Framework Probabilístico para Design de Mistura: Uma formulação novel da mistura de dados de SFT como minimização de energia em um MRF, equilibrando explicitamente representatividade (termos unários) e diversidade (termos pareados).
2. Métricas de Afinidade Baseadas em Comportamento: A proposta de computar relações de tarefas usando divergências de distribuição preditiva (JSD/PMI) de modelos ajustados para tarefa única, indo além da similaridade semântica estática para capturar dependências funcionais de tarefas.
3. Otimização com Base Teórica:
   • Derivação de uma solução de forma fechada para misturas contínuas via condições KKT.
     ção da submodularidade fraca para a variante orçada discreta, suportando a seleção gulosa com garantias teóricas.
4. Validação Empírica: Demonstração de ganhos de desempenho consistentes em múltiplas famílias de modelos (Qwen2-7B, Llama-2-7B) e benchmarks (BIG-Bench Hard, MMLU, etc.) comparado a heurísticas e baselines de seleção de estado da arte.

Resultados Experimentais

O framework foi avaliado em 309 tarefas de instrução usando modelos Qwen2-7B e Llama-2-7B com orçamentos de 25K e 50K instâncias.

• Ganhos de Desempenho: O TaskPGM (usando tanto as variantes PMI quanto JSD) superou consistentemente a amostragem uniforme, o tamanho proporcional (EPM) e os baselines aleatórios.
  • No Qwen2-7B (orçamento de 25K), a variante PMI melhorou o desempenho em quase 3% sobre os baselines heurísticos em benchmarks de raciocínio complexo.
  • No Llama-2-7B, a mistura derivada de PMI superou o Uniform e o EPM em 2% a mais de 4% no BBH.
• Comparação com SOTA: O TaskPGM permaneceu competitivo ou superior ao LESS (um forte baseline baseado em função de influência). Notavelmente, o TaskPGM alcança isso com um único checkpoint de propósito geral, enquanto o LESS constrói modelos distintos e orientados ao alvo para cada benchmark.
• Eficiência: O TaskPGm é computacionalmente mais eficiente que o LESS. O cálculo de similaridade é um custo amortizado de etapa única (~16–34 horas de GPU), após o qual as misturas para qualquer orçamento podem ser derivadas em segundos via solução de forma fechada. Em contraste, o LESS requer ~320 horas de GPU para construir repositórios de gradiente para cada alvo.
• Escalabilidade: Aumentar o orçamento de 25K para 50K gerou ganhos consistentes para o TaskPGM, enquanto métodos heurísticos (Uniform, EPM) exibiram comportamento fragmentado e não robusto conforme os orçamentos cresciam.

Significância e Alegações

O artigo afirma que o TaskPGM fornece uma abordagem principista, interpretável e eficiente para a mistura de dados de SFT.

• Interpretabilidade: Os pesos aprendidos e os mapas de calor de similaridade revelam clusters de tarefas, redundância e tarefas "atratoras", oferecendo insights sobre por que certas fontes ajudam.
• Generalização: Diferente de métodos que otimizam para alvos específicos, o TaskPGM produz uma mistura única que melhora o desempenho em um amplo espectro de tarefas downstream (raciocínio, compreensão de linguagem, conhecimento).
• Robustez: Os ganhos são impulsionados pela seleção funcional de tarefas baseada no comportamento do modelo, em vez de variância favorável de amostragem, como evidenciado pelos melhorias consistentes através de diferentes modelos base e métricas de similaridade.

Os autores concluem que a mistura baseada em princípios e similaridade é essencial para o ajuste fino de instrução escalável, muitas vezes igualando ou excedendo misturas heurísticas com substancialmente menos dados. Eles reconhecem limitações em relação ao pool fixo de tarefas e ao custo computacional da estimativa de afinidade, sugerindo aproximações escaláveis como uma direção para trabalhos futuros.