Learning the ropes of the susceptibility technique by replicating part of “Patterning”
Introduction
A main thing that makes alignment hard is the slipperiness of generalization: two models that perform the same on a training distribution may implement distinct algorithms, meaning they’ll generalize differently on inputs outside that distribution. It would be nice to gain control over which algorithm a model implements when multiple choices could achieve the same training loss. Wang & Murfet (2026) introduced patterning, a method that uses concepts from singular learning theory (SLT) to achieve just that sort of control. In particular, patterning let me characterize kinds of data points in the training set that, when overrepresented, bias the model to learn different internal structures. In this post I describe my replication of one of their experiments and my results from another similar experiment. The code for these experiments is at https://github.com/henrylhale/patterning-brackets-repro.
What are we trying to do?
We want to find settings where a model might learn one of a few algorithms for a task, and force it to learn one or the other by picking training examples to overrepresent in training.
Replication task
The first task, the one from Wang & Murfet (2026), is binary classification of even-length strings of parentheses of length at most 40. The setting is originally from Li et al. (2025).
We label a string TRUE if it was balanced in the usual way, like “()”, “(()())”, or “()((()))”.
We label a string FALSE if it was not so balanced, like “)(”, “())(()”, or “((”.
On the one hand, a wrong heuristic you might use on this problem is the Equal-Count algorithm, which outputs TRUE for strings with as many open parens as close parens.
On the other hand, we call the correct algorithm for this problem Nested.
We can construct a data set where Equal-Count was just as accurate as Nested by removing the examples where Nested was right and Equal-Count was wrong: the ones where the count was equal, but the nesting was wrong, like “)(” or “())(()”.
Suppose we train a model on this modified data set to high accuracy. We can then measure its out-of-distribution accuracy on the removed examples. If the OOD accuracy was high, then the model was implementing Nested, and if the OOD accuracy was low, then the model was implementing Equal-Count.
Wang & Murfet, following Li et al., trained populations of small transformers on this task. When you train a population of such models to high in-distribution accuracy with a fixed hyperparameter configuration and different seeds for the data set shuffle and weight initialization, the models end up with some distribution of OOD accuracies. Finding a good baseline to use patterning on meant finding hyperparameters that result in a bimodal distribution of OOD accuracies, i.e. one where many models are clearly implementing Nested or Equal-Count. The goal was then to find strings to overrepresent in the data set such that this distribution of OOD accuracies collapsed to a single mode.
Extension task
The other task I looked at, extending Wang & Murfet, is a binary classification of even-length strings with the token vocabulary ‘(’, ‘)’, ‘{’, ‘}’.
We label a string TRUE if it was balanced in the way one would expect from programming languages, like “()”, “{({})}”, or “(){{()}}”.
We label a string FALSE if it was not so balanced, like “)(”, “({)}”, or “({”.
One wrong heuristic you might use on this problem is the Forgetful algorithm, which outputs TRUE for strings that are balanced if you “forget the structure” and read the strings as if they used a single type of bracket.
Let’s call the correct algorithm for this problem 2-Nested.
Proceeding in the same way as in the original problem, I constructed a limited data set where Forgetful was as accurate as 2-Nested by removing the examples where 2-Nested was right and Forgetful was wrong.
I wasn’t able to find hyperparameters that resulted in a bimodal distribution of models between these two algorithms: many configurations saw all models collapse to Forgetful, and the best I could do was a unimodal distribution around 60-70% OOD accuracy, with a few models at 0% OOD accuracy and no models near 100% OOD accuracy.
However, applying the patterning method, I was able to get the mean OOD accuracy to move in both directions. It doesn’t seem like the movement was enough to be significant.
How do we find data points to repeat?
I’ll give a description of the theory at two levels of detail. The first is a sketch of how the pieces fit together logically. The second gives more motivation and more explicit definitions, but still black-boxes most of the derivations. Wang and Murfet (2026) gave a more detailed and complete account of the theory.
Sketch
- We have a model parameterized by weights \(w \in W\), and sequence-token pair training examples \((x, y)\).
- Consider some arbitrary “observables”: functions on weight space \(\varphi_i: W \to \mathbb{R}\). Let \(\mu^\infty_i\) be the expectation of \(\varphi_i\) as the number of training examples goes to infinity.
- Consider an arbitrary perturbation \(q_h\) of the ground truth distribution \(q\), parameterized by coordinates \(h\).
- We can define a matrix of susceptibilities \(\chi\), of shape [number of observables] x [dim h], such that \[d\mu^\infty = \chi d h.\] The susceptibility is the derivative of \(\mu\) with respect to \(h\), evaluated at \(h=0\) (i.e. when the data distribution has not been perturbed at all).
- One useful parameterization Wecan take is to think of \(q_h\) as a sum of perturbations towards constant distributions \(\delta(x, y)\). On this parameterization, \(\chi\) has a column for every training data point; its entries can be denoted \(\chi_{xy}^{\varphi_i}\).
- Say Wewanted to effect some change in \(\mu^\infty\); call it \(d\mu^\infty_{\text{TARGET}}\). It can be shown that the optimal intervention in \(h\) (i.e., the optimal way to overrepresent data points) to achieve this is \[d h_{\text{OPT}} = \chi^\dagger d\mu^\infty_{\text{TARGET}},\] where \(\chi^\dagger\) is the Moore–Penrose pseudoinverse.
- Unrelatedly, at each local minimum of the loss landscape, Wecan compute a quantity \(\lambda\), the local learning coefficient, and it turns out that models prefer to converge to minima where \(\lambda\) is lower.
- This quantity is not tractable to compute, but it admits an approximation \(\hat{\lambda}\) that Wecan estimate with an “SG-MCMC sampler” over a certain posterior distribution. Also, \(\hat{\lambda}\) is defined even away from local minima.
- The same sampler over the same distribution is used in computing the components of \(\chi\).
- The key insight of Wang and Murfet in this experiment was that one can define observables \(\varphi_{\text{N}}\) for Nested and \(\varphi_{\text{EQ}}\) for Equal-Count such that \(\mu^\infty_{\text{N}}\) approximates \(\hat{\lambda}\) (and thus also \(\lambda\)) at points in weight space near models that implement Nested, and likewise for \(\mu^\infty_{\text{EQ}}\) near models that implement Equal-Count.
- The hypothesis was then that applying data perturbations that minimize or maximize \(\Delta\mu := \mu^\infty_{N} - \mu^\infty_{EC}\) should respectively increase or decrease the chance that the model learns Nested rather than Equal-Count. (The signs go this way because models prefer low-LLC basins, so making \(\mu^\infty_N\) small relative to \(\mu^\infty_{EC}\) is what makes the Nested basin relatively more attractive.)
- It can be shown that the contribution to \(\Delta\mu\) from a data point \((x, y)\) is proportional to \(\chi_{xy}^N- \chi_{xy}^{EQ}\), a “susceptibility gap”.
- So Wecould sample some medium number of points from the data set, compute susceptibility gaps for them, and find interpretations for the top- and bottom-susceptibility-gap sets of data points.
- Retraining on new synthetic data sets with more points like those should then bias the models towards learning one algorithm or the other.
A note on sampling
Stochastic gradient Markov chain Monte Carlo (SG-MCMC) is a family of techniques for approximating expectations \(\mathbb{E}_\pi[\varphi(w)]\) under a target distribution \(\pi\) that is too complicated to sample from directly. The idea is to construct a Markov chain \(w_0, w_1, \ldots, w_T\) whose stationary distribution is \(\pi\); once the chain has mixed, the time average \(\frac{1}{T - T_b}\sum_{t=T_b}^{T} \varphi(w_t)\) over the tail of the chain approximates the desired expectation. The “stochastic gradient” piece refers to using a minibatch loss in place of the full-dataset loss inside the chain’s update rule, which makes SG-MCMC tractable on large training sets.
The specific algorithm I used is stochastic gradient Langevin dynamics (SGLD), with update rule
\[ w_{t+1} = w_t - \tfrac{\varepsilon}{2}\,\nabla\!\bigl(n\beta\, \tilde{L}_n(w_t) + \tfrac{\gamma}{2}\|w_t - w^*\|_2^2\bigr) + \sqrt{\varepsilon}\,\eta_t,\qquad \eta_t \sim \mathcal{N}(0, I), \]
where \(\tilde{L}_n\) is the minibatch loss estimate, \(\varepsilon\) is the step size, and the term \(\frac{\gamma}{2}\|w - w^*\|_2^2\) localizes the chain to a neighborhood of \(w^*\). As \(\varepsilon \to 0\) the stationary distribution of this chain is the localized Gibbs posterior,
\[p(w;\, w^*, \beta, \gamma) \propto \exp\bigl(-n\beta L_n(w) - \tfrac{\gamma}{2}\|w - w^*\|_2^2\bigr).\]
Both the LLC estimator \(\hat{\lambda}\) and the per-sample susceptibilities \(\chi_{x}\) defined below are expectations under this same posterior, so a single SGLD chain at some model checkpoint \(w^*\) is enough to compute both.
The local learning coefficient
Let \(p(y|x, w)\) be a sequence-prediction model parameterized by weights \(w \in W\) that predicts tokens \(y\) given a sequence of tokens \(x\).
For a point \(w\) in weight space and a sequence-token pair \((x, y)\), define
\[ \ell_{xy}(w) := -\log p(x|y, w).\]
Then considering a series of sequence-token pairs \((x, y)\) drawn i.i.d. from some ground truth distribution \(q(x, y)\), define the empirical loss on \(n\) data points at \(w\),
\[ L_n(w) = \frac{1}{n}\sum_{i=1}^{n} \ell_{xy}(w), \]
which is an approximation of the theoretical loss at \(w\)
\[ L(w) = -\mathbb{E}_{q(x,y)} \log(p(x|y, w)).\]
At a local minimum of loss \(w^*\), the local learning coefficient \(\lambda(w^*)\) is defined as follows: let
\[ \text{vol}(t, w^*) = \int_{w \in B \text{ such that } |L(w^*) - L(w)| < t} dv, \] and let \[ \lambda(w^*) = \lim_{t \to 0^{+}}{\log_2[\text{vol}(\tfrac{1}{2}t, w^*)/\text{vol}(t, w^*)]}.\]
This quantity measures “how degenerate” the loss landscape is in the neighborhood of \(w^*\): if there are redundant dimensions of weight space near \(w^*\), then its basin is flatter and its LLC is lower. The slogan I have been using for myself when thinking about the LLC is “how fast does the bath drain”.
The central theoretical result of singular learning theory is that as the number of training examples increases, models are likely to end up near those local minima of the loss landscape whose LLCs are lower.
To validate this empirically requires a computationally tractable estimate of the LLC. Consider
\[\hat{\lambda}(w^*) := n\beta\bigl[\mathbb{E}^{\beta}_{w|w^*, \gamma}[L_n(w)] - L_n(w^*)\bigr],\]
where the expectation \(\mathbb{E}^{\beta}_{w|w^*, \gamma}\) is taken with respect to the localized Gibbs posterior \(p(w; w^*, \beta, \gamma)\) defined above.
Lau et al. (2023) found that \(\hat{\lambda}\) is a faithful approximation of \(\lambda\) in simple cases where the formally defined LLC is computable. Hoogland et al. (2024) showed that models empirically tend to converge to regions of weight space where the value of \(\hat{\lambda}\) is lower, as one would expect for \(\lambda\). My understanding is that the faithfulness of \(\hat{\lambda}\) has not been formally verified, but these two findings provide good evidence for it.
In addition to being faithful and computable, \(\hat{\lambda}\) has the advantage of being defined even away from local minima of loss.
So the goal is to pick data examples to overrepresent in training such that \(\hat{\lambda}\) goes down in the basin of weight space where our preferred algorithm lives.
Susceptibilities and patterning
[ I removed this section from this draft of the post; I should have a version I’m happy with in at most a few days ]
Doing the experiments
These were the steps I took in each of the two experiments. There is a notebook in the associated repo for each of these steps for each experiment.
- Found hyperparameters that produced a baseline bimodal distribution of OOD accuracies with modes near 0 (models implementing the heuristic algorithm) and 1 (models implementing the correct algorithm). In the replication experiment, I just used the hyperparameters from the original experiment.
- Found hyperparameters for the SGLD sampler that produced an equilibrium where the draws were probably “typical” points in the loss landscape as described in Timaeus’s guide to SG-MCMC sampling. I again just verified that the hyperparameters from the original worked, but I did a sweep anyway to learn the skill.
- Used the sampler to compute susceptibility gaps for 1024 random data points with respect to the 3 models with the most extreme OOD accuracies in each direction; gave an interpretation of the data points with the top susceptibility gaps. Defined algorithms to generate more synthetic examples similar to these interpretable sets of points.
- Retrained a new population of models with synthetic data added in and checked the new OOD accuracies.
Replication experiment
Step 1: Architecture and baseline run
The architecture was Karpathy’s minGPT transformer with 3 layers, 4 attention heads and 64-dimensional embeddings. We say the model predicts TRUE for a string if its next-token logit for ) is greater than its next-token logit for (. Training used Adam with lr = 1e-4, weight decay = 1e-3, \(\beta = (0.9, 0.999)\), and batch size 8 for 5 epochs over a 200,000-sample dataset (100k well-formed, 100k neither well-formed nor equal-count). All hyperparameters were taken from Li et al.’s code repository and Wang & Murfet’s Appendix B.1, but finding bs=8 was the longest debug step in this project. Li et al.’s default transformer.yaml sets bs=8, but neither paper called it out as critical, so for a while I was using bs=32, which resulted in the models learning only Nested.
In addition to collecting the final OOD accuracy, I also followed Li et al. and collected the minimum and maximum OOD accuracies achieved during training (taking a snapshot every 500 gradient steps and conditioning on in-distribution accuracy of at least 0.9).
The baseline population of models came from 125 (init seed, shuffle seed) pairs with no repeated seed values. The resulting OOD distribution was bimodal at the endpoints as desired.
Step 2: Tuning the SGLD sampler for susceptibility computations
As described above, computing susceptibilities requires sampling the localized tempered Gibbs posterior
\[ p(w; w^*, \beta, \gamma) \propto \exp\!\bigl(-n\beta\, L_n(w) - \tfrac{\gamma}{2}\,\|w - w^*\|_2^2\bigr) \]
via SGLD. I used the devinterp SGLD sampler implementation, devinterp.optim.sgld.SGLD. The hyperparameters for the sampler are the step size, \(\varepsilon\), the strength of localization to \(w^*\), \(\gamma\), and the quantity \(n\beta\), where \(\beta\) is the inverse temperature and \(n\) is the dataset size. For a second time I just took the hyperparameters from Wang and Murfet; they’re in Appendix C.1:
- \(\varepsilon = 3 \times 10^{-6}\)
- \(\gamma = 500\)
- \(n\beta = 100\)
- SGLD minibatch size 32
- 100 burn-in steps, then 500 draws per chain with no thinning
To validate that these transferred to my checkpoints, I generated loss traces for 2 chains at each of a few hyperparameter values, imitating the visuals in the (very helpful!) Timaeus sampling guide (2026).
Here is a loss trace of a single chain at the final hyperparameters:
Step 3: Computing susceptibility gaps on random data points
[I’ll probably add a more complete explanation of the susceptibility calculation here by 5/5 or so]
I sampled 1,024 evaluation points (512 True + 512 False) uniformly from the training distribution and used this same eval set across every sampler chain on every model in the population.
For each trained model \(w^*\) in the baseline population I ran 4 SGLD chains at the hyperparameters fixed in Step 2. With Wang and Murfet’s choice of observables \(\varphi_N(w)\) and \(\varphi_{EQ}(w)\), I aggregated across chains and seeds as follows. Sorting the population by final OOD accuracy, I took \(\chi_{\text{N}, xy}\) = mean over the top-3 seeds (the Nested-cluster representatives, OOD \(\to 1\)), and \(\chi_{\text{EC}, xy}\) = mean over the bottom-3 seeds (the Equal-Count-cluster representatives, OOD \(\to 0\)). The susceptibility gap is then
\[ \Delta\chi_{xy} \;:=\; \chi_{\text{N},\,xy} \;-\; \chi_{\text{EC},\,xy}. \]
Here are my plots of the Dyck paths for the strings with the most extreme values of \(\Delta\chi\). They pretty much agree with the interpretations given in Wang & Murfet.
Step 4: Retraining
Since my characterizations of the points with the highest susceptibility gaps agreed with those of Wang & Murfet, I followed their recipes for generating new synthetic data:
- Almost-Nested dataset: start with the baseline dataset, remove 36,600 random False examples, add 18,300 unique i.i.d.-sampled strings that consist of a valid string followed by
)), each appearing twice. - Almost-Equal dataset: start with the baseline dataset, remove 67,100 random False examples, add 19,000 unique i.i.d.-sampled strings where the difference between the numbers of types of parentheses is at most 2 and the Dyck path has taken at least as many steps below 0 as above 0 at every prefix, each appearing four times.
These recipes upweight strings that the opposite algorithm gets wrong. Almost-Nested samples are FALSE strings with a correct prefix, and Almost-Equal samples are FALSE strings whose total open/close counts nearly match. I trained 125 models on the Almost-Nested dataset and 120 on the Almost-Equal dataset.
As in Wang & Murfet, Almost-Nested produced a clean shift: 123 of 125 retrained seeds had OOD accuracy < 0.05. The population trained on Almost-Equal is also consistent with Wang & Murfet: the mean OOD accuracy increased, but the OOD accuracy of a majority of models remained below 0.5.
Extension experiment
Step 1: Architecture and baseline run
The extension task used the same minGPT architecture and training pipeline as the replication, with two differences. First, the vocabulary expanded from (, ) to (, ), {, }. Second, the FALSE-class generator was “i.i.d. uniform over the four brackets, rejected if the result is depth-balanced”, and the held-out OOD set was again the set of strings where the candidate algorithms disagree: strings that are depth-balanced but have type-mismatches, like ({)}.
I first did a quick sweep over the number of layers in the model and the weight decay to find a configuration with a bimodal OOD distribution analogous to the first experiment, training only 5 models at each cell:
The only configuration I found that didn’t collapse to Forgetful was \(L=1, H=4, E=64, \mathrm{wd}=0.02, \mathrm{bs}=32\). The resulting OOD accuracy distribution was unimodal at OOD ≈ 0.6–0.8.
I didn’t realize that the distribution was unimodal at first because I didn’t generate the above figure until later. I thought it was bimodal because of the outliers on the far left of the figure, so I proceeded with the rest of the method from the replication.
Later, as a sanity check, I trained 3 models with the complete data set of strings of parentheses and braces, and their OOD accuracies were 0.918, 0.926, and 0.926, more than 0.1 higher than that of any model trained on the limited dataset. So it’s not that the 1-layer architecture is incapable of learning 2-Nested.
Step 2: The same SGLD hyperparameters were fine
I re-used the replication SGLD hyperparameters \((\varepsilon, \gamma, n\beta) = (3 \times 10^{-6},\, 500,\, 100)\) on the new checkpoints. A single-chain diagnostic and small \((\varepsilon, \gamma)\) grid sweep on the new architecture confirmed the trace was well-behaved at the same values.
Step 3: Interpreting the strings with extreme susceptibility gaps (they were the very long and very short ones)
I computed \(\Delta\chi := \chi_{\text{2D}} - \chi_{\text{1D}}\) on the 100-seed (L=1, wd=0.02) population, with \(\chi_{\text{2D}}\) averaging over the top 3 seeds and \(\chi_{\text{1D}}\) averaging over the bottom 3 seeds by OOD accuracy. I used the same 1024-sample eval set protocol as in the replication.
The clearest interpretable pattern was that string length correlates strongly with \(\Delta\chi\):
The statistical significance of this length pattern is non-trivial. Under the baseline dataset’s \(n = 2 + 2 \cdot \text{Bin}(19, 0.5)\) length distribution, the modal lengths 20 and 22 each get ≈ 11 % of the mass; the rare-tail lengths 10 and 32 each get ≈ 0.74 %.
Step 4: Retraining
I retrained on the two susceptibility-derived slices with n_copies = 10 (an order of magnitude above the replication setting), to compensate for the smaller per-sample \(\Delta\chi\) magnitudes seen here than in the replication setting: short-FALSE samples (predicted: push toward Forgetful, lowering OOD) and long-TRUE samples (predicted: push toward 2-Nested, raising OOD).
The OOD accuracy shifts appeared in the predicted directions but with very small magnitudes. In particular, the magnitude is in the same ballpark as the natural seed-to-seed variability, so in my book this experiment does not provide a huge amount of evidence.
AI use statement
Since the main purpose of this post is to serve as an application material for fellowships I’m currently applying to, it behooves me to be clear about which parts of this work were done by Claude. Things I did: Came up with the extension problem, specified the code, directed the debugging, wrote 50% of the body text of this writeup, rephrased many a Claude-ism in the other 50%. Some debugging ideas that were mine were: noticing Claude wrote a training loop that trained on the test set, and figuring out that the batch size was supposed to be 8. It was also my idea to try including 10× as many synthetic examples in the retraining step of the extension experiment, to try to account for the lower susceptibility gaps in that setting.
Things Claude did: Wrote all the code, made all the figures, wrote the first draft of 50% of this writeup, handled formatting of this page, answered a lot of my questions about the math.
Producing this writeup also taught me a lot about the limitations of working with Claude. When I was first sending baseline training runs off to validate the extension experiment, I was having trouble finding hyperparameters that yielded a bimodal distribution of OOD accuracies. Eventually, Claude just claimed that one of the runs was bimodal because it had 2 outliers. I was fooled by this for long enough to get through the rest of the steps for the extension experiment (otherwise I would have stopped there and tried something else). I only realized my mistake when I plotted the histogram of OOD accuracies for this post. I’m glad to have caught this error in this relatively low-stakes setting! I feel more prepared to be wary of the outputs of Claude going forward.