Resumo Técnico: F-TIS – Aproveitando Modelos Diversos em GRPO Colaborativo

Declaração do Problema

O Otimização de Política Relativa em Grupo (GRPO) emergiu como um método dominante para o pós-treinamento de Modelos de Linguagem de Grande Escala (LLMs), particularmente para aprimorar capacidades de raciocínio. Diferentemente da Otimização de Política Proximal (PPO), o GRPO elimina a necessidade de um modelo de valor separado, calculando vantagens relativas a um grupo de conclusões geradas para um único prompt. No entanto, a natureza auto-regressiva da geração em LLMs cria um gargalo computacional significativo durante a fase de inferência.

Para abordar isso, trabalhos anteriores exploraram a distribuição do passo de geração entre múltiplos nós. Abordagens distribuídas existentes assumem majoritariamente modelos homogêneos (arquiteturas e parâmetros idênticos) para manter condições "on-policy". Essa suposição é impraticável para cenários de treinamento descentralizado, onde participantes possuem recursos e preferências heterogêneos. Em tais ambientes, surgem três tipos de heterogeneidade:

1. Tamanho do Modelo: Nós podem executar modelos com contagens de parâmetros drasticamente diferentes (ex.: 1B vs. 3B).
2. Especialização do Modelo: Nós podem possuir modelos de tamanho similar, mas com especializações distintas (ex.: generalista vs. programador).
3. Parâmetros Treináveis: Nós podem empregar diferentes estratégias de ajuste fino (ex.: ajuste fino completo vs. Ajuste Fino Eficiente em Parâmetros como LoRA).

Quando modelos heterogêneos colaboram sem ajuste, os dados de treinamento tornam-se altamente off-policy. A divergência entre a política geradora ($\pi_{gen}$) e a política de atualização ($\pi_{\theta}$) introduz ruído destrutivo, levando ao colapso da política ou a desempenho significativamente degradado em comparação ao treinamento em isolamento.

Metodologia: Amostragem por Importância Filtrada e Truncada (F-TIS)

Os autores propõem o F-TIS, um framework unificado projetado para permitir treinamento colaborativo GRPO estável entre modelos heterogêneos, minimizando ao mesmo tempo a sobrecarga de comunicação. A metodologia combina dois mecanismos-chave: Amostragem por Importância Truncada (TIS) e Filtragem de Amostras.

1. Amostragem por Importância Truncada (TIS)

O GRPO padrão assume dados on-policy. Para lidar com amostras off-policy, os autores adaptam a função de perda usando amostragem por importância. Eles comparam três abordagens:

• NoIS (Sem Amostragem por Importância): Trata todas as amostras como on-policy (linha de base para falha).
• VIS (Amostragem por Importância Pura): Aplica a razão de importância $\frac{\pi_{\theta}}{\pi_{gen}}$ a cada token.
• TIS (Amostragem por Importância Truncada): Trunca a razão de importância em uma constante $C$ (ex.: $C=2$) para impedir que pesos extremos desestabilizem o treinamento. A razão é movida para fora da soma interna sobre os tokens.

A formulação da perda TIS é:
$$\mathcal{L}_{GRPO} = \frac{1}{G} \sum_{i=1}^{G} \frac{1}{|a_i|} \sum_{t=1}^{|a_i|} \min \left( \frac{\pi_{\theta}(a_{i,t}|p \circ a_{i,<t})}{\pi_{gen}(a_{i,t}|p \circ a_{i,<t})}, C \right) \cdot \min \left( R_{i,\theta} \hat{A}_i, \text{clip}(R_{i,\theta} \hat{A}_i, 1-\epsilon, 1+\epsilon) \right)$$
Onde $R_{i,\theta}$ é a razão desconectada e $\hat{A}_i$ é a vantagem do grupo.

2. Filtragem de Amostras

Os autores observam que amostras off-policy com vantagens negativas ($\hat{A}_i < 0$) podem amplificar tokens improváveis, levando a conclusões sem sentido. Eles introduzem um mecanismo de filtragem baseado em:

• Vantagem: Amostras com $\hat{A}_i < 0$ são descartadas para o passo de atualização.
• Divergência KL: Amostras onde a divergência KL entre a política atual e a política geradora excede um limiar $g$ ($D_{KL}(\pi_{\theta} \parallel \pi_{gen}) > g$) também são filtradas.

Amostras filtradas ainda são usadas para calcular a vantagem do grupo (para manter o ranking relativo do grupo), mas são excluídas da atualização de gradiente.

3. O Framework F-TIS

O F-TIS integra TIS e filtragem. O termo de vantagem $\hat{A}_{t,i}$ na função de perda é definido como 0 se a amostra for filtrada (ou seja, se $\hat{A}_i \leq 0$ OU $D_{KL} > g$).

• Eficiência de Comunicação: O framework requer comunicação mínima, transmitindo apenas tokens e log-probabilidades (aproximadamente 8 bytes por token), tornando-o adequado para redes descentralizadas.

Contribuições Principais

1. Identificação de Desafios de Heterogeneidade: O artigo demonstra empiricamente que o GRPO padrão falha em configurações descentralizadas heterogêneas (NoIS), levando à degradação de desempenho em comparação ao treinamento isolado.
2. Algoritmo F-TIS: Um paradigma de treinamento inovador que aproveita com sucesso amostras off-policy de modelos diversos (diferentes tamanhos, especializações e configurações PEFT) sem colapsar a política.
3. Eficiência de Comunicação: A abordagem mantém baixa sobrecarga de comunicação, exigindo apenas log-probabilidades e tokens, ao contrário de métodos que requerem sincronização completa do modelo ou trocas pesadas de modelos de valor.
4. Estudos de Ablação: Os autores analisam sistematicamente o impacto do limiar de filtragem $g$ e comparam F-TIS com F-VIS (Amostragem por Importância Pura Filtrada), demonstrando a superioridade da truncagem em estágios posteriores do treinamento.

Resultados Experimentais

Os autores avaliaram o F-TIS no conjunto de dados GSM8k para treinamento e no conjunto de dados MATH-500 para teste fora da distribuição (OOD), usando modelos Qwen2.5 (1.5B e 3B, variantes Base e Coder).

• Convergência: O F-TIS alcança convergência final do modelo idêntica ao treinamento puramente on-policy (homogêneo) em todas as configurações heterogêneas.
• Generalização (OOD): Em várias configurações, o F-TIS superou o treinamento on-policy isolado em tarefas OOD. Notavelmente:
  • O modelo 3B Coder emparelhado com um modelo Base melhorou seu desempenho no MATH-500 em 5,2%.
  • O modelo 3B Coder emparelhado com um modelo Coder menor melhorou em 12%.
  • O modelo 3B PEFT melhorou em 7% quando treinado colaborativamente.
• Limiar de Filtragem ($g$): Estudos de ablação revelaram que valores menores de $g$ beneficiam inicialmente modelos menores (1.5B) ao remover ruído aleatório, enquanto modelos maiores (3B) beneficiam-se de $g$ mais alto (ex.: 50) em estágios posteriores do treinamento, permitindo maior exploração do espaço.
• Vertical vs. Horizontal: O artigo conclui que a colaboração vertical (um nó gera o grupo completo para um prompt) é superior à colaboração horizontal (nós geram subconjuntos). A colaboração horizontal introduz viés no cálculo de vantagem em relação à média do enxame, degradando o desempenho.

Significado e Alegações

O artigo posiciona o F-TIS como um passo fundamental em direção ao RL colaborativo descentralizado para LLMs. Os autores afirmam que seu trabalho está entre os primeiros a estudar e permitir com sucesso o RL colaborativo entre modelos heterogêneos.

As principais alegações incluem:

• Viabilidade: É possível treinar modelos diversos juntos em uma única execução GRPO sem o colapso de desempenho tipicamente associado a dados off-policy.
• Eficiência: O método é eficiente em comunicação, tornando-o viável para sistemas descentralizados com restrições de rede variáveis.
• Raciocínio Aprimorado: Além de igualar o desempenho on-policy, o uso colaborativo de amostras off-policy pode atuar como um regularizador ou impulsionador de exploração, levando a melhor generalização em tarefas fora da distribuição (até 12% de melhoria) em comparação ao treinamento de modelos em isolamento.
• Praticidade: O framework acomoda restrições do mundo real onde usuários possuem diferentes capacidades de hardware (tamanhos de modelo) e necessidades especializadas (modelos especialistas ou PEFT), permitindo que contribuam para um objetivo de aprendizado compartilhado sem exigir infraestrutura idêntica.