InfiAlign: Uma Estrutura Escalável e Eficiente em Amostras para Alinhar LLMs e Aprimorar Capacidades de Raciocínio

InfiAlign: A Scalable and Sample-Efficient Framework for Aligning LLMs to Enhance Reasoning Capabilities

August 7, 2025

Autores: Shuo Cai, Su Lu, Qi Zhou, Kejing Yang, Zhijie Sang, Congkai Xie, Hongxia Yang

cs.AI

Resumo

Modelos de linguagem de grande escala (LLMs) têm demonstrado habilidades impressionantes de raciocínio em uma ampla gama de tarefas complexas. No entanto, aprimorar essas capacidades por meio de pós-treinamento continua a ser um processo que demanda muitos recursos, especialmente em termos de dados e custo computacional. Embora esforços recentes tenham buscado melhorar a eficiência amostral por meio da curadoria seletiva de dados, os métodos existentes frequentemente dependem de estratégias heurísticas ou específicas para tarefas, o que dificulta a escalabilidade. Neste trabalho, apresentamos o InfiAlign, um framework de pós-treinamento escalável e eficiente em termos de amostras que integra o ajuste fino supervisionado (SFT) com a Otimização de Preferência Direta (DPO) para alinhar LLMs visando um raciocínio aprimorado. No núcleo do InfiAlign está um pipeline robusto de seleção de dados que automaticamente seleciona dados de alinhamento de alta qualidade a partir de conjuntos de dados de raciocínio de código aberto, utilizando métricas de qualidade multidimensionais. Esse pipeline permite ganhos significativos de desempenho enquanto reduz drasticamente os requisitos de dados e permanece extensível a novas fontes de dados. Quando aplicado ao modelo Qwen2.5-Math-7B-Base, nosso modelo SFT alcança desempenho equivalente ao DeepSeek-R1-Distill-Qwen-7B, utilizando apenas aproximadamente 12% dos dados de treinamento, e demonstra forte generalização em diversas tarefas de raciocínio. Melhorias adicionais são obtidas por meio da aplicação do DPO, com ganhos particularmente notáveis em tarefas de raciocínio matemático. O modelo alcança uma melhoria média de 3,89% nos benchmarks AIME 24/25. Nossos resultados destacam a eficácia de combinar a seleção de dados fundamentada com o pós-treinamento em todas as etapas, oferecendo uma solução prática para alinhar modelos de raciocínio de grande escala de maneira escalável e eficiente em termos de dados. Os checkpoints do modelo estão disponíveis em https://huggingface.co/InfiX-ai/InfiAlign-Qwen-7B-SFT.

English

Large language models (LLMs) have exhibited impressive reasoning abilities on a wide range of complex tasks. However, enhancing these capabilities through post-training remains resource intensive, particularly in terms of data and computational cost. Although recent efforts have sought to improve sample efficiency through selective data curation, existing methods often rely on heuristic or task-specific strategies that hinder scalability. In this work, we introduce InfiAlign, a scalable and sample-efficient post-training framework that integrates supervised fine-tuning (SFT) with Direct Preference Optimization (DPO) to align LLMs for enhanced reasoning. At the core of InfiAlign is a robust data selection pipeline that automatically curates high-quality alignment data from open-source reasoning datasets using multidimensional quality metrics. This pipeline enables significant performance gains while drastically reducing data requirements and remains extensible to new data sources. When applied to the Qwen2.5-Math-7B-Base model, our SFT model achieves performance on par with DeepSeek-R1-Distill-Qwen-7B, while using only approximately 12% of the training data, and demonstrates strong generalization across diverse reasoning tasks. Additional improvements are obtained through the application of DPO, with particularly notable gains in mathematical reasoning tasks. The model achieves an average improvement of 3.89% on AIME 24/25 benchmarks. Our results highlight the effectiveness of combining principled data selection with full-stage post-training, offering a practical solution for aligning large reasoning models in a scalable and data-efficient manner. The model checkpoints are available at https://huggingface.co/InfiX-ai/InfiAlign-Qwen-7B-SFT.