This project was done as part of the Iliad Fellowship. L.J. did this work while at EleutherAI.

From functions to data

That we can learn efficiently from our environment at all tells us something about (1) the world and (2) ourselves: (1) Regularities in the world are simple enough to be picked up from limited experience. (2) Our inductive biases favour precisely those regularities. Learning, in this case, is effective because the two are matched. Considering a neural network with a fixed architecture lets us probe the second half of this match directly, by asking which (classes of) functions are easy or difficult for them to learn.

Neural networks learn to approximate a target function by minimising a loss over the training set We may ask: which (classes of) target functions are easy or difficult to learn for the model? This will depend on the loss function and the choice of the training set We say that the model generalises if it successfully learns to approximate beyond the training set By choosing the subset or, as we will see, by weighting it, we can control the generalisation behavior and hence what is learned. In this setting, our question becomes: which (classes of) data are easy or difficult to learn and generalise from?

The data-to-parameter map

Let's make this more formal. We assume: a dataset , a space of model parameters, a per-sample loss Given , we define the space , where assigns a weighting to each data point ^[1] To avoid confusion we call the parameters and reserve the word weights for the data weights. We define the data-to-parameter map as where corresponds to the parameters the of model after being trained on the target for steps, starting from the initialisation Everything else about training (architecture, optimiser, schedule) is fixed and is part of We write for when the initialisation is clear from context. This map is central to understanding and steering how data shapes a model: setting removes a data point, and setting upweights it.

Let's say we also have a differentiable evaluation function (or ) that captures a behaviour we care about. For the rest of this post we take to be the test loss, but emphasise that could be more general, e.g. representing the output of an eval. ^[2]

A good understanding of lets us interpret from the perspective of the model (and the associated training recipe): How does the model generalise from specific subsets? How does the dataset pair with the model's bias to produce a solution to the task?

The local view

A normal training run is given by setting and evaluating , i.e. we assign the same weight to every data point. This points out a fundamental limitation: Evaluating at a single costs a full training run ^[3] . So in practice we have little leverage to understand what looks like globally.

One way around this is to look at the local behaviour about , through the derivatives This allows us to approximate a leave-one-out effect: What do the parameters look like when we omit the data point ? Computing this derivative is thus cheaper compared to computing the leave-one-out retrainings. Variants of this have been done before: influence functions and susceptibilities fit this framework by changing what exactly means, see here for more details.

A basic limitation is that this is purely local:

• Higher-order terms. We don't know how far the first order approximation is from the actual effect of removing a point

• Group effects. The data weights may behave very differently away from : Once we remove a specific combination of data points, phenomena can appear or vanish that matter a lot for

Metagradient descent

To move beyond , we use metagradient descent (MGD) (Engstrom et al., 2025), an iterative outer optimisation over : (Here we denote by the data weights at meta-optimisation step , as a superscript to emphasise meta-optimisation as opposed to optimisation). In our version, weights each example's loss along the whole run, and is held fixed across outer steps. Each step re-linearises at the current , so the path can pick up group effects that the single derivative at misses.

MGD optimises for a single initialisation : every outer step differentiates through This means that the optimal can depend substantially on the randomness in the initialisation (and mini-batch if we don't keep it fixed) or it may transfer across different parameter initialisations

Experiments

Setup

We adopt the familiar modular addition setup in Power et al. (2022) and Nanda et al. (2023): a one-layer transformer is trained to compute the sum of two elements of the cyclic group for We train on a fixed of the input pairs and evaluate on the remaining pairs. Let denote the average of the validation loss over a set of training checkpoints (see Appendix A for details). We run steps of MGD on the per-example data weights (as distinct from the network parameters ), initialised at , minimising (see Appendix B for details). We identify the grok epoch as the first epoch where the test accuracy , and the delay is its gap from the train-accuracy- epoch. We read off which solution a grokked network has learned as its Fourier mode set (Appendix C).

Results

We summarise and highlight three observations here and refer to the appendix for more experiments we ran.

MGD moves grokking earlier

We run MGD from three parameter initialisations , and (seeds , and ). In each, the meta-objective decreases steadily and the replayed runs grok to earlier (Figure 5, Appendix D). To test transfer, we use each final (for seed ) as fixed data weights for ten distinct parameter initialisations , none of which were used in the meta-optimisation, each paired with a unit-weighted () baseline from the same initialisation. The -weighted runs grok to faster than their paired baselines ( on average); at each initialisation, all three -weighted runs grok before the baseline (Figure 1, full statistics in Table 1, Appendix D). Appendix E reports training on the most up-weighted (largest ) points alone.

Figure 1: -weighted training from ten new parameter initialisations vs. unit-weighted baselines. (a) Test accuracy (faint are the individual initialisations, bold is the mean). (b) Vertical gap between open and filled circles is the grokking delay.

Each biases the run towards the same solution

Gradient descent with the ten unit-weighted baselines finds ten different mode sets, with no mode shared by five or more runs (Figure 2a), so without weighting, the solution depends on the initialisation. The optimised weightings override this: the dominant set under is exactly at all ten initialisations, and under exactly at eight of ten, with mode missing in one and mode added in another (Figure 2c,d), while imposes its set only partially (Figure 2b).

Imposition is positively correlated with the grokking speedup: the data weights that produce the greater speedups over the parameter initialisations are also the more successful at transferring their Fourier mode bias (Figures 5a and 1b).

Figure 2: Rows are the initialisation seeds, columns are the Fourier modes, shade is the mode's share of the run's axis-Fourier logit energy, markers are the modes dominant in runs.

Each is spectrally structured

We conjecture that if imposes a mode set during training, then those modes should also be salient in the -D Fourier decomposition of over the grid. To check this, we embed on the grid (zero off the train set) and take its -D Fourier transform. The power at the axis sites and measures how much varies at frequency in alone or in alone, for the same modes as in Figure 2.

We report power as a z-score against a null that shuffles the values among the training points, so that a large z-score reflects the arrangement of the data weights rather than the split or the distribution of data weight values. For and , the strongest values in Figure 3 are axis sites of imposed modes: mode reaches and respectively (Appendix F). For , which imposes its set only partially, no imposed mode gets a high z-score.

Since many imposed modes are indistinguishable from chance and large values also appear outside the imposed sets, we take this as suggestive evidence that encodes its solution in Fourier terms, not as showing that the solution can be read off the spectrum. Off-axis, the spectrum of is dominated by the diagonal (weightings by answer class ), while for and there are large off-axis values whose structure we leave to future work.

Figure 3: z-scores of the -D Fourier power against the shuffle null, circles are the axis sites of the imposed modes.

Limitations.

Continuing the meta-optimisation well past steps eventually drives some data weights negative, at which point the inner minimisation becomes unbounded below and has no minimum. The vectors we use stop before this point. Each is also specific to the trajectory and observable it was optimised for, and the cost of exact metagradients limited how many variants we could run. A natural extension is to investigate whether continued meta-optimisation, once the negative data weight problem is addressed, makes the solution fully readable from the spectrum, and what role the off-axis structure plays.

Conclusion

We have defined the data-to-parameter map to formalise the relation between data and the model's bias to learn. This in turn frames the process of aligning models as a specification problem: which data mixture will induce a desired trait (Pepin Lehalleur et al., 2025)?

Since is costly to evaluate, studying its properties globally is intractable. On the other hand, understanding it locally is feasible. We suggest studying a compromise of the two by performing metagradient descent on the data weights

In a toy setting (modular addition), we already observed some interesting properties:

1. Iterative MGD can move grokking considerably earlier.

2. The solutions transfer across initialisations.

3. These solutions carry and transfer biases across initialisations: it seems that specific data sets bias models to learn a specific set of Fourier modes.

We see this as indicative that some datasets carry implicit biases, which can be discovered using MGD.

Acknowledgements

We are grateful to CoreWeave for providing the compute resources.

To reference this work, please cite it as:

@misc{exploring_data_landscape_2026,
  author      = {Louis Jaburi and Benjamin Shaffrey and Xavier Poncini},
  title       = {Exploring the data landscape via metagradient descent},
  year        = {2026},
  month       = {October},
  url         = {https://www.lesswrong.com/posts/iduSpeG2Y6kAf4uHk/exploring-the-data-landscape-via-metagradient-descent},
}

References

1. A. Power, Y. Burda, H. Edwards, I. Babuschkin, and V. Misra. Grokking: Generalization beyond overfitting on small algorithmic datasets. arXiv:2201.02177, 2022.

2. N. Nanda, L. Chan, T. Lieberum, J. Smith, and J. Steinhardt. Progress measures for grokking via mechanistic interpretability. ICLR, 2023. arXiv:2301.05217.

3. T. Kumar, B. Bordelon, S. J. Gershman, and C. Pehlevan. Grokking as the transition from lazy to rich training dynamics. ICLR, 2024. arXiv:2310.06110.

4. M. Andrychowicz, M. Denil, S. Gómez Colmenarejo, M. W. Hoffman, D. Pfau, T. Schaul, B. Shillingford, and N. de Freitas. Learning to learn by gradient descent by gradient descent. NeurIPS, 2016. arXiv:1606.04474.

5. L. Metz, N. Maheswaranathan, J. Nixon, C. D. Freeman, and J. Sohl-Dickstein. Understanding and correcting pathologies in the training of learned optimizers. ICML, 2019. arXiv:1810.10180.

6. L. Engstrom, A. Ilyas, B. Chen, A. Feldmann, W. Moses, and A. Madry. Optimizing ML training with metagradient descent. arXiv:2503.13751, 2025.

7. S. Pepin Lehalleur, J. Hoogland, M. Farrugia-Roberts, S. Wei, A. Gietelink Oldenziel, G. Wang, L. Carroll, and D. Murfet. You are what you eat – AI alignment requires understanding how data shapes structure and generalisation. arXiv preprint, 2025. arXiv:2502.05475.

Appendices

The appendices cover the choice of observable (Appendix A), the experimental setup and training (Appendix B), the solution readout (Appendix C), the meta-optimisation runs and their speedups (Appendix D), training on subsets selected by each compared with random subsets (Appendix E), and the spectral analysis of the (Appendix F).

Appendix A: Choice of observable

The observable is the average of the validation loss over a log-spaced grid of checkpoints, The justification for this is that the endpoint loss saturates after grokking and its metagradient shrinks significantly. The checkpoint average on the other hand retains gradient signal wherever the transition occurs, which becomes especially important if the transition is moved earlier due to successful metagradient optimisation. This idea can be found in Andrychowicz et al. (2016) where they replace the endpoint meta-objective with a trajectory average to ensure the outer problem (the meta-optimisation) retains signal along the whole trajectory. In Metz et al. (2019), it is shown that long unrolls can make even this averaged objective highly non-smooth, which they address with variational smoothing.

Appendix B: Details of the setup and training

We used a transformer with one attention layer and one MLP layer with , , ReLU, no LayerNorm and learned positional embeddings. We used a single head rather than the four-head model of Nanda et al. (2023). The model takes as input a pair , presented as tokens with logits read at the final position, and must maximise the logit of the sum We trained on a fixed of the input pairs () with full-batch AdamW at a constant learning rate of with , so one epoch is one gradient step. Writing for the set of all held-out pairs, the test accuracy is In order to make the metagradient optimisation less expensive we make use of the finding in Kumar et al. (2024) that the grokking transition can be advanced or postponed by adjusting various hyperparameters. Accordingly, we use a weight decay of in order to obtain a sharp grokking transition within a window of epochs (memorisation occurs around epoch and grokking between epochs and with test accuracy at chance for roughly the first half of the delay window, Figure 4). We use a single seed to fix the train/test split across all experiments, so that each data weight always refers to the same training point. We only vary the parameter initialisation of the model. For the MGD steps, we split the held-out pairs in half: a meta-validation half, whose loss enters the observable that the meta-optimisation minimises, and a meta-test half, whose loss is used to verify that the meta-optimisation is not overfitting. This split is used only inside the meta-optimisation. The test accuracy above and the grok epoch are always computed over all held-out pairs. For the outer loop we use Adam on the data weights with a step size on normalised, clipped metagradients.

Figure 4: The three base training runs (parameter initialisations , and ) at weight decay : train accuracy saturates around epoch in all three, and the grok epochs are , and , within the horizon

Appendix C: Solution readout

As shown in Nanda et al. (2023), a network grokked on modular addition computes its logits, to a very good approximation, as a small Fourier sum in , with only a small number of the distinct frequencies carrying substantial amplitude. We fit this to a trained model's logits greedily, adding frequencies in descending order of energy until is the number of frequencies the fit adds. We then define the dominant frequencies as the fitted frequencies with energy at least of the strongest mode's. These form the mode set , with , which indicates the solution that the training run converged to.

Appendix D: Meta-optimisation curves

Figure 5a shows the meta-objective at each of the outer steps for the three parameter initialisations. It decreases steadily in all three runs, least for seed , whose vector imposes its mode set only partially. Figure 5b shows, for , the validation-loss curve of the replayed training run at each outer step. As the data weights are updated, the drop in validation loss that marks the grokking transition moves earlier. The grok epoch of the final -weighted run is , against for the unit-weighted run at the same initialisation ( earlier, see Table 1).

Figure 5: MGD at the three parameter initialisations ( outer steps). (a) (solid) per outer step; the dashed curves show the same average over the meta-test half (), which tracks closely at all three seeds, so the meta-optimisation is not overfitting to the meta-val half. (b) Per-step inner val-loss curves (). The grok epoch moves from to (seed ), to (seed , shown in b) and to (seed ).

	grok epoch	delay	speedup (mean)	speedup (range)
			–	–
				–
				–
				–

Table 1: Grok epoch and delay (mean s.d. over the ten parameter initialisations ) and the grok-epoch speedup over the paired unit-weighted baseline.

Appendix E: Subset retraining

From each we construct subsets of the training data such that we obtain the following:

• Helpful. The most up-weighted points.
• Unhelpful. The most down-weighted points.
• Mixed. The most helpful and the most unhelpful.
• Random. A random subset of points.

We retrain on each subset alone (all data weights set to ) for , with ten different parameter initialisations for 100,000 epochs.

The helpful subsets generalise at sizes where the controls do not. At the helpful subset of reaches test accuracy at initialisations while every control condition is at (: , : ). At all three helpful subsets reach at , against at most for any control. Grokking separates in the same way: at the helpful subsets of and grok at (every control at most ), and at the median grok epoch of the helpful runs that grok ranges from k to k across the three vectors, against k to k for the control runs that grok (Figure 6). The strength of the effect follows the vectors: , the weakest vector in the earlier sections, is also the weakest here. The unhelpful subsets perform no worse than random at any and grok frequently at , so the down-weighted points appear to be redundant rather than harmful. Below the sizes at which they grok, the helpful subsets get stuck on stable partial solutions. Running far beyond epochs shows these plateaus are permanent rather than slow (Figure 7). These sizes refer to this particular train set and not to train sets in general.

Figure 6: Fraction of the ten runs per cell reaching test accuracy (a) and (b) vs. subset size . Colours are the three vectors, line styles are the subset conditions; the random controls (black dashed) are shared across vectors.

Figure 7: Sub-critical helpful-subset plateaus are arrested, not slow.

Appendix F: Per-mode axis power

Figure 8 gives the per-mode results behind the spectral analysis in the main text. The left panels show each field on the grid; held-out pairs are white, since they carry no data weight. The right panels show the axis power for each of the modes as a z-score against the shuffle null, with each vector's imposed modes highlighted and marked for reference. For and , mode is among the largest peaks ( and ), and for mode also stands out (). The remaining imposed modes have axis power indistinguishable from chance. For , none of the imposed modes stand out, and for all three vectors several strong peaks fall outside the imposed sets.

Figure 8: Left: the fields on the grid (red means up-weighted, blue means down-weighted, white means unit-weighted or held out). Right: excess power at the axis sites per Fourier mode, in units of chance (z-score against the shuffle null, dotted line at ). Highlighted bars and markers are each vector's imposed modes.

1. In general, we could also take We chose to stick with the current formulation to exclude the case where we optimise a negative loss function. ↩︎

2. In our experiments averages the validation loss over several checkpoints, so strictly returns the trajectory. This was an engineering choice to make the experiments more stable. See Appendix A for details. ↩︎

3. Albeit at a lower cost if is very sparse. ↩︎