Prismatic Synthesis: Gradient-based Data Diversification Boosts Generalization in LLM Reasoning

We sincerely thank the reviewer for the positive feedback and constructive questions! We are glad that you find our idea to be nice and well-motivated, our analyses to be thorough, our work to be empirically strong and well-written! Below, we address each of your questions:

Controlled baselines for data generation?

We appreciate the reviewer’s concern. We would like to gently refer to Section 3.2.2, where we ask whether diversification in gradient space is truly necessary — conducting controlled experiments via naive scaling & heuristic diversification under the same condition as in Prismatic Synthesis. In addition to the existing results, here we add an additional baseline that leverages embedding diversity. In the same controlled setup as in our main setup, we leverage embedding-based clustering instead of gradients in Prismatic Synthesis:

The results at 10k scale is shown above. We find that Embedding-based diversification, just like the other baselines from Section 3.2.2, does not lead to improvements comparable to Prismatic Synthesis. Overall, these results demonstrate the high utility of gradient information in capturing diverse reasoning patterns.

Prospect of “online proxy model”

We sincerely appreciate this great suggestion! In this work we did focus on a fixed proxy model, in particular because even the impact of diversity on a static dataset has not been widely understood. But updating the proxy model throughout the course of model training would be an interesting and reasonable next step to explore. We believe this would holistically improve the relevance of gradient information we get through the proxy model, particularly in an active learning setup (e.g. dynamically selecting diverse subsamples during training).

Relative cost of gradient entropy

As expected, much of the operation we perform in G-Vendi computation is similar to canonical gradient descent. For example, we do forward propagation once as in the typical mini-batch gradient descent. The only difference is that for G-Vendi, we backpropagate for each sample loss individually, as we want to represent each sample with their respective gradient vectors. Other overheads such as computing the entropy from collected gradient features are constant, thus is negligible compared to the gradient collection stage.

We thank the reviewer again for thoughtful feedback!

We thank the reviewer for both the positive and thoughtful comments! We are happy that you find our work to be of significant benefit for LLM data curation, our experiments to be well-designed, and the overall presentation to be clear. Below, we address each of your concerns.

Intermediate dataset between DiverseReason and DiversePersona

We agree that adding another dataset in between the two extremes could further strengthen the findings of our analyses. Instead of adding a heuristic pipeline that may lie in the middle of the two datasets, we take a straightforward approach of taking 50%-50% mixture of the two as an intermediate dataset with 10k samples. The results are shown above. The performance of the mixture goes in between the two datasets, with consistent trend as in the original finding (G-Vendi captures relative diversity in reasoning patterns, while embedding operates in an opposite direction).

In addition, we note that the experiments in Table 2 are designed to analyze the feature-level difference between gradient and embedding diversity, by following the two extreme data generation pipelines for didactic purposes. The sensitivity of G-Vendi against baseline metrics is best shown in our main correlation analyses in Section 2.3.1 and Section B, where we consider a realistic data pool and sample many subsets in a continuous spectrum of diversity.

Larger proxy models?

For the practical applicability of G-Vendi, we assess that using a small proxy model (of ≤ 1B scale) is a fundamental necessity — as larger models would lead to proportionately more resources, just to collect the gradients. In this sense, the fact that G-Vendi can highly correlate with performance even while using a small proxy model is a strong appeal for our approach!

Additionally, for other aspects of our experiments such as the main results for Prismatic Synthesis (Section 3.2) or ablation study on student models (Section 2.3.4), we leverage diverse set of models both in terms of model family (Llama and Qwen) and scale (up to 8B).

Again, thank you for the positive and thoughtful feedback!

We thank the reviewer for the positive assessment of our work and the thought feedback! We are glad you found our work to tackle fundamental problem in LLM development, our G-Vendi to be theoretically well grounded and effective, and our work to be of real utility. Below, we address each of your concerns.

How many iterations were used?

We do not fix the number of iterations, but limit the number of new samples generated at each iteration to be 10k (that is, after generating 10k samples, we move on to the subsequent step 3 in line 245). We iterate this process until we reach the target size of 1M. We will add this to the implementation details section in the revised version.

On the notion of task similarity

Thank you for pointing this out. You are correct that our setup assumes relevance between the training and evaluation tasks — e.g., training on math problems to improve model’s math reasoning skills — while not assuming any access to the evaluated benchmarks. This level of loose relevance is both realistic and practical, particularly in the context of curating reasoning data for LLMs — as even the state-of-the-art reasoning datasets (e.g. Open-R1, MetaMath, SynLogic) specializes to specific domains in improving reasoning capabilities.

But then, is reasoning data all about aggregating more datapoints for a specific domain? Our experiments suggest the opposite — even when focusing on a specific task, the diversity of reasoning patterns significantly impacts model performance. While it is true that we use Open-R1 as seed with the goal of improving math reasoning, we still see a significant improvement from not only the Open-R1 itself, but also larger datasets with manual curation, by strategically improving diversity. Overall, these results show that domain relevance, while generally important, is only an auxiliary factor behind our strong results.

In terms of task breadth, we apply Prismatic Synthesis to two widely different tasks, namely generative math reasoning tasks based on long CoT, and discriminative NLI task with limited label space. We are already working on further application of Prismatic Synthesis beyond these two domains, which we believe is beyond our current scope that primarily introduce data diversity measure in a principled framework.

Potential failure modes of G-Vendi

Extending our discussion in Appendix D, we note that G-Vendi is not a "universal” metric that can predict whether one model would perform better on a specific benchmark. Importantly, our measure of model performance (with which G-Vendi strongly correlates) is an aggregation of metrics across multiple benchmarks, rather than targeting a specific test set. If one’s goal is to optimize to a specific benchmark, minimizing the distance between training distribution and the benchmark distribution could be more effective than increasing gradient diversity. But we view this as a difference by design, rather than an inherent limitation of G-Vendi.

We also note that the gradient representation can often be hard to interpret (unlike examples shown in the appendix). We will include hard-to-interpret examples in our revision as well.

Prismatic Synthesis Overhead

We iterate gradient collection and clustering on every 10k generations. Since the projected gradient prior to the current steps are already saved in a storage, we only need to collect the gradients for the new 10K samples using a 0.5B proxy model. For clustering, we specifically implement Lloyd’s algorithm on GPU; since prior samples are already almost clustered, adding new 10K samples to the pool typically does not require more than 3 E-M steps in the Lloyd’s algorithm. Overall, we find that gradient collection on 1 H100 node typically takes less than 1 minute per step, and the clustering takes no more than 1.5 minutes per step in our experiments. Therefore, the computational overhead in Prismatic Synthesis is negligible compared to the the cost of distilling from larger models (which could take more than 1 hour on 1 H100 node for 10K samples with 32K max sequence length), which all baseline curation approaches share.

In addition, we show in Table 4 that Prismatic Synthesis is substantially more model-efficient — our dataset outperforms state-of-the-art datasets generated via R1-671B, which is 21 times larger than our data generator (R1-32B). The efficiency gain from using a order-of-magnitude smaller teacher far outweighs that of the computational overhead describe above, adding to the practicality of our approach.

G-Vendi Scalability

As suggested by the reviewer, we measure the computational requirements to compute G-Vendi score. For the dataset size of 100k on 1 H100 node, G-Vendi computation took approximately 8m 30s. Most of the runtime is spent on iterating over the dataset and encoding each sample, which even the embedding-based baselines need to go through as well. Notably, aggregating Vendi score out of the collected gradients took less than 20 seconds, despite its theoretical complexity of O(d^2|D|). For memory requirements, the computation took no more than 20GB of GPU ram. This contrasts with using a 7B embedding model (as done in our baseline setup), which took 14GB of ram just to load the model on the GPU — indicating the efficiency of using small proxy model in G-Vendi. We appreciate the suggestion and will include this in the revision.

Scale Constraints

We would like to note that the nature of our analyses requires hundreds of iterations of experiments under a controlled setup, which can realistically be implemented with reasonable size models. Importantly, as we primarily investigate diversity of training data, we study the effect of data scale with respect to diversity, incorporating experiments from small controlled setup (e.g. 10k) to a large scaling long CoT training (e.g. 1M). But we also agree that understanding the impact of data diversity additionally with relation to model scale (e.g. > 70B) is an interesting and less-explored agenda, which we hope our work sets a foundation for.

How was the projection dimension chosen?

The dimension was empirically chosen based on our preliminary study on the correlation between G-Vendi scores with different projection dimensions on NLI data pool.

The results on 10k subsets are shown above. We find that dimensions beyond 1024 are very highly correlated — that is, using 1024-dim projection leads to 8x reduction in storage requirement than 8192, while almost perfectly preserving the relative positions of subsets. But we also find that further decreasing the projection size could harm the correlation against larger dimensions. Based on these findings we set projection dimension to 1024 in the subsequent experiments.

Clustering Parameters

Thank you for the suggestion. We conduct 2 ablation studies on the impact of values of $k$, the number of clusters (defined as ratio of number of all samples) and $N$, the number of sparsest clusters to leave in Prismatic Synthesis. Using our math data generation pipeline, we generate 50k datasets with different values of k and N, and report the yield (ratio of samples passing the cluster filtering stage) and the G-Vendi score. When ablating $k$, we fix $N = k / 2$ as in our default setup.

The results are shown in the above tables. We first find that our pipeline is robust to the specific values of $k$, exhibiting no conspicuous trend in G-Vendi or Yield. In addition, the result shows that the value of $k = 1\%$ leads to small improvement in the G-Vendi compared to larger values. On the impact of N, we see a tradeoff between yield / G-Vendi -- lower N leads to lower yield (i.e. needs more samples to get to a fixed data size), but achieves higher diversity on a fixed data size. These two effects ultimately offset each other in a long run (thus as reasonable values of $N$ would suffice), but we find that taking the middle value of $k/2$ adds to the stability by balancing the yield and diversity improvement.

We thank the reviewer again for the positive and constructive suggestions!