New metric reveals how data diversity drives LLM reasoning—and a framework to generate better training data

Training language models on diverse data leads to better generalization, but identifying what kind of diversity actually matters has proven elusive. A new paper from arXiv proposes that the answer lies in analyzing the gradients that flow through a model during training.

Measuring diversity through gradients#

Conventional diversity metrics often rely on surface-level differences—like paraphrasing or semantic variation—but these don't always correlate with how well models perform on out-of-distribution tasks. The researchers introduced G-Vendi, a metric that measures diversity by computing the entropy of gradients induced across a training dataset. Rather than examining the data itself, G-Vendi looks at how differently the model "sees" each example through its learned representations.

Remarkably, G-Vendi achieves a strong correlation (Spearman's ρ ≈ 0.9) with out-of-distribution performance across both natural language inference and math reasoning benchmarks, despite using a small off-the-shelf proxy model to compute the gradients.

Prismatic Synthesis: targeted data generation#

Building on this insight, the team developed Prismatic Synthesis, a framework for generating synthetic training data that deliberately targets underrepresented regions in gradient space. Rather than producing data randomly, the approach focuses on filling gaps where the model's gradient responses are sparse or redundant.

Testing on reasoning tasks showed consistent improvements in model performance as more synthetic data was added—not just on in-distribution benchmarks, but crucially on unseen, out-of-distribution test sets. The distilled 7B model trained with this approach, called PrismMath-7B, outperformed a larger baseline (R1-Distill-Qwen-7B) on 6 out of 7 challenging benchmarks, despite the baseline being trained on data from a 671B model versus their 32B source.

Scale and methodology#

The findings rest on over 300 controlled training runs, carefully matched for data volume and quality to isolate the effect of diversity itself. This scale of experimentation provides confidence that the correlation between G-Vendi and generalization is robust rather than an artifact of confounding variables.