Resumo Técnico: PartitionSel via Correspondência de Gradiente com Restrição de Matroide de Partição

Declaração do Problema

O treinamento de Grandes Modelos de Linguagem (LLMs) em dados heterogêneos requer o equilíbrio entre velocidade de convergência e cobertura em diversos domínios. Embora mini-lotes (mini-batches) grandes melhorem a convergência, eles incorrem em um overhead significativo de memória de GPU. Soluções existentes para treinamento eficiente em termos de memória geralmente dependem da seleção de subconjuntos de amostras dentro de mini-lotes. No entanto, quando os dados abrangem múltiplos domínios heterogêneos, as abordagens atuais enfrentam duas limitações primárias:

1. Seleção Independente: Métodos que selecionam amostras independentemente em cada domínio falham em considerar a redundância entre domínios e potenciais conflitos de gradiente.
2. Reponderação Baseada em Proxy: Métodos de mistura contínua de domínios (ex: DoReMi) dependem do treinamento de modelos auxiliares de proxy para aprender pesos ideais de domínio. Isso introduz um overhead computacional substancial e exige garantir que o proxy reflita fielmente a trajetória de otimização do LLM principal.

O desafio central é realizar a seleção de mini-lote entre domínios (cross-domain minibatch selection) que maximize a utilidade de treinamento enquanto respeita restrições de capacidade por domínio, sem o custo computacional de modelos de proxy auxiliares.

Metodologia: PartitionSel

Os autores propõem o PartitionSel, uma abordagem discreta e combinatória para a seleção de mini-lotes consciente de domínios. O método formula o problema como a maximização de uma utilidade única, guiada por validação, de correspondência de gradiente sujeita a restrições de matroide de partição que codificam orçamentos por domínio.

1. Função Objetivo

A seleção visa maximizar uma função de utilidade $f(S)$ para um subconjunto $S$ do mini-lote atual $B_t$. A utilidade é definida com base na redução esperada na perda de validação após um passo de descida de gradiente.

• Correspondência de Gradiente (Gradient Matching): A utilidade equilibra o alinhamento com sinais de validação contra a diversidade intra-subconjunto para limitar a redundância.
• Formulação de Protótipo Ponderado: Diferente de abordagens gananciosas anteriores (ex: GREATS) que utilizam seleção binária, o PartitionSel otimiza um vetor de pesos contínuos $w$ para o subconjunto selecionado. A função objetivo é:
  $$f(S) = \max_{w \ge 0, \text{supp}(w) \subseteq S} \left( \langle w, \bar{\mu}_{\theta_t} \rangle - \frac{\eta_t}{2} \langle w, K_{\theta_t} w \rangle \right)$$
  onde $\bar{\mu}_{\theta_t}$ representa o alinhamento agregado de gradiente com o conjunto de validação, e $K_{\theta_t}$ é a matriz de Gram dos gradientes que captura a similaridade (redundância).
• Monotonicidade e Submodularidade: Os autores provam que esta função objetivo ponderada é monotônica não-decrescente e fracamente submodular. Isso contrasta com o objetivo binário do GREATS, que foi observado como não-monotônico e não estritamente submodular.

2. Restrições (Matroide de Partição)

Para lidar com a heterogeneidade de domínio, a seleção é restringida por um matroide de partição. O conjunto base (o mini-lote) é particionado em domínios disjuntos $V_1, \dots, V_C$. Um subconjunto $S$ é válido se respeitar os limites de cardinalidade por domínio $\kappa_c$:
$$|S \cap V_c| \le \kappa_c \quad \forall c$$
Esta formulação induz uma mistura de domínios implícita ao nível do lote sem exigir otimização de pesos contínuos ou modelos de proxy.

3. Algoritmo de Otimização

Como a função objetivo é fracamente submodular, os autores empregam um algoritmo de Busca por Correspondência Ortogonal (OMP - Orthogonal Matching Pursuit) (Algoritmo 2) para encontrar uma solução aproximada com garantias prováveis.

• Seleção Iterativa: Começando de um conjunto vazio, o algoritmo adiciona iterativamente uma amostra por vez.
• Contração de Matroide: Em cada etapa, ele identifica um conjunto independente maximal no matroide contraído que maximiza a massa de gradiente positiva agregada, então amostra um elemento uniformemente deste conjunto.
• Reajuste de Peso (Weight Refitting): Após adicionar uma amostra, os pesos contínuos $w$ para o suporte atual são reotimizados usando o Gradiente Ascendente Projetado Acelerado (APGA).
• Garantia Teórica: O algoritmo alcança uma garantia de aproximação de $(1 + \tilde{C}_1/c_{2\kappa})^{-2}$ em relação à solução ótima, onde a razão depende das constantes de Suavidade Restrita (RSM) e Concavidade Forte Restrita (RSC) da função objetivo.

Principais Contribuições

1. Formulação Unificada e Acoplada de Domínio: O PartitionSel trata a seleção entre domínios como um único problema de maximização de utilidade sob restrições de matroide de partição. Ele recupera a regra GREATS como um caso especial binário, mas a melhora por ser provadamente monotônico e fracamente submodular.
2. Ponte Discreta para Métodos Contínuos: A abordagem fornece uma conexão principada entre a seleção de subconjuntos discretos e estratégias de reponderação contínua (como DoReMi). Ela induz uma mistura de domínios discreta, ao nível de amostra, em cada etapa de treinamento sem o overhead de treinar modelos de proxy.
3. Algoritmo Eficiente com Garantias: O algoritmo proposto baseado em OMP oferece garantias de aproximação prováveis para a maximização da função objetivo fracamente submodular sob restrições de matroide.

Resultados Experimentais

Os autores avaliaram o PartitionSel no ajuste fino (fine-tuning) dos modelos Qwen2.5 e Llama-3 através de dois benchmarks: MetaMathQA (raciocínio matemático) e Mol-Instructions (geração de moléculas).

• Ganhos de Desempenho: O PartitionSel superou consistentemente as bases fortes, incluindo:
  • GREATS: Uma base de correspondência de gradiente.
  • ID (Domínio Independente): Selecionando amostras independentemente por domínio.
  • IWD: Uma seleção de domínio independente relaxada com pesos contínuos.
  • COLM: Um método de seleção de coreset eficiente em memória.
  • DoReMi: Um método de mistura de dados contínua que requer modelos de proxy.
• Redução de Conflito de Gradiente: A análise mostrou que o PartitionSel reduz significamente o número de pares de gradientes conflitantes (onde $\langle g_i, g_j \rangle < 0$) dentro de um lote em comparação com a seleção independente (ID). Isso indica que o acoplamento entre domínios leva a atualizações de treinamento mais compatíveis.
• Perda de Validação: No MetaMathQA, o PartitionSel alcançou uma perplexidade logarítmica de validação menor em comparação tanto com o DoReMi quanto com outras bases.

Significância e Alegações

O artigo afirma que o PartitionSel oferece uma alternativa principada e computacionalmente eficiente às estratégias existentes de treinamento consciente de domínio.

• Eficiência: Ao evitar o treinamento de modelos de proxy auxiliares, ele reduz o overhead computacional enquanto mantém os benefícios da otimização consciente de domínio.
• Robustez: O método demonstra ganhos robustos em diferentes arquiteturas de modelo (Qwen2.5, Llama-3) e tarefas (matemática, química).
• Solidez Teórica: A formulação transforma um problema de seleção anteriormente não-monotônico em um problema monotônico e fracamente submodular, permitindo o uso de algoritmos com garantias de aproximação teóricas.
• Mistura Implícita: Ele induz com sucesso uma mistura de dados eficaz ao nível do lote através de seleção discreta, preenchendo a lacuna entre a seleção de subconjuntos discretos e métodos de reponderação contínua.

Os autores concluem que o PartitionSel é uma estratégia eficaz para treinar LLMs em diversos domínios, reduzindo conflitos de gradiente e melhorando o desempenho downstream sem a complexidade da otimização baseada em proxy.