Resumo Técnico: SPOKES – Otimizando para a Seleção de Dados de Pré-treinamento Diversos

1. Declaração do Problema

A diversidade de dados é um fator crítico na construção de corpora de pré-treinamento, pois reduz a redundância e melhora o desempenho em tarefas downstream sob orçamentos de dados fixos. No entanto, otimizar a diversidade é inerentemente difícil porque é uma propriedade de nível de conjunto que depende de interações entre pontos de dados, em vez de exemplos individuais. A otimização direta da diversidade é computacionalmente intratável devido à natureza combinatória da seleção de subconjuntos, particularmente na escala de pré-treinamento.

Abordagens existentes normalmente dependem de proxies ou aproximações, como agrupamento de grão grosso (por tópicos, habilidades ou clusters não supervisionados) ou deduplicação semântica. Embora esses métodos garantam a representação dos clusters, eles dependem de partições de baixa resolução e frequentemente falham em capturar informações em conhecimentos de cauda longa ou sub-representados. Além disso, os métodos atuais frequentemente tratam métricas de diversidade (como o escore G-Vendi) como ferramentas de avaliação post-hoc em vez de objetivos de otimização. O desafio central abordado por este trabalho é como extrair subconjuntos diversos de corpora de pré-treinamento de larga escala para melhorar resultados downstream sem depender de heurísticas que falham em escalar ou capturar diversidade fina.

2. Metodologia: SPOKES

Os autores introduzem o SPOKES (Spokes: Optimizing for Diverse Pretraining Data Selection), um framework escalável que otimiza diretamente a diversidade usando o escore G-Vendi dentro de um framework de diversificação probabilística.

2.1 Objetivo Central

O método formula a seleção de dados como um problema de otimização global que maximiza conjuntamente a qualidade e a diversidade. Dado um conjunto de dados de $n$ exemplos, o objetivo é selecionar um subconjunto $S$ de tamanho $k$ que maximize:
$$\max_{S \subseteq [n], |S|=k} \alpha \ln \left( \frac{1}{k} \sum_{i \in S} q_i \right) + (1 - \alpha) \ln \text{Vendi}(K_S)$$
Onde:

• $q_i$ representa o escore de qualidade do exemplo $i$ (ex: de um classificador FineWeb-Edu).
• $K_S$ é o kernel de similaridade restrito ao subconjunto selecionado.
• $\alpha \in [0, 1]$ é um parâmetro ajustável que controla o trade-off entre qualidade e diversidade.

2.2 Diversidade Baseada em Gradiente (G-Vendi)

Ao contrário dos métodos tradicionais que usam embeddings estáticos, o SPOKES opera no espaço de gradiente. Ele computa gradientes $\nabla_\theta \ell(x_i)$ em relação aos parâmetros do modelo $\theta$ para cada ponto de dado usando um modelo proxy. O escore G-Vendi mede a diversidade examinando a entropia do espectro da matriz de similaridade construída a partir desses gradientes. Esta abordagem incentiva a ortogonalidade entre vetores de representação, reduzindo a redundância nas direções de atualização e promovendo passos de otimização mais independentes e informativos.

2.3 Otimização via Gradiente Descendente Exponenciado

Para resolver o problema combinatório de seleção de subconjuntos, o SPOKES emprega uma estratégia de otimização relaxada:

1. Relaxação Contínua: Em vez de seleção discreta, um vetor de pesos não negativo contínuo $w$ é atribuído a cada ponto de dado, restrito ao simplex ($w \in \Delta_n$).
2. Construção do Kernel: Um kernel dependente de peso $K(w)$ é definido onde as entradas são escalonadas por $\sqrt{w_i w_j}$.
3. Otimização: O objetivo é otimizado usando gradiente descendente exponencial. Em cada iteração, os pesos são atualizados via $w_i^{(t+1)} \propto w_i^{(t)} \exp(-\eta g_i^{(t)})$ e normalizados para evitar o colapso sobre poucos exemplos de alto peso.
4. Recuperação Discreta: Após a convergência, o subconjunto final é recuperado selecionando os top-$k$ entradas do vetor de pesos aprendido.

2.4 Melhorias de Escalabilidade

Para tornar esta abordagem viável na escala de pré-treinamento, os autores implementam várias estratégias de eficiência:

• Gradientes Truncados: Em vez de computar gradientes de modelos completos, o método os aproxima usando apenas as últimas $n$ camadas do transformer (tipicamente 2 camadas). Isso reduz significativamente o custo computacional enquanto preserva a estrutura do kernel de similaridade.
• Projeção Johnson-Lindenstrauss (JL): Para lidar com vetores de gradiente de alta dimensão (ex: 200M de parâmetros), os autores aplicam projeções aleatórias (matrizes de Rademacher) para reduzir a dimensionalidade. Resultados empíricos mostram que mesmo projeções de baixa dimensão (ex: $k=1024$) mantêm alta estabilidade de subconjunto (95,5% de sobreposição com referências de alta dimensão).
• Paralelização: O conjunto de dados pode ser particionado em subconjuntos disjuntos otimizados independentemente, gerando aceleração linear.

3. Principais Contribuições

1. Framework SPOKES: Introdução de uma técnica escalável e fundamentada para otimização de diversidade que opera eficientemente na escala de conjuntos de dados de pré-treinamento modernos, indo além do agrupamento heurístico.
2. Otimização Direta de G-Vendi: Demonstração de que o escore G-Vendi pode ser diretamente otimizado via gradiente descendente exponencial para extrair subconjuntos altamente diversos, em vez de servir meramente como uma métrica post-hoc.
3. Otimização Conjunta de Qualidade-Diversidade: Prova de que a otimização conjunta de qualidade e diversidade produz resultados superiores em comparação com a otimização de apenas um dos objetivos ou o uso de baselines padrão.
4. Melhorias de Eficiência: Validação de que a computação de gradiente truncado e a redução de dimensionalidade via projeção JL permitem que o método escale para corpora massivos sem custos computacionais proibitivos.

4. Resultados Experimentais

O método foi avaliado em dois corpora de pré-treinamento de larga escala: DCLM (subconjunto Dolmino, agressivamente filtrado) e FineWeb (menos agressivamente filtrado). Modelos foram pré-treinados do zero em uma arquitetura LLaMA-1B usando 100 bilhões de tokens.

4.1 Métricas de Diversidade

• Aumento do Escore G-Vendi: Em um subconjunto de 500k amostras, o SPOKES alcançou um aumento de +489 pontos no escore G-Vendi em comparação com a amostragem aleatória. Em contraste, a deduplicação semântica (SemDeDup) alcançou apenas um aumento de +7 pontos.
• Diversidade de Batch: Durante o treinamento, os batches selecionados pelo SPOKES mantiveram scores de diversidade média consistentemente mais altos (aumentando de 315 para 520 no DCLM) em comparação com a amostragem aleatória.

4.2 Desempenho Downstream

A avaliação foi conduzida em 10 benchmarks de inglês do conjunto OLMES.

• Diversidade-Apenas: O SPOKES (diversidade-apenas) melhorou o desempenho médio downstream em +0,4 (DCLM) e +0,5 (FineWeb) pontos sobre a amostragem aleatória.
• Otimização Conjunta: Os resultados mais fortes foram alcançados pela otimização conjunta de qualidade e diversidade. O SPOKES alcançou ganhos de +1,5 (DCLM) e +1,4 (FineWeb) pontos sobre a amostragem aleatória, superando todos os baselines, incluindo deduplicação semântica e filtragem de apenas qualidade.
• Forças Complementares: Os resultados indicam que qualidade e diversidade capturam propriedades distintas, mas complementares. Enquanto a filtragem de qualidade se destaca em benchmarks intensivos em conhecimento (ex: MMLU), a otimização de diversidade melhora o raciocínio de senso comum (ex: HellaSwag). A otimização conjunta une essas lacunas, performando competitivamente em todas as tarefas.

5. Significância e Alegações

O artigo afirma que a otimização direta da diversidade é tanto tratável em escala quanto praticamente benéfica. Ao se afastar de aproximações de grão grosso e heurísticas, o SPOKES demonstra que:

• A diversidade é um eixo crítico: Mesmo sem supervisão explícita de qualidade, otimizar para diversidade isoladamente gera ganhos consistentes, sugerindo que uma representação de gradiente diversa da distribuição de dados melhora a generalização.
• Métricas de qualidade são insuficientes: Métricas de qualidade padrão não capturam totalmente a utilidade dos dados. A otimização conjunta da diversidade pode fornecer sinais de treinamento complementares que levam a um melhor desempenho downstream, mesmo quando o subconjunto selecionado possui um escore de qualidade média menor do que um baseline de apenas qualidade.
• Escalabilidade: O método proposto extrai com sucesso subconjuntos densos e altamente diversos de corpora de pré-treinamento massivos, oferezendo uma alternativa viável aos pipelines de seleção existentes que dependem de deduplicação semântica ou amostragem aleatória.

Os autores posicionam o SPOKES como uma ferramenta para construir subconjuntos de alta qualidade e diversos, enfatizando que os benefícios da diversidade se traduzem de forma confiável através de datasets com diferentes regimes de filtragem e perfis de ruído.