🤖 machine learning

Hierarchical Data Selection via Manifold Coverage and Sparse Feature Coverage in LLM Post-training

The paper proposes MASS, a hierarchical data selection framework that improves LLM post-training efficiency and performance by first grouping data via low-dimensional manifold coverage and then selecting high-quality subsets through sparse feature coverage, consistently outperforming existing baselines and often matching full-data training with significantly smaller datasets.

Original authors: Peng Sun, Yi Yang, Antong Zhang, Chunxiao Li, Yanbo Wang, Dianbo Liu, xin chen, Kai Yu, Lu Chen, Tianfan Fu

Published 2026-08-19

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Peng Sun, Yi Yang, Antong Zhang, Chunxiao Li, Yanbo Wang, Dianbo Liu, xin chen, Kai Yu, Lu Chen, Tianfan Fu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

In the rapidly evolving world of artificial intelligence, a persistent challenge has emerged as models grow larger and more capable: the sheer volume of data required to teach them. To learn effectively, these systems need vast libraries of examples, but gathering and processing this information is expensive and time-consuming. Researchers have long sought a way to identify the most valuable pieces of information within these massive collections, hoping to train powerful models using only a small, high-quality fraction of the available data. The prevailing idea was that if you could measure how different two pieces of data are from one another, you could pick a diverse set that covers all the necessary ground. However, this approach often stumbled because the mathematical spaces used to measure these differences were too crowded and noisy, mixing broad topics with tiny, irrelevant details and random errors.

A team of researchers has now proposed a new way to solve this problem, treating the selection of training data not as a flat search, but as a layered process of discovery. They argue that to find the best data, one must first understand the broad landscape of the information and then zoom in on the specific, fine-grained details that matter. By building a system that first groups data by its major themes and then carefully selects samples to cover the subtle variations within those themes, they have created a method that consistently outperforms existing techniques. In tests involving complex visual and reasoning tasks, this new approach allowed models trained on just a tiny fraction of the data to perform as well as, or even better than, models trained on the entire dataset.

The researchers, led by Peng Sun and colleagues from several universities, developed a method they call MASS, which stands for Manifold Aware Sparse Selection. To understand why this is necessary, imagine trying to organize a massive library where every book is described by a single, incredibly long list of keywords. If you try to find similar books by comparing these lists directly, you might end up grouping a book about a cat with a book about a car simply because both lists happen to contain the word "fast," even though the topics are completely different. The original lists are too cluttered with minor details and random noise to show the true structure of the library. The team realized that the data used to train modern AI models behaves similarly; the mathematical representations of these data points are so high-dimensional and complex that simple comparisons fail to capture their true relationships.

To fix this, the team designed a two-step process that mimics how a human might explore a new territory. First, they use a tool called a dense autoencoder to compress the complex data into a simpler, lower-dimensional form. Think of this as taking a high-resolution, three-dimensional map of a mountain range and flattening it into a clear, two-dimensional topographic map that shows the major peaks and valleys without the clutter of every single rock and bush. This step allows the system to identify the broad, major semantic groups within the data, effectively sorting the library into distinct sections like "science," "history," and "fiction" based on the dominant themes. This coarse grouping ensures that the selected data covers the main regions of the knowledge space.

Once the data is organized into these broad groups, the second step begins. Here, the researchers use a different tool, a sparse autoencoder, to look for the specific, fine-grained details within each group. While the first step handled the big picture, this step focuses on the unique characteristics that distinguish one example from another within the same category. It is like walking through the "science" section of the library and ensuring that the selected books cover not just general physics, but also specific sub-topics like quantum mechanics, astrophysics, and thermodynamics, rather than picking five books that all happen to be about the same narrow subject. This tool is designed to pick out only the most important features, ignoring the noise and redundancy that often confuse other methods.

The researchers tested this approach on two challenging datasets: one focused on general visual instructions and another on complex reasoning tasks involving images and text. They compared their method against nine other popular data selection techniques, including random selection and methods based on similarity or importance scores. The results were striking. Across different amounts of data, ranging from five percent to fifteen percent of the total pool, the MASS method consistently produced models that performed better than those trained with any other selection strategy. In several instances, the models trained with the MASS-selected subset actually surpassed the performance of models trained on the full dataset, achieving higher accuracy on benchmarks for visual understanding and logical reasoning.

The study also included a series of careful checks to understand why the method worked so well. The researchers found that skipping the first step and trying to group data directly from the original, complex representations led to poorer results, confirming that the initial compression was essential for finding the right structure. Similarly, using only the second step without the initial grouping failed to capture the full breadth of the data. The combination of both steps was crucial, with the first providing a stable framework and the second ensuring rich, detailed coverage. They also discovered that the system was robust to changes in the specific tools used for analysis, working well even when different types of data representations were employed.

One of the key insights from the work is that the quality of the data matters just as much as its diversity. The MASS system incorporates a quality score, allowing it to avoid selecting poor-quality examples even if they are unique. This ensures that the final subset is not only diverse but also reliable. The researchers noted that while their method requires some initial computational effort to set up the grouping and feature extraction, this cost is small compared to the savings gained by training on a much smaller dataset. The time spent preparing the data is a one-time investment that pays off by reducing the massive resources usually needed for full-scale training.

Ultimately, this work suggests a fundamental shift in how we think about teaching artificial intelligence. Instead of assuming that more data is always better, or that simply picking diverse examples is sufficient, the researchers show that the structure of the data itself holds the key. By respecting the natural hierarchy of information—first understanding the broad landscape and then filling in the specific details—AI models can learn more efficiently. The findings offer a practical path forward for developers who need to train powerful systems without access to infinite computing power or endless streams of data, proving that with the right selection strategy, a little data can indeed go a long way.

Technical Summary: Hierarchical Data Selection via Manifold Coverage and Sparse Feature Coverage in LLM Post-training

Problem Formulation

As the scale of supervised fine-tuning (SFT) data for large language models (LLMs) and vision-language models continues to grow, selecting high-value subsets from massive candidate pools has become critical for reducing training costs while maintaining or improving model performance. Existing data selection methods typically rely on measuring diversity directly within the original high-dimensional embedding space. However, this approach suffers from a fundamental limitation: geometric metrics in the raw embedding space often entangle dominant semantic directions, fine-grained supervision differences, and local noise.

The authors observe that for complex instruction data, the original embedding space exhibits high global linear complexity (requiring hundreds of PCA components to explain variance) but possesses a much lower effective dimensionality in local neighborhoods. This suggests that data distributions are better characterized as globally complex but locally low-dimensional nonlinear manifolds. Consequently, performing flat diversity measurement in the original space mixes global semantic structure, local fine-grained differences, and task-irrelevant perturbations into a single selection criterion, leading to suboptimal subset selection.

Methodology: MASS

To address these limitations, the paper proposes MASS (Manifold Aware Sparse Selection), a method that reformulates data selection as a coarse-to-fine hierarchical coverage problem. The core hypothesis is that a selected subset should first preserve coverage over major semantic manifold regions (coarse-grained) and then cover fine-grained supervision features within each region (fine-grained).

The MASS pipeline consists of three main stages:

Dual View Encoding:
For each candidate sample $z_i = (q_i, a_i)$ (input and response), two embeddings are constructed using a frozen embedding model:
- Condition Embedding ( $e^c_i$ ): Encodes only the input $q_i$ . This is used for learning the principal manifold structure.
- Supervision Embedding ( $e^s_i$ ): Encodes the concatenated pair $(q_i, a_i)$ . This captures the complete supervision signal and is used for fine-grained feature extraction.
Coarse-Grained Grouping via Dense Autoencoder (DAE):
A Dense Autoencoder is trained on the condition embeddings to learn stable, low-dimensional principal manifold coordinates.
- Training: The DAE is trained with a perturbation strategy (adding spherical noise to inputs) to ensure robustness. It compresses high-dimensional inputs into a low-dimensional latent space (set to 32 dimensions in main experiments).
- Grouping: The learned latent coordinates are used to cluster the data (e.g., via K-Means) into coarse-grained semantic groups. This step aims to capture the dominant semantic directions while suppressing redundant dimensions and local noise.
Fine-Grained Selection via TopK Sparse Autoencoder (SAE):
Within each coarse-grained group, a TopK Sparse Autoencoder is employed to extract sparse supervision features from the supervision embeddings.
- Feature Extraction: The SAE maps inputs to an overcomplete feature space, applies ReLU, and retains only the top $K$ activations (e.g., $K=64$ out of 131,072 features).
- Greedy Selection: Samples are selected sequentially within each group. The selection score combines two factors:
  - SAE Coverage Gain: A metric measuring how many new sparse features a candidate sample covers relative to the currently selected subset, normalized by feature activation magnitude.
  - Quality Score ( $\eta_i$ ): An external quality score (e.g., from a lightweight multimodal model) to avoid selecting low-quality samples merely for diversity.
- The algorithm greedily selects the sample with the highest combined score until the group's budget is filled.

Key Contributions

The paper outlines three primary contributions:

Problem Reformulation: The authors reformulate SFT data selection as a hierarchical coverage problem, explicitly separating coarse-grained principal manifold coverage from fine-grained supervision feature coverage.
MASS Framework: They propose the MASS algorithm, which integrates a Dense Autoencoder (DAE) for constructing stable principal manifold groups and a TopK Sparse Autoencoder (SAE) for within-group sparse feature coverage, augmented by external quality signals.
Empirical Validation: The method is validated across multiple datasets (Vision-Flan, LLaVA-CoT), data budgets (5%, 10%, 15%), embedding sources, and target models, demonstrating consistent effectiveness and robustness.

Experimental Results

The authors evaluated MASS on Vision-Flan (general instruction tasks) and LLaVA-CoT (reasoning tasks) using target models like LLaVA-V1.5-7B and Llama-3.2-11B-Vision-Instruct.

Performance: MASS consistently outperformed nine strong baselines (including Random, XMAS, COINCIDE, SemDeDup, PRISM, and EL2N) across all data budgets.
Surpassing Full Data: In several settings, MASS achieved performance that matched or surpassed full-data training using only a small subset (e.g., on Vision-Flan with a 10% budget, MASS achieved an Average Relative Performance (ARP) of 102.12, exceeding the full-data baseline of 100.00).
Ablation Studies:
- DAE vs. Raw Embeddings: Replacing DAE clustering with raw embedding clustering resulted in lower performance (101.59 vs. 102.12 ARP), confirming that principal manifold coordinates provide more reliable coarse grouping.
- Complementarity: Using only DAE (coarse grouping) or only SAE (fine coverage) yielded lower results than the combined approach, demonstrating that the two components are complementary.
- Configuration Sensitivity: While 32 dimensions were used in main experiments, sensitivity analysis revealed that a latent dimension of 64 achieved the best ARP (102.71), suggesting a moderate compression dimension better balances structure preservation and noise suppression. Additionally, a TopK sparsity of 64 yielded optimal results. Larger SAE feature dimensions (131,072 vs. 65,536) improved performance, indicating the need for sufficient feature capacity.
- Quality Signals: Removing quality scores reduced performance significantly, while changing the scoring model yielded consistent improvements, showing the robustness of the quality signal integration.
- Robustness: MASS maintained superior performance over random selection across different target model scales and families (e.g., LLaVA-13B and Qwen2-VL-2B).

Significance and Limitations

Significance:
The paper claims that MASS provides a new paradigm for embedding-based data selection. By acknowledging that complex instruction data resides on nonlinear manifolds with varying local dimensions, MASS avoids the pitfalls of flat geometric selection. It demonstrates that a hierarchical approach—preserving the global manifold structure while ensuring fine-grained feature diversity—leads to more efficient and effective data selection, often outperforming full-data training with significantly fewer samples.

Limitations:
The authors modestly acknowledge several limitations:

Dependency on External Signals: MASS relies on external embedding sources and quality scoring models. If these signals are severely biased, the selected subset may be affected.
Candidate Pool Constraints: The method operates within a given candidate pool. If the pool itself lacks certain task types or supervision patterns, MASS cannot compensate for these fundamental distributional gaps.
Task-Oriented Selection: The current work focuses on selecting subsets from general data collections for SFT. It does not address how to select subsets specifically for distinct downstream tasks, leaving that as an area for future exploration.