Guiding LLM Post-training Data Engineering with Model Internals from Sparse Autoencoders

Large language models (LLMs) have revolutionized natural language processing by leveraging massive pretraining on diverse corpora. However, to achieve state-of-the-art performance on complex tasks, post-training fine-tuning, especially via reinforcement learning (RL), has become essential. RL enables models to optimize for specific objectives, such as alignment with human preferences or task-specific rewards. Effective data engineering during this phase—selecting which samples to train on, their ordering, and batching—is critical to maximize learning efficiency and final model quality.

Traditional post-training data engineering relies on external feedback signals including human annotations, verifier outcomes, or rollout success rates. While useful, these signals are costly to obtain, often sparse, and may not fully capture the intrinsic structure of the data as perceived by the model. Concurrently, advances in mechanistic interpretability have revealed that model internal activations encode rich semantic and structural information about input data. Sparse Autoencoders (SAEs) have emerged as a powerful tool to decompose dense LLM hidden states into sparse, disentangled features, offering a fine-grained lens into model internals.

Previous work has leveraged model internals for pretraining data selection or supervised fine-tuning, but their role in post-training RL data engineering remains underexplored. This gap motivates the current study to harness SAE activations as intrinsic signals for guiding data engineering operations, aiming to improve training efficiency and model performance.

Post-training data engineering faces three intertwined challenges: ensuring batch diversity to cover a wide range of semantic and reasoning patterns, constructing curricula that order samples from easy to hard to facilitate progressive learning, and filtering out noisy or off-distribution samples to maintain data quality. Existing methods predominantly depend on external signals such as human difficulty annotations, rollout accuracy, or preference labels, which are expensive and may not scale well.

Moreover, these external signals often provide coarse or delayed feedback, limiting their effectiveness in fine-grained data selection and ordering. The question arises whether intrinsic model signals—specifically, sparse and interpretable activations extracted via SAE—can reliably encode these data properties and be operationalized into concrete data engineering steps. Addressing this is crucial for developing more adaptive, efficient, and interpretable post-training pipelines for LLMs.

This work introduces several key innovations:

• �� Application of Sparse Autoencoders to extract sparse, disentangled feature activations from LLM hidden states, providing a structured and interpretable representation space for data characterization.

• �� Joint modeling of three intrinsic data properties—diversity, difficulty, and quality—directly from SAE activations, moving beyond reliance on external or scalar feedback.

• �� Mapping these intrinsic properties to concrete data engineering operations: SAE-space clustering and moderate batch mixing for diversity-driven batch construction; ElasticNet-based difficulty prediction and cluster-wise calibration for curriculum learning; and linear classification over SAE features for quality-based data filtering.

• �� Introduction of a moderate batch mixing strategy that balances within-batch gradient coherence and cross-cluster coverage, optimizing training dynamics.

• �� Demonstration that a single SAE encoder trained on a smaller model can transfer effectively across model scales and RL algorithms, highlighting engineering efficiency and generalizability.

• �� SAE Representation Extraction: For each training sample, extract hidden activations from the 27th layer of the LLM separately for prompt and solution spans. Encode these activations using a pretrained SAE, then aggregate via mean and max pooling to capture sustained and localized patterns, resulting in a 960-dimensional sparse feature vector concatenated with shallow metadata (e.g., length, TeX ratio).

• �� Diversity-Driven Batch Construction: Perform MiniBatchKMeans clustering in SAE feature space to partition samples into K clusters reflecting semantic and reasoning structures. Construct batches by grouping samples within clusters and apply moderate batch mixing by swapping a small tail portion of samples between adjacent batches from different clusters but similar difficulty and length, balancing gradient coherence and coverage.

• �� Difficulty-Driven Curriculum Ordering: Train an ElasticNet regressor on a small labeled subset (3,000 samples) to predict continuous difficulty scores from SAE features. Calibrate scores cluster-wise using shrinkage-based corrections to normalize scales. Sort samples within clusters by calibrated difficulty to form easy-to-hard trajectories. Globally interleave batches across clusters stage-wise.

• �� Quality-Driven Data Filtering: Train a linear classifier over SAE features to estimate the probability of a sample belonging to the target distribution, using labeled source data. Filter samples by thresholding or top-K ranking to exclude noisy or off-distribution data before curriculum construction.

• �� Training and Evaluation: Integrate SAERL with GRPO and DAPO RL algorithms, train on DeepMath-103K dataset with batch size 128, evaluate on six mathematical reasoning benchmarks (GSM8K, AMC23, MATH500, MinervaMath, OlympiadBench, AIME24) using Avg@8 and Pass@8 metrics.

Experiments are conducted on the DeepMath-103K dataset, a large-scale mathematical reasoning corpus with annotated topic labels and difficulty scores. Two model scales, Qwen2.5-Math-1.5B and 7B, are used to test scalability. Baselines include vanilla GRPO and DAPO without curriculum, Difficulty Curriculum Learning using external difficulty labels, rollout accuracy-based curriculum methods, and hidden state compression-based data selection.

Evaluation metrics include average accuracy (Avg@8) across five benchmarks and pass rate (Pass@8) on the competition-level AIME24 dataset. Ablation studies dissect the contributions of difficulty sorting, batch mixing, and cluster grouping. Training efficiency is measured by the number of steps to reach target accuracy. Additional experiments assess SAE activations’ ability to predict data diversity (topic labels), difficulty (Spearman correlation), and quality (Pearson correlation), as well as their effectiveness in filtering noisy samples from a mixed data pool.

SAERL achieves a 3.00% average accuracy improvement over vanilla GRPO on Qwen2.5-Math-1.5B and reduces training steps by 20% to reach target accuracy. Consistent improvements are observed across model scales and RL algorithms, with SAERL outperforming external difficulty label-based curricula and rollout accuracy methods. Ablations confirm that difficulty-based curriculum ordering is essential, while moderate batch mixing and cluster-first grouping further enhance performance. Batch diversity analysis reveals a concave relationship between mixing strength and performance, with moderate mixing yielding optimal results. SAE activations enable accurate prediction of data diversity (topic label linear probe accuracy significantly above majority baseline), difficulty (Spearman ρ up to 0.749), and quality (Pearson r improved to 0.3715). In noisy data selection, SAE-based linear classifiers achieve ROC-AUC of 0.9911, effectively filtering off-distribution samples.

SAERL is directly applicable to post-training RL of LLMs, particularly in structured reasoning domains like mathematics. Its lightweight SAE encoder facilitates rapid deployment and cross-model scalability, making it attractive for industrial training pipelines aiming to optimize data selection and curriculum without extensive external labeling. By improving training efficiency and final model accuracy, SAERL can reduce computational costs and accelerate model iteration. Future extensions could adapt the framework to code generation, agentic RL, multi-modal tasks, and general instruction following, broadening its impact across AI applications.

The study focuses on mathematical reasoning with verifiable rewards, limiting demonstrated applicability to other domains. Difficulty and quality proxies rely on limited labeled data and source distribution labels, hindering fully unsupervised operation. Theoretical understanding of SAE-space distances’ causal impact on training dynamics is incomplete, necessitating further gradient-level analyses. Additionally, the optimal batch mixing parameter may vary by task and model, requiring adaptive tuning.