InfiAlign: Un marco escalable y eficiente en muestras para alinear modelos de lenguaje grandes y mejorar capacidades de razonamiento

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

Resumen

Los modelos de lenguaje de gran escala (LLMs) han demostrado capacidades de razonamiento impresionantes en una amplia gama de tareas complejas. Sin embargo, mejorar estas capacidades mediante el entrenamiento posterior sigue siendo intensivo en recursos, particularmente en términos de datos y costos computacionales. Aunque esfuerzos recientes han buscado mejorar la eficiencia de muestreo mediante la curación selectiva de datos, los métodos existentes a menudo dependen de estrategias heurísticas o específicas de tareas que dificultan la escalabilidad. En este trabajo, presentamos InfiAlign, un marco de entrenamiento posterior escalable y eficiente en muestreo que integra el ajuste fino supervisado (SFT) con la Optimización Directa de Preferencias (DPO) para alinear LLMs y mejorar su razonamiento. En el núcleo de InfiAlign se encuentra una robusta canalización de selección de datos que cura automáticamente datos de alineación de alta calidad a partir de conjuntos de datos de razonamiento de código abierto utilizando métricas de calidad multidimensionales. Esta canalización permite ganancias significativas en el rendimiento mientras reduce drásticamente los requisitos de datos y sigue siendo extensible a nuevas fuentes de datos. Cuando se aplica al modelo Qwen2.5-Math-7B-Base, nuestro modelo SFT logra un rendimiento comparable con DeepSeek-R1-Distill-Qwen-7B, utilizando solo aproximadamente el 12% de los datos de entrenamiento, y demuestra una fuerte generalización en diversas tareas de razonamiento. Se obtienen mejoras adicionales mediante la aplicación de DPO, con avances particularmente notables en tareas de razonamiento matemático. El modelo logra una mejora promedio del 3.89% en los benchmarks AIME 24/25. Nuestros resultados resaltan la efectividad de combinar la selección de datos basada en principios con el entrenamiento posterior en todas las etapas, ofreciendo una solución práctica para alinear modelos de razonamiento de gran escala de manera escalable y eficiente en datos. Los puntos de control del modelo están disponibles en 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.