PRISM: Desmistificando a Retenção e a Interação no Meio do Treinamento

PRISM: Demystifying Retention and Interaction in Mid-Training

March 17, 2026

Autores: Bharat Runwal, Ashish Agrawal, Anurag Roy, Rameswar Panda

cs.AI

Resumo

Apresentamos o PRISM, um estudo empírico abrangente sobre as escolhas de design durante o treinamento intermediário de grandes modelos de linguagem. Por meio de experimentos controlados com sete modelos base, abrangendo quatro famílias (Granite, LLaMA, Mistral, Nemotron-H), dois tipos de arquitetura (Transformer denso e híbrido attention-Mamba) e escalas de 3B a 24B de parâmetros, demonstramos que o treinamento intermediário com aproximadamente 27B de *tokens* de alta qualidade produz ganhos consistentes de +15 a +40 pontos em matemática, +5 a +12 pontos em código e +6 a +13 pontos em benchmarks de ciências, preservando o desempenho geral. O *pipeline* completo PRISM para RL (Aprendizado por Reforço) melhora a macro-média em seis benchmarks de raciocínio de menos de 12 para 29-42 (uma melhoria de 3-4x), enquanto o RL aplicado diretamente à maioria dos modelos base permanece substancialmente menos eficaz, com pontuações AIME próximas de zero. A composição dos dados é mais importante no treinamento intermediário do que no RL: incluir dados científicos durante o treinamento intermediário desbloqueia ganhos de +17 a +28 pontos no GPQA-Diamond durante o RL, enquanto alterar a mistura de dados no RL produz diferenças inferiores a 2 pontos. Mecanicamente, o treinamento intermediário reestrutura densamente mais de 90% dos pesos do modelo, enquanto o RL faz refinamentos esparsos e frontais em aproximadamente 5% dos parâmetros. A análise de representações (CKA) confirma que o RL preserva consistentemente a geometria representacional do treinamento intermediário (CKA acima de 0,998) entre arquiteturas. Crucialmente, o RL aplica alterações de peso idênticas independentemente do ponto de partida, mas só tem sucesso em modelos com treinamento intermediário, consistente com a ideia de que este coloca o modelo em uma configuração a partir da qual o RL pode melhorar efetivamente o desempenho. Nossos resultados demonstram que o treinamento intermediário com retenção de conhecimento é altamente eficaz para o aprimoramento confiável do raciocínio e fornecem orientações práticas para a concepção de *pipelines* robustos de treinamento intermediário.

English

We present PRISM, a comprehensive empirical study of mid-training design choices for large language models. Through controlled experiments across seven base models spanning four families (Granite, LLaMA, Mistral, Nemotron-H), two architecture types (dense Transformer and attention-Mamba hybrid), and scales from 3B to 24B parameters, we show that mid-training on approximately 27B high-quality tokens yields consistent gains of +15 to +40 points on math, +5 to +12 points on code, and +6 to +13 points on science benchmarks while preserving general performance. The full PRISM to RL pipeline improves macro-average across six reasoning benchmarks from under 12 to 29-42 (a 3-4x improvement), whereas RL applied directly to most of the base models remains substantially less effective, with AIME scores near zero. Data composition matters most at mid-training, not RL: including science data during mid-training unlocks +17 to +28 point GPQA-Diamond gains during RL, while changing the RL mix produces less than 2 point differences. Mechanistically, mid-training densely restructures over 90% of model weights, while RL makes sparse, front-loaded refinements to approximately 5% of parameters. Representation analysis (CKA) confirms that RL consistently preserves mid-training's representational geometry (over 0.998 CKA) across architectures. Crucially, RL applies identical weight changes regardless of starting point, yet only succeeds on mid-trained models, consistent with mid-training placing the model in a configuration from which RL can effectively improve performance. Our results demonstrate that retention-aware mid-training is highly effective for reliable reasoning enhancement and provide practical guidance for designing robust mid-training pipelines.