On the creation of narrow AI: hierarchy and nonlocality of neural network skills
Authors: Eric J Michaud
Asher Parker-Sartori
Max Tegmark
Venue: NeurIPS 2025 poster ()
Keywords: safety
mechanistic interpretability
theory
curriculum learning
pruning
connectionism
unlearning
Primary area: other
TL;DR: We investigate some basic questions about how neural networks learn and represent skills that are relevant to the problem of creating narrow AI systems.

Abstract: We study the problem of creating strong, yet narrow, AI systems. While recent AI progress has been driven by the training of large general-purpose foundation models, the creation of smaller models specialized for narrow domains could be valuable for both efficiency and safety. In this work, we explore two challenges involved in creating such systems, having to do with basic properties of how neural networks learn and structure their representations. The first challenge regards when it is possible to train narrow models from scratch. Through experiments on a synthetic task, we find that it is sometimes necessary to train networks on a wide distribution of data to learn certain narrow skills within that distribution. This effect arises when skills depend on each other hierarchically, and training on a broad distribution introduces a curriculum which substantially accelerates learning. The second challenge regards how to transfer particular skills from large general models into small specialized models. We find that model skills are often not perfectly localized to a particular set of prunable components. However, we find that methods based on pruning can still outperform distillation. We investigate the use of a regularization objective to align desired skills with prunable components while unlearning unnecessary skills.

On the creation of narrow AI: hierarchy and nonlocality of neural network skills
Eric Michaud 1.3- Asher Parker-Sartori Max Tegmark 1.3
1 Department of Physics, Massachusetts Institute of Technology
2 Department of EECS, Massachusetts Institute of Technology
The NSF AI Institute for Artificial Intelligence and Fundamental Interactions
Abstract
We study the problem of creating strong, yet narrow, Al systems. While recent AI progress has been driven by the training of large general-purpose foundation models, the creation of smaller models specialized for narrow domains could be valuable for both efficiency and safety. In this work, we explore two challenges involved in creating narrow AI systems, having do with basic properties of how neural networks learn and structure their representations. The first challenge regards when it is possible to train narrow models from scratch. Through experiments on a synthetic task, we find that it is sometimes necessary to train networks on a wide distribution of data to learn certain narrow skills within that distribution. This effect arises when skills depend on each other hierarchically, and training on broad distribution introduces a curriculum which substantially accelerates learning. The second challenge regards how to transfer particular skills from large general models into small specialized models. We find that model skills are often not perfectly localized to particular set of prunable components. However, we find that methods based on pruning can still outperform distillation. We investigate the use of regularization objective align desired skills with prunable components while unlearning unnecessary skills.
Introduction
Today, the most competent AI systems in any particular domain are general systems that are relatively competent in every domain. The best models at math and coding are also broadly knowledgeable about very diverse array of topics, from Roman history to home cooking recipes medical diagnostics. And when comain-speccif models are created today, they are typically general foundation models fine-tuned particular task, rather than new models trained from scratch [1, 2). though with some notable exceptions [3]. This state of affairs is convenient and powerful, since a single general model be used for variety of applications [4].
However, generality has downsides for both efficiency and safety. For instance, Al systems used as coding assistants possess large amount of knowledge which is never needed in those applications. Instead, we might like to use smaller, specialized networks which preserve the coding knowlede ledgledge of general systems without the same breadth of irrelevant knowledge. Narrow systems may also pose fewer safety risks than general systems. For instance, narrow systems may have fewer dangerous capabilities that pose CBRN risks [5]. be easier understand mechanistically [6, 7], be casier to verify properties of [8-.0]. More imaginatively. for systems to operate autonomously in the world requires large breadth of skills, and ecosystem of narrow "tool AI" systems may therefore reduce loss-of-control risks and better support human agency over the long term [11].
In this work, we investigate some basic questions about how neural networks learn and implement skills that are relevant to the problem of creating narrow Al systems. As summarized in Figure 1,
"ericjm@ mit.edu
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

(A)
(B)
coding
math
Distributed Representations
Less-distributed Representations
mathi odlencee
Neurots
Neuson
science
other
Figure 1: We study two challenges to making strong, narrow- purposse AI models. (A): Data may have hierarchical structure. If skills have a hierarchical dependence, where some skills are only learnable after more primitive skills are learned first, then it sometimes may be necessary to train on broad distribution of data to learn certain narrow skills within that distribution. These dynamics may mean that general -purpooe models must be trained to achieve good performance on some domains. (B): Model features are distributed. By default, skills may not be localizable to a particular set of model components (c.g. neurons). In this case, pruning of model components won't precisely retain wanted skills and remove unwanted skills from models. We explore methods for aligning the model features relevant to particular domains with smaller subset of model components while unlearning others.
focus main themes. First, consider the question of when it is possible train well- performing neural networks from scratch on a narrow data distribution. Through experiments on a novel synthetic task with hierarchical structure, find that it can be necessary to train networks a broad distribution to efficiently learn narrow tasks within that distribution. These results contribute growing literature on how task structure influences neural network learning dynamics [12-15], but with special relevance to the problem of creating narrow AI. Second, we consider the question of whether one can use pruning to turn broad networks into smaller narrow ones. We find that the nonlocality of network representations to prunable model components poses challenge for this goal. While distributed representations have been extensively studied the context neural network interpretability [16] and in the classical connectionist literature [17, 18], we study how this property of neural network computation impacts pruning and unlearning [5] for creating narrow AI. Our specific contributions are as follows:
We describe synthetic task, compositional multitask sparse parity (CMSP), extending the multitask sparse parity task of Michaud et al. [19]. We find that networks trained on this task exhibit extremely strong curriculum learning effects, where it is necessary to train broad distribution of tasks in order to learn certain other tasks.
We study pruning and unlearning on our networks trained on CMSP. We observe that tasks are often distributed and distinct subtasks are entangled, making pruning an imperfect strategy for "narrowing" the breadth of network skills. However, we find that a simple group-sparsity regularization objective can be used sparsify networks and unlearn skills.
We perform an empirical comparison of methods for creating narrow systems on MNIST and in LLMs. We tentatively find that methods based pruning outperform distillation and training networks from scratch for the creation of smaller, more narrow systems.?
This work is organized as follows. In Section 2, we briefly describe methods. In Section we perform a detailed (Section 3.1), study network training dynamics on it (Section 3.2), find distributed representations in these networks (Section 3.3), and use a regularization method for pruming and unlearning these networks (Section 3.4). We then compare various methods for creating narrow models, studying MNIST in Section and language models in Section 5. We discuss related work in Section 6 and conclude in Section 7.
study of training and pruning the CMSP task. We define the task
"Code for the experiments in this paper be found at: https://githum com/ ttt/ cnichhad/narror

---

2 Methods
Pruning: We aim to preserve the performance of model f(;;0) 0) on distribution Dx while pruning model components. Let be a collection of parameter indices corresponding to prunable components of the model (e.g. the in-weights and out-weights of neuron), and let G denote the collection of all such groups. After ablating a specific g. we denote the new parameters 0- To perform pruning, for each g G. we compute an ablation score 8, (x.)) [=(f 0), (xf 03), y)]l. the absolute change in the model's expected loss L after pruning g. We sort groups by their ablation score and prune greedily the desired sparsity from lowest to highest ablation score. Where feasible, we estimate 8y empirically by manually ablating the group across many samples from DN. When this is computationally intractable, we instead use linear estimate which refer to as attribution at
score after [20, 21]: 8g- Eieg
BL (-0:) = 28 (0g ----),, 0)} where aL/30 is the model's gradient 88,
on the distribution Dx. which can be computed once and reused for all E G.
Regularization: We also experiment with making networks more prunable by performing additional training with "group lasso" regularization penalty on the model weights [22-25]. The group lasso
penalty R is the LI norm of the L2 norms of cach parameter group: R(O) = EyEG v EiEn 02. This
penalty incentivizes the weights to become sparse at the level of entire groups g. When we perform "group lasso training" on a distribution Dx, we minimize the loss D (x;; 0), y) + )-OO)- Distillation: We also explore distilling knowledge from teacher model using the standard algorithm employed in [26], minimizing the KL divergence between student and teacher output distributions.
3 Case study: compositional multitask sparse parity
In this section, we conduct a detailed study of both curriculum learning and pruning on simple synthetic task, which we call compositional multitask sparse parity (CMSP).
3.1 Defining compositional multitask sparse parity (CMSP)
The compositional multitask sparse parity (CMSP) task is a simple extension of the spurse parity task recently studied in [27] and the multitask sparse parity task studied in [19], described below.
Barak et al. [27] studied neural network training dynamics on the sparse parity (SP) task. This is a binary classification problem on binary strings of length n, where the label of a given sample is the parity of the bits at subset / of k indices; y = o (=__)=[4] [=[]]. Strings I are sampled uniformly. Barak et al. observed that the loss curve for neural networks trained on this task exhibits sharp drop after an initial plateau, a case of the sort of "emergence" which has been observed in LLMs [28, 29] and in other toy settings [30].
Michaud et al. [19] extended the task to multitask sparse parity (MSP). In the MSP problem, input strings consist of m + bits, where the first m bits are called "control bits" and the last bits are called "task bits". These leading m control bits encode which "subtask" must be solved in each problem instance. Instead of a single length-k set of indices I, collection of m sets of task bit indices (typically of equal length k) is chosen {L1..... }. For each sample, only one control bit is ON (1) while the others are OFF (0). If control bit is ON, then that sample's 's label Michaud et al. [19] found that when subtasks are
is the parity of bits 4; =
power law distributed in frequency, where the probability that control bit is ON is pa i (a(+1)) the mean loss exhibits power-law scaling [31, 32] while individual sub- tasks are learned at different times proportional to their frequency [33]. They conjecture that LLM learning dynamics may be similar.
CMSP Samples: m=3, nw9, km3
x
100101001000
One of the main limitations of the MSP task, and the associated model of neural scaling from [19], is that subtasks are independent. However, the world, and the problem of learning from it, intuitively has a hierar- chical structure- in order to learn to do certain concepts and tasks, we must first learn simpler concepts and tasks. Accordingly, humans are taught curriculum, where simpler concepts and tasks are learned
Atomic
100010100011 1
010110111110
001101101110 0
Compositional
1100001010111
1101 110111001010 0
01 0111100011100 1 111000110100
first, and then composed into more complex ones later. To capture this structure, we now introduce the toy task compositional multitask sparse parity (CMSP). We show some CMSP samples to the right:
Control Bits
Taak Bits
3

---

CMSP is similar to MSP, except that (I) we require subtask indices to be disjoint: 1; N = 0 if i # j. and (2) multiple control bits can now be ON the same time, in which case the label the parity of the bits in the union of the indices for each subtask. If control bits i, are both ON, the label the parity of the task bits at indices T; U J.. We call samples for which only one control bit is ON "atomic" and samples for which multiple control bits are ON "composite" So if k = 3 for all subsets 4, then on atomic samples networks must compute the parity of input bits, on composite samples with two control bits ON the label is the parity of 6 input bits, and so on. Above, we illustrated some CMSP data samples with m = 3, n = and k = 3.
3.2 Learning dynamics on CMSP
We find that neural network training on CMSP exhibits extremely strong curriculum learning effects. To denote different types of CMSP samples, we list the ON control bits for those samples. So, given a choice of m, n,, k and /,,... Im Dio denotes all samples for atomic subtask O, and 110,1.2.3) denotes all samples from the composite subtask where the first four control bits are ON. For each subtask, there are 2 possible samples in that subtask, since the m. control bits are fixed but the n task bits are free.
We first train ReLU MLPs with 1-2 hidden layers of width 128 with the Adam optimizer on CMSP samples with m 4. n. 64, k = 4, and 2000 samples per task (atomic or composite) per batch. We show the learning dynamics of these networks in Figure 2. We find that when we train on a dataset containing both atomic and composite samples D(o) U Diy U (22) ((3 U (0.1.233) in equal proportion, atomic tasks are learned before composite ones.
Something interesting happens when we remove atomic samples from the dataset: we find that learning composite tasks takes dramatically longer. When we train on Doo UD U ((2) UD33) U 1,233 with 10000 samples per batch (2000 per subtask), across 40 seeds, we find that 27/40 networks converge composite subtask D (00,1.2.3) within 2 x 10 samples, and the networks that
Training on atomic and composite tasks, composite task learned after atomic tasks
1.00
Subtask
Loss (bits)
0.75
(0)
(1)
0.50
(2)
(3)
0.25
(0,1.2.3)
0.00
107
10a
Training samples
Training on atomic and composite tasks
Training on composite task only
Loss on (0,1,2,3)
1.00
Loss on (0,1,2,3)
1.00
0.75
0.75
0.50
0.50
0.25
0.25
(0,1,2,3)
(0.1,2.3)
0.00
0.00
10a
10"
10c
10'
10"
Training samples
Training samples
Figure 2: Training dynamics compositional multitask sparse parity. Top: training dynamics for
single network trained on four atomic subtasks (0), [1), (2). (3). and their composition [0,1,2.3), Bottom left: loss on compositional subtask (0.1.2,3). when training on atomic subtasks and their composition, as in the top subplot, across 40 network seeds (depth 3, width 128). Bottom right: loss on compositional subtask (0,1,2.3), when only training on samples from that composite task, without also training on atomic tasks, across 40 network seeds (depth 3, width 128). We find that removing the atomic tasks prevents our networks from learning the composite task. We report the minimum loss within the previous 100 steps of training to filter out loss spikes.

---

do converge typically converge within 2 x 10 samples (bottom left). However, when we train just on 1..33, with a batch size of 2000 samples, we find that 0/40 networks converge within 2 x 10 samples in Figure 2 (bottom right). It is much more efficient to train on broader distribution in order to learn subtask D00.1.2.3) than just training on D (00,1,2.3) on its own.
This result makes sense given the exponential hardness of learning parities [34]. Since the atomic task bits are disjoint in CMSP, networks can compute the parity of the composite tasks by computing the parity of features learned for the atomic tasks, reducing the complexity of the problem. For instance, one could imagine that 3 networks would compute the parity of the atomic tasks in the first layer, then the composite tasks in the next layer. Curiously though, we find that single-hidden-le (depth-2) 2) networks can still take advantage of a curriculum in Appendix A. We leave mechanistic analysis of how networks learn these tasks to future work.
These sorts of dynamics may be part of the explanation for why large-scale general-purpose models perform so strongly at many narrow, valuable tasks. We emphasize, however, that the toy task where we observe these currieulum effects is fairly contrived. is not clear what extent there are similarly strong effects on real-world tasks, and in many domains it is in fact possible to train narrow specialized models by only training on data in that domain. For instance, self-driving cars do not need to be trained on a broad corpus of text like LLMs.
Our work shows that for some types of tasks, it may be necessary to train broad data distribution in order to learn some subtasks efficiently. Accordingly, we now turn to the question of how to efficiently transfer knowledge from large models into smaller specialized ones.
3.3 Nonlocal representations in CMSP networks
If the circuitry for some subtasks were localized to a particular set of neurons, and the circuitry for other subtasks were localized to different neurons, then the task of specializing broad networks into narrow ones would be trivial. One could simply prune away neurons (or other model components, e.g. attention heads in transformers) associated with some subtasks, and keep others. However, often the situation seems more complicated than this ideal, which will be focus for the rest of this paper.
related problem has recently been studied in neural network interpretability, where it has been observed that individual computational units in neural networks, such as neurons, are polysemantic, activating across a wide variety of unrelated inputs [35, 36). Accordingly, many assume that the true model "features" do not align with architectural components like neurons. Multiple explanations have been proposed for this phenomenon. One is simply that the model architecture does not always "privilege" a particular basis [37, 38], though other incidental reasons for polysemanticity have also been proposed [39]. Another explanation of polysemanticity is the superposition hypothesis [40, 36). The superposition hypothesis suggests that the need to represent more features than there are dimensions or neurons prevents features from being represented as orthogonal directions in the feature space, and therefore all features cannot be aligned with standalone model dimensions neurons. Recently, studies that use sparse autoencoders to identify monosemantic model features have found that most features are highly distributed across large number of dimensions of activation space [16].
In our CMSP networks, we observe a related problem. By default, without any explicit regularization, there is no incentive for the network to localize certain circuits cleanly into a particular set of neurons.
We train networks with two hidden layers on dataset of CMSP samples with m 6, n = 18, & with two different skill trees: (0), (1), (2), (0.1.2), and (3), (4), (5), (3,4.5) until convergence. In Figure 3, we visualize the connectivity of 2-hidden-layer MIP trained on this dataset, and see that the network is densely connected without obvious structure.
We attempt to prune this network to retain performance on subtask [0.1,2] while unlearning (3.4.5). With MLPs, our groups of parameters g are the in in-weights, bias, and out-weights for each hidden layer neuron, so that |G] is the total number of hidden neurons. We compute ablation scores on the distribution D = (0..1.2) (estimated on 2000 samples) and prune greedily, as described in Section 2. In Figure 3, we show accuracy on subtask [0,1,2) vs. sparsity with this pruning strategy. When applied naively, network performance tends to degrade quickly, since without explicit regularization the network is not optimized to be naively prunable (see in Appendix Figure 8 and Figure 9 for more on how ablation scores for each subtask vary across neurons). However, when we perform an additional 1000 steps of training on D 1(0,1,2) (5000 samples per batch) after pruning (after either removing neurons from the architecture pinning pruned weights at zero) we recover
5

---

"Pretrained", unregularized network
Network after group-lasso on (0,1,2}
Training on subtask (0,1,) with
group lassa regularizat lon on the weights
Pruning on {0,1,2) ablation scores
Pruning on {0,1,2) ablation scores
1.0
1.0
0.9
0.9
Accuracy
0.8
Accuracy
0.8
0.7
0.7
0.6
0.6
0.5
0.5
0.4
0.4
0.1
0.8
0.9
0.95
0.98
0.1
0.8
0.9
0.95
0.98
Sparsity
Sparsity
Subtask {0,1,2), recovery training
Subtask (0,1,2), 1K recovery steps
Subtask (3,4.5), no recovery training
Subtask (3,4,5), 1 recovery steps
Figure 3: Top: We visualize the connectivity of h-hhiden--lyerr MLPs trained on CMSP distribution with subtasks (0). (1), (2). (0,1,2), (3), (4), (5), (3,4,5), visualized before (left) and after (right) regularizing network weights with the group lasso sparsity penalty while training on subtask (0,1,2). We find that network connectivity becomes sparse after regularizing. Negative weights are shown in blue and positive ones in red, with width proportional to the norm of the weight. Bottom: we show how pruning affects task performance subtasks (0,1,2) and [3,4,5] at varying sparsity levels. We prune neurons based on the absolute change that ablating them has on the loss on subtask [0,1,2]. We find that subtasks here nonlocal and entangled in the "pretrained" network (left). As prune neurons according to their relevance on subtask (0,1,2), at the sparsity at which performance subtask (0,1,2) accuracy drops below 98% (green line),
still recover on subtask (3,4,5) with a small amount of additional training. Thus naive pruning here has not completely and robustly unlearned subtask (3,4.5). However, after regularizing the weights (right), we find that not only we can more aggressively prune the network, but we have also robustly unlearned subtask (3,4,5). Note that the degree to which subtasks nonlocal and entangled in the pretrained networks depends on seed and width, and we show a variety of additional curves in Figure 10.
performance
performance higher sparsity levels. Since we training just on composite samples, 1000 steps is not enough to re-learn this task on its own, so if we can recover performance, that will be because the mechanisms for the task somewhat preserved after pruning. We observe that often, the tasks (0.1,2) and [3,4.5) are entangled when we prune as aggressively as we can while being able to recover performance on subtask [0,1.2), we can often still recover some performance on subtask (3,4,5). However, for our CMSP networks the degree of entanglement is highly seed-dependent, and sometimes the subtasks are disentangled enough that pruning as aggressively as possible on one subtask does in fact robustly unlearn the other. We show pruning curves across seeds and widths in Appendix Figure 10.
3.4 Regularizing "narrow" networks
We find that we can use simple regularization penalty simultaneously unlearn some tasks while incentivizing the network to move its features to less distributed, allowing for more aggressive pruning. As described in Section 2, simply perform additional training on 1,0,1. with group

---

Total Hidden Neurons vs Datapoints on MNIST Even Digits
Individual Group Lasso Runs
100
50k
98
480
Datapoints
20k
57.3
07-3
10.3
27.3
Accuracy (%)
96
120
Final Neuron Count
10k
97.
92
Training from Scratch
57.397.5
Distillation
90
Group Lasso Pruning
57.4
88
Attr. Prune then Recover
2k
Accuracy
Points are speced by -22 dalapoints
15
30
60
120
240
480
960
1920
1920 960 480 240 120 60 30
15
Total Hidden Neurons
Neurons Remaining
Figure 4: Left: We compare the performance of distillation, training from scratch, and two pruning approaches for creating small networks that classify MNIST even digits. Pruning- approaches Pareto-dominate distillation, achieving high compression ratios with fewer datapoints. All points are averaged over individual training runs. Right: When pruning using group lasso, it often helps to first prune rapidly, degrading performance, and then recover performance with no regularization. Each line represents single training run, with a new point logged every 2,000 datapoints. Darker lines correspond to more aggressive pruning runs that attain lower final neuron counts.
lasso sparsity penalty on the network weights. We aim for this penalty to "clean up" circuitry [41] not relevant to prediction on D ((0.1.2) while also sparsifying the weights across groups of parameters, allowing for easier pruning.
We apply this regularization to CMSP networks, training with penalty strength X = 10-3 with a batch size of 2000 for 10000 steps. With this regularized network, we then re-compute neuron ablation scores and prune. In Figure we find that, in CMSP networks, this method is effective unlearning skills and making the network more prunable. We find that we can prune more aggressively while retaining performance on subtask (0,1.2), and are not able recover performance on {3,4,5] at any sparsity level. These results complete our case study on CMSP: in order to efficiently learn compositional tasks, we train on broad distribution of tasks. Then, we can regularize the network to turn this broad network into a much smaller one which nevertheless retains the broad network's 's performance on the target compositional task.
Pruning vs distillation: MNIST
We now consider the problem of creating narrow systems in more natural domains, first on MNIST. As our narrow subtask, we choose only even digits from the original MNIST datasct [42]. We compare the resources required to achieve good performance on this narrow task when (1) training from scratch on the narrow task, and (2) distilling models from a general teacher on this task, (3) using group-lasse regularization prune large general model and (4) using attribution- based pruning and then recovery training on a large general model. Our teacher model is a ReLU MLP with two hidden layers each of width 1200, in [26], and achieves 98.7% accuracy the test set. When pruning, we use this same teacher model as our initial model and prune its hidden neurons to create a smaller network.
When using distillation (2), we use the approach of Hinton et al. [26] with T = 20. When pruning with the group lasso penalty (3), we use À values ranging from 0.001 to 0.008. Unlike in Section 3.4, we regularize and prune simultancously, pruning neurons when their L2 norm drops below 0.05. When the number of remaining neurons drops below a target threshold. we remove the pruning penalty and continue training to recover lost performance during pruning. When we prume up front and then separately recover performance (4),
attribution scores as described in Section 2.
To compare methods, we require that each method reach a test-set accuracy 97 percent, and we then plot the frontier of neuron count versus datapoints subject to that threshold. As seen in Figure 4 (right), it is often optimal to first aggressively prune the network down to the desired size and further train it until it reaches the requisite accuracy.
7

---

Pruning curves: regularized vs. unregularized
Recovery training after pruning
3i30 midd
ploup lase Amou focassss
group lanas A-2 stops
group teme
proup
group 1e0
praned
Loss (nats)
Sucdd pelfomargg
Loss (mat))
3
2
0.0
0.4
101
102
10a
10*
Sparsity
Recovery steps
Figure 5: Left: Neuron sparsity vs. loss curve for LLMs finetuned with group lasso regularization with varying À for 70k steps, as compared to the base network. Regularization flattens the sparsity vs. loss curve, at the cost of slightly degrading model performance. Right: After pruning our networks to 30%, 63%, and 80% sparsity for runs with À of 5e-4, 3e-4, and 1e-3, respectively, performance with additional training. We find that we can recover performance lost during pruning, including in the network that pruned without first using group-lasso regularization training. The plot shows an exponential moving average of batch losses, with individual batch losses shaded in back.
recover
We find that while group lasso pruning is highly sensitive to choices in hyperparameters, both pruning methods Pareto-dominate distillation and training from scratch, especially high neuron counts. Moreover, while other methods cannot consistently bridge the 97 percent accuracy threshold with fewer than 25 hidden neurons, aggressive pruning can consistently shrink the network's size to lower absolute limit.
5 Pruning distillation: LLMs Python documents
next study LLMs. As our narrow task Dv we choose next-token prediction on Python documents in the GitHub Code Dataset,
We first prune the neurons in the MLP blocks of Llama-3.2-1B [43]. Later, we will also consider pruning residual stream dimensions, which involves pruning all model parameters that "read to" or "write from" dimension of the residual stream [37]. We use attribution scores when pruning, and show that the attribution scores correlate moderately with true ablation scores in Appendix Figure 11. In Figure 5 (left), we show neuron sparsity vs loss curves. When pruning neurons naively, we see that loss increases quickly at low sparsity levels. We also experiment with applying group lasso regularization while further training on Python documents (learning rate of 2e-6, max length 512, 18 documents per batch) and find that this training does indeed level out the sparsity loss curve, albeit slight cost to the loss. Fortunately, we find that we can recover lost performance after pruning by doing small amount of additional training Python documents in Figure 5 (right). We find that despite the loss increasing substantially after naively pruning, we can also quickly recover that lost performance, and overall this strategy seems to be more efficient than using group lasso training. We therefore next compare naive pruning + recovery training against distillation and training networks from scratch.
We train networks with the Llama [43] architecture of varying shape and size. We use a learning rate of 5e-4, sequence length of 1024, and batch size of 64. For distillation, we use Llama-3. 1-- as teacher with 7 For pruning recovery, we prune Llama-3.2-11 to varying levels of neuron and residual stream sparsity, shown in Appendix Table 3. In Appendix Figure 12 we show learning curves for these three approaches. In Figure 6, we tentatively find that pruning substantially outperforms training from scratch and distilling model from scratch on the data-paramed frontier. Given

---

target narrow network size and a fixed data budget, if one already has access to a general model, it appears be more efficient to prune that model than it is to perform distillation.
We also experiment with pruning random model compo- nents instead of components selected by attribution score and find that random pruning performs similarly [44]! In Appendix Figure 13, we find that after moderate number of recovery steps, attribution pruning and random pruning result the same recovered performance. This result in
Frontier for loss=1.7 nats
5b
Scratcn hant baggts Scrtch w/ distibtionn Pruring Lama-J32-40
Training Tokens
2b
1b
line with our earlier discussion of nonlocality, and the em- pirical findings of Bricken al. [16], that monosemantic features are distributed widely across model components, While some studies have found some geometric similarity between functionally similar features [45, 46), in our case does not seem like the relevant features for our task are localized into a set of "Python" vs. "non-Python" neurons, at least that attribution pruning identifies.
500m
200m
100m
200m
500m
Parameters
Figure 6: Frontier of network parame- ters vs. data required to achieve certain cross-entropy loss Python documents. For the task of creating a LLM special- ized on Python documents, we find that pruning Llama-3.2. 2B-1B and then perform- ing recovery training is much more effi- cient than training LLMs from scratch or distilling LLMs from scratch on the soft targets of Llama-3.2-8B.
We lastly evaluate whether neuron pruning can robustly unlearn downstream tasks. We consider three pruning methods applied to Llama-3.2 at sparsity lev- els of 30%, 63%, and 80%: (i) random neuron pruning, (ii) attribution-based pruning on Python code, and (iii) models trained with group-lasso regularization on Python documents from Figure 5. We evaluate each pruned model on four benchmarks: counterfact [47], A12-ARC (ARC-
Unlcarning:
Easy) [48], WMDP-Bio, and WMDP-Cyber [49] for haz- ardous knowledge, We found that the base model per- formed poorly on WMDP-Chess
For each model, we first measure baseline accuracy after pruning, then fine-tune on each benchmark's 's training set with batch size 4, sweeping learning rates between le-6 and le-4 and reporting the best-performing configuration the test set. In Table 1, we find that all pruning techniques, including random pruning, robustly unlearn non- non-Python skills with roughly similar effectiveness, except that at low sparsity, attribution pruning does not achieve much unlearning. This result contrasts somewhat with our findings on CMSP networks in Sections 3.3 and 3.4, where training with the group lasso penalty led to stronger unlearning than naive pruning.
Table I: Unlearning: We prune Llama-3.2-ID using various techniques to create narrow Python- specific models, then evaluate other tasks to assess whether pruning led to unlearning. Accuracies are reported for the pruned model before finetuning and after finetuning on a train set of each benchmark (non-FT- FT). / indicates successful unlearning (FT score 2 10% lower than base).
Method
Spar.
counterfact
A12-ARC
WMDP Bio
WMDP Cyber
Base
0%
0.18 +0.49
0.65 +0,67
+0.62
0.35 +0.59
Random
30%
0.00 +0.50
0.24 +0.25
0.22 >0.27V
0.27
Attribution
30%
0.12-0.66
0.35
0.36-00.56
0.27-0.555
Group Lasso
30%
0.06++0.47
0.25-+
0.24-+0.31
0.26-0..27
Random
63%
0.00 -0.45
0.25 00.25/
0.23 0..27
Attribution
63%
0.02+0.52
0.25-+0300
0.23-40.34/
0.26-+.322
Group Lasso
63%
0.03++0.45
0.25-
0.24-+0.29
0.27-0.30V
Random
80%
0.00-
0.26 :0.25/
0.26->0.28/
0.23 00.26
Attribution
80%
0.00-+0.47
0.25-+.225
0.23-4.27V
0.26-0.28V
Group Lasso
80%
0.02
0.25-0.225
0.24-40229
0.26
6 Related Work
Distillation. Many works have built on the original distillation work of Hinton et al. [26], seeking to transfer intermediate representations [50, 51] and applying these techniques to language models [52, 53), often after pruning [54, 55). Relevant to our discussion, Ture et al. [56] showed that pretraining the student before distillation can substantially improve results.

---

Pruning, Even early approaches to pruning used second-order methods for pruning weights [57, 58], whereas our "attribution" scores are first-order. When training vision models, Zhou al. [25] and Wen et al. [24] used a group lasso penalty like we do, albeit while training on the full data distribution. Sanh et al. [59] proposed "movement pruning" to prune weights during transfer learning.
variety of works have applied structured pruning to LLMss [60, 61], including Xia et al. [62] who develop a method for task-specific pruning. Vural and Erdogdu [63] provide evidence that pruning boosts sample efficiency by directing gradient updates to a sparse set of important directions.
Task Structure and Learning Dynamics. Several works have lately investigated the relationship between task structure and learning dynamics [12-15]. Liu et al. [64] also briefly study task similar to CMSP to show that hierarchical relationships between tasks cause what they call "domino" learning dynamics. We also highlight the "Skill it!" paper of Chen et al. [65], who empirically investigate hierarchical learning dependencies for real L.M tasks.
Machine Unlearning. Many approaches to unlearning have been proposed [66-70]. Guo et al. [71] study how applying unlearming fine-tuning to different model components affects unlearning success,
inspired by a mechanistic understanding of how knowledge is retrieved in LLMs. Lee et al. [72] show that distillation can be applied after unlearning fine-tuning to more robustly remove unwanted capabilities. Highly relevant is the work of Cloud et al. [73], who apply "gradient routing" during training to localize capabilities to different model components, allowing pruning to be used for unlearning.
7 Discussion
This work explored some aspects of neural network learning which are relevant to the creation of narrow AI systems. Much work remains to be done in determining how relevant our findings are to real-world networks in various domains. We first emphasize that CMSP is a toy task, an existence proof that curriculum learning effects can be quite strong, but it is still unclear whether these dynamics occur on more natural datasets. For our pruning experiments, we emphasize that our greedy attribution pruning approach is quite simple, and we cannot rule out that more sophisticated pruning strategies would be more successful in preserving some skills while unlearming others [74]. Also, we did not optimally scale hyperparameters in our LLM experiments, so the performance of each of our models in Section 5 is likely imperfect. While we found that pruning a pre-existing Llama model was more efficient than training a model from scratch on Python documents, this could be due to our imperfect training setup, and not because one needs to train on a broad distribution of data to learn Python documents due to curriculum effects.
Limitations:
In this work, we have studied some potential challenges involved in creating narrow AI systems, having to do both with the structure of data and the structures learned internally by neural networks. Underlying this work is perspective from mechamistic interpretability, that neural networks compute a variety of sparse features [35], each with a distributed representation [36, 16, 46), and that these features are computed from each other hicrarchically in circuits [75-77]. First, we found that in order to learn certain complex features, we may have to first train on a broad set of samples which encourage the learning of simpler features. Second, because features are distributed across model components, is nontrivial problem to move set of task- specific features and circuits into a smaller network, especially while unlearning unneeded circuits.
Conclusion:
While a neural network's computation across the whole data distribution may be quite complex, we hope that the computation that networks perform on any particular task will be reducible to something less complex. That less complex computation, whatever it is, might be interpretable as a circuit [76, 77], or reduce to a simple program [78], or, as we studied here, could be instantiated in a much smaller network. However, this hope, and the question of whether it is possible to create narrow- purpose versions of today's 's models, probes at a basic question about the nature of the intelligence of these models. If their apparent generality indeed results from their having learned large, diverse set of crystallized, task-specific circuits, then we ought in principle to be able to create competent, specialized versions of these models by just transferring the relevant task-specific circuits. However, if the intelligence of our models, and intelligence more generally, is better understood as resulting from single unified algorithm, then the basic prospect of creating narrow Al systems that are as strong truly general ones could be challenge.
10

---

Acknowledgments and Disclosure of Funding
We thank Ziming Liu, Josh Engels, David D. Baek, Tom McGrath, and Jamie Simon for helpful conversations and feedback. E.J.M. was supported by the NSF via the Graduate Research Fellowship Program (Grant No. 2141064) and under Cooperative Agreement PHY-2019786 (IAIFI).
References
[1] ZZ Ren, Zhihong Shao, Junxiao Song, Huajian Xin, Haocheng Wang, Wanjia Zhao, Liyue Zhang, Zhe Fu, Qihao Zhu, Dejian Yang, al. Deepseck-proverr-2.. Advancing formal mathematical reasoning via reinforcement learning for subgoal decomposition. arXiv preprint arXiv:2504.21801, 2025.
[2] Paul Kassianik, Baturay Saglam, Alexander Chen, Blaine Nelson, Anu Vellore, Massimo Aufiero, Fraser Burch, Dhruv Kedia, Avi Zohary, Sajana Weerawardhena, al. Llama-3.1- foundationai satiiiooiiisccuurutlllmb---sssss
technical report. arXiv preprint arXiv:2504.21034 2025.
[3] Binyuan Hui, Jian Yang, Zeyu Cui, Jiaxi Yang, Dayiheng Liu, Lei Zhang, Tianyu Liu, Jia- jun Zhang, Bowen Yu, Keming Lu, et al. Qwen2. 5-coder technical report. arXiv preprint WxXiv:2409122886 2024.
[4] Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx. Michael Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportunities and risks of foundation models. arXiv preprint arXiv:2108 07258, 2021.
[5] Fazl Barez, Tingchen Fu, Ameya Prabhu, Stephen Casper, Amartya Sanyal, Adel Bibi, Aidan O' Gara, Robert Kirk, Ben Bucknall, Tim Fist, et al. Open problems in machine unlearning for safety. arXiv preprint arXiv:250I.0m 244.0995,. 2025.
[6] Leonard Bereska and Efstratios Gavves. Mechanistic interpretability for ai safety-a review. arXiv preprint arXiv:2404. 14000..4402. 2024.
[7] Lee Sharkey, Bilal Chughtai. Joshua Batson, Jack Lindsey, Jeff Wu, Lucius Bushnag, Nicholas Goldowsky-Dile Stefan Heimersheim, Alejandro Ortega, Joseph Bloom, al. Open problems in mechanistic interpretabilityy arXiv preprint arXiv:2501. 100::2501116466, 2025.
[8] Max Tegmark and Steve Omohundro. Provably safe systems: the only path to controllable agi. arXiv preprint arXiv:2309.0199: 2023.
[9] David Dalrymple, Joar Skalse. Yoshua Bengio, Stuart Russell, Max Tegmark, Sanjit Seshia, Steve Omohundro, Christian Szegedy, Ben Goldhaber, Nora Ammann, et al. Towards guar- anteed safe ai: A framework for ensuring robust and reliable ai systems. arXiv preprint 24X::22005.06624, 2024.
[10] Max Tegmark, Sören Mindermann, Vanessa Wilfred, and Wan Sie Lee. The singapore consensus on global ai safety research priorities. Conference report, 2025 Singapore Conference on AI: International Scientifie Exchange on Al Safety, Singapore, April 2025. URL https: //file. go. gov. sg/sg- consenssus-ai-safeetyy.... -ddf.
si-stty
[11] Eric Drexler,
Reframing superintelligence: Comprehensive ai services as general intelligence. Future of Humanity Institute Technical Report, 2019. Available at https: / //www fhi. ox. ac
[12] Rahul Ramesh, Ekdeep Singh Lubana, Mikail Khona, Robert Dick, and Hidenori Tanaka. Compositional capabilities of autoregressive transformers: A study on synthetic, interpretable tasks. arXiv preprint arXiv:2311 12997, 2023.
[13] Maya Okawa, Ekdeep S Lubana, Robert Dick, and Hidenori Tanaka. Compositional abilities emerge multiplicativel:: Exploring diffusion models on synthetic task. Advances Neural Information Processing Systems, 36:50173-50195 2023.

---

[14] Core Francisco Park, Maya Okawa, Andrew Lee, Ekdeep S Lubana, and Hidenori Tanaka. Emergence of hidden capabilities: Exploring learning dynamics in concept space. Advances in Neural Information Processing Systems, 37:84698-887777 84729, 2024.
[15] Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. Sgd learning on neural net- works: leap complexity and saddle-to-saddlt dynamics. In The Thirty Sixth Annual Conference on Learning Theory, pages 2552- 2623. PMLR, 2023.
[16] Trenton Bricken, Adly Templeton, Joshua Batson, Brian Chen, Adam Jermyn, Tom Con- erly, Nick Turner, Cem Anil, Carson Denison, Amanda Askell, Robert Lasenby, Yifan Wu, Shauna Kravec, Nicholas Schiefer, Tim Maxwell, Nicholas Joseph, 7ac Hatfield-Dodds, Alex Tamkin, Karina Nguyen, Brayden McLean, Josiah Burke, Tristan Hume, Shan Carter, Tom Henighan, and Christopher Olah. Towards monosemanticity: Decomposing language models with dictionary learning. Transformer Circuits Thread, 2023, https: //transformer--ee circuitt ppwp2023/mooosemaaaiaaaa wtweeeaanntii-ftueeteeeeeeeeeeee index.html, html.
[17] Paul Smolensky. Tensor product variable binding and the representation symbolic structures in connectionist systems. Artificial intelligence, 46(1-2):15-2116, 1990.
[18] Geoffrey E Hinton. Learning distributed representations of concepts. Proceedings of the Annual Meeting of the Cognitive Science Society, volume 8, 1986.
[19] Eric Michaud, Ziming Liu, Uzay Girit, and Max Tegmark. The quantization model of neural scaling. Advances in Neural Information Processing Systems, 36:28699-28772 28722, 2023,
[20] Neel Nanda. Attribution patching: Activation patching at industrial scale, 2023. URL https: patching.
/ www.
[21] Aaquib Syed, Can Rager, and Arthur Conmy. Attribution patching outperforms automated circuit discovery. arXiv preprint arXiv:2310. 10X14:2300.10348, 2023.
[22] Robert Tibshirani. Regression shrinkage and selection via the lasso. Journal of the Royal Statistical Society Series B: Statistical Methodology, 58(1):267-28) 1996.
[23] Ming Yuan and Yi Lin. Model selection and estimation in regression with grouped variables. Journal of the Royal Statistical Society Series B: Statistical Methodology, 68(1):49-00) -7,4496765400 2006.
[24] Wei Wen, Chunpeng Wu, Yandan Wang, Yiran Chen, and Hai Li. Learning structured sparsity in deep neural networks. Advances in neural information processing systems, 29, 2016.
[25] Hao Zhou, Jose M Alvarez, and Fatih Porikli. Less more: Towards compact enns. In Computer Vision-ECCV 2016: 14th European Conference, Amsterdam, The Netherlands, October 1 11-14, 2016, Proceedings, Part IV 14, pages 662-677. Springer, 2016.
[26] Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in a neural network. arXiv preprint arXiv:
2015.
[27] Boaz Barak, Benjamin Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang Hidden progress in deep learning: Sgd learns parities near the computational limit. Advances in 2022.
Neural Information Processing Systems, 35:21750-211
[28] Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, et al. Emergent abilities of large language models. arXiv preprint arXiv:2206.00 076682 2022.
[29]
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodae Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads. Transformer Circuits Thread, 2022. http://raaafommmeeeecciiisippppp homtp//iiii2iiioiotoooeeaannnnn and-iduction-n- heads/index.html.
12

---

[30] Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt, Progress measures for grokking via mechanistic interpretability. arXiv preprint arXiv:2301.05217, 2023.
[31] Joel Hestness, Sharan Narang, Newsha Ardalani, Gregory Diamos, Heewoo Jun, Hassan Kianinejad, Md Mostofa Ali Patwary, Yang Yang, and Yangi Zhou. Deep learning scaling is predictable, empirically, arXiv preprint arXiv: / /7 124099 2017.
[32] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361 2020.
[33] Nayara Fonseca, Scok Hyeong Lee, Chris Mingard, Ard Louis, et al. An exactly solvable model for emergence and scaling laws in the multitask sparse parity problem. Advances in Neural Information Processing Systems, 37:39632-39966 2024.
[34] Itamar Shoshani and Ohad Shamir. Hardness of learning fixed parities with neural networks. arXiv preprint arXiv:2501.008// 2025.
[35] Chris Olah, Alexander Mordvintsev, and Ludwig Schubert, Feature visualization. Distill, 2017. http://iittt1lllppb/22//fftttte
doi: 10.23915/distill.oo //diitill.00000
[36] Nelson Elhage, Tristan Hume, Catherine Olsson, Nicholas Schiefer, Tom Henighan, Shauna Kravec, Zac Hatfield-Dodds, Robert Lasenby, Dawn Drain, Carol Chen, Roger Grosse, Sam McCandlish, Jared Kaplan, Dario Amodei, Martin Wattenberg, and Christopher Olah. Toy mod- els of superposition. Transformer Circuits Thread, 2022. attps://ttet http:////aaaafooeerrrrreee circuttss pub/2022 2ty//t0222/too0e//e//eeee2o hetee///oommoee///eeeeeeeee www.mmdddemlin.. /ttp//in.x....//// html.
[37] Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Haafiied-Dodds Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. mathematical framework for transformer circuits. Transformer Circuits Thread, 2021. https: http::/traaaffrrnnn ormer- circuits, pwwppppp/ hamb/2221/framoo/// httt///rrawwrr//eeeee/eeaeeaeeee in///.nddxxx/// html.
[38] Nelson Elhage, Robert Lasenby, and Christopher Olah. Privileged bases in the transformer residual stream. Transformer Circuits Thread, page 24, 2023.
[39] Victor Lecomte. Kushal Thaman, Rylan Schaeffer, Naomi Bashkansky, Trevor Chow, and Sanmi Koyejo. What causes polysemanticity? an alternative origin story of mixed selectivity from incidental causes. arXiv preprint arXiv:2312. 03096. 2023.
[40] Sanjeev Arora, Yuanzhi Li, Yingyu Liang. Tengyu Ma and Andrej Risteski. Linear algebraic structure of word senses, with applications polysemy. Transactions of the Association for Computational Linguistics, 6:483-495 2018.
[41] Vikrant Varma, Rohin Shah, Zachary Kenton, János Kramár, and Ramana Kumar. Explaining grokking through circuit efficiency. arXiv preprint arXiv:2309.02390, 2023.
[42] Yann LeCun. The mnist database of handwritten digits, 1998. URL httpp//yann. lecun. com/exdb/ /nnis//nnis.......
[43] Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The llama 3 herd of models. arXiv preprint ar 2407:22407. 2024.
[44] Shuyao Xu, Liu Jiayao, Zhenfeng He, Cheng Peng, and Weidi Xu. The surprising effectiveness of randomness in Ilm pruning. In Sparsity in LLMs (SLLM): Deep Dive into Mixture of Experts, Quantization, Hardware, and Inference.
[45] Yuxiao Li, Eric J Michaud, David D Baek, Joshua Engels, Xiaoqing Sun, and Max Tegmark. The geometry of concepts: Sparse autoencoder feature structure. Entropy, 27(4)3344, 2025.
13

---

[46] Adly Templeton, Tom Conerly, Jonathan Marcus, Jack Lindsey, Trenton Bricken, Brian Chen, Adam Pearce Craig Citro, Emmanuel Ameisen, Andy Jones. Hoagy Cunningham. Nicholas I. Turuer Callum McDougall, Monte MacDiarmid, C. Daniel Freeman, Theodore R. Sumers, Edward Rees, Joshua Batson, Adam Jermyn, Shan Carter, Chris Olah, and Tom Henighan. Scaling monosemanticity: Extracting interpretable features from claude 3 sonnet. Transformer Circuits Thread, 2024. URL html.
circuits ptt.iitss.ub//222/ pub/
scaling- www.ooomntiiiittonnneeeoooo
[47] Kevin Meng. David Bau, Alex Andonian, and Yonatan Belinkov. Locating and editing factual associations in gpt. Advances in neural information processing systems, 35:17359- 17372, 2022.
[48] Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv: 883.05557, 2018.
[49] Nathaniel Li, Alexander Pan, Anjali Gopal, Summer Yue, Daniel Berrios, Alice Gatti, Justin D
Li, Ann-Kathrin Dombrowski, Shashwat Goel, Long Phan, et al. The wmdp benchmark: Measuring and reducing malicious use with unlearning. arXiv preprint mrXi::2403.05218. 2024.
[50] Adriana Romero, Nicolas Ballas, Samira Ebrahimi Kahou, Antoine Chassang, Carlo Gatta, and
Yoshua Bengio. Fitnets: Hints for thin deep nets. arXiv preprint arXiv: 2014.
[51] Siqi Sun, Yu Cheng, Zhe Gan, and Jingjing Liu. Patient knowledge distillation for bert model compression. arXiv preprint a9Xiv:/908.0000 00355, 2019.
[52) Xiaoqi Jiao, Yichun Yin, Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, and Qun Liu. Tinybert: Distilling bert for natural language understanding. arXiv preprint arXiv:1909.10351, 2019.
[53] V Sanh. Distilbert, distilled version of bert: smaller, faster, cheaper and lighter. arXiv preprint arXiv:1910.01108,, 2019.
[54] IS McCarley, Rishav Chakravarti, and Avirup Sil. Structured pruning of a bert-based question answering model. arXiv preprint arXiv: 1eI01999000..... 2019.
[55] AI Meta. Llama 3.2: Revolutionizing edge and vision with open, customizable models. Meta Al Blog. Retrieved March 2025, 2024.
[56] Iulia Turc, Ming-Wes Wei Chang, Kenton Lee, and Kristina Toutanova. Well-read students learn better: On the importance of pre-training compact models. arXiv preprint arXiv: (xxiv:/908. met::::0809090 2019.
[57] Yann LeCun, John S. Denker, and Sara A. Solla. Optimal brain damage. In David S. Touretzky, editor, Advances in Neural Information Processing Systems 2, pages 598-605. Morgan Kaufmann, 1990. URL httpp:::/t/toeeeeeeeeete 6G98220acC67091115551188777666 htw.:/12277766---Aaaaaeaee .html.
[58] Babak Hassibi and David G. Stork. Second order derivatives for network pruning: Optimal brain surgeon. In Stephen J. Hanson, Jack D. Cowan, and C. Lee Giles, editors, Advances in Neural Information Processing Systems 5, pages 164-171. Mor- gan Kaufmann, 1993. URL https: httpss:/prooeeedinggs www.eddiiiinnnnnnnneeeeeeeee nww.innnni... cc/ pptpp/ap//////aaaa// neettoo prum
[59] Victor Sanh, Thomas Wolf, and Alexander Rush. Movement pruning: Adaptive sparsity by fine-tuning. Advances in neural information processing systems, 33:20378-20338 2020.
[60] Mengzhou Xia, Tianyu Gao, Zhiyuan Zeng, and Danqi Chen. Sheared laama: Accelerating language model pre-training via structured pruning. arXiv preprint arXiv:2310.0669) 2023.
[61] Xinyin Ma, Gongfan Fang, and Xinchao Wang. language models. Advances in neural information processing systems, 36:21702-2111 227200 2023.
On the structural pruning of large
[62] Mengzhou Xia, Zexuan Zhong, and Danqi Chen. Structured pruning learns compact and accurate models. arXiv preprint arXiv:2204.004 2022.
14

---

[63] Nuri Mert Vural and Murat A. Erdogdu. Pruning is optimal for learning sparse features in nigh-dimensions, 2024. URL https: / //txxxxxvvv///////////a/////aav/ org/abs/2406 08658.
[64] Ziming Liu, Yizhou Liu, Eric Michaud, Jeff Gore, and Max Tegmark. Physics of skill learning. arXiv preprini arXiv:2501. 2025.
[65] Mayee Chen, Nicholas Roberts, Kush Bhatia, Jue Wang, Ce Zhang, Frederic Sala, and Christo- pher Ré. Skill-it! a data-driven skills framework for understanding and training language models, Advances in Neural Information Processing Systems, 36:36000-36040, 2023.
[66] Yinzhi Cao and Junfeng Yang. Towards making systems forget with machine unlearning. In 2015 IEEE symposium on security and privacy, pages 463-480, IEEE, 2015.
[67] Lucas Bourtoule, Varun Chandrasckaran, Christopher A. Choquette-Choo Hengrui Jia, Adelin Travers, Baiwu Zhang, David Lie, and Nicolas Papernot. Machine unlearning. arXiv preprint arXiv: 19/2.03///////////// 2019, doi: 11..550WrXii.1112.0353 URL httph:ttt//artt org/ abs/ 1912. 03817.
[68] Yuanshun Yao, Xiaojun Xu, and Yang Liu. Large language model unlearning. Advances in Neural Information Processing Systems, 37: 17:105425-105475 2024.
[69] Jiaao Chen and Diyi Yang. Unlearn what you want forget: Efficient unlearning for Ilms. arXiv preprini urXiv:2310.201555 2023.
[70] Kang Gu, Md Rafi Ur Rashid, Najrin Sultana, and Shagufta Mehnaz. Second-order infor- mation matters: Revisiting machine unlearning for large language models. arXiv preprint arXiv:2403. 10557, 2024.
[71] Phillip Guo, Aaquib Syed, Abhay Sheshadri, Aidan Ewart, and Gintare Karolina Dziugaite. Mechanistic unlearning: Robust knowledge unlearning and editing via mechanistic localization, 2024. URL https: //wxxv/./vvv///////////////aaa org/abs httg/bbb22410 /2400. 12949.
[72] Bruce W. Lee, Addie Foote, Alex Infanger, Leni Shor, Harish Kamath, Jacob Goldman- Wetzlerr Bryce Woodworth, Alex Cloud, and Alexander Matt Turner. Distillation robustifies unlearning, 2025. URL https: //tp:////xxvvv www.aas.////eeeeooooot....t..aee /206. 06277
[73] Alex Cloud, Jacob Goldman- -eezzlnn-eezzler, Evžen Wybitul, Joseph Miller, and Alexander Matt Turner Gradient routing: Masking gradients to localize computation in neural networks. arXiv preprint wrXiv:2410.04322 2024.
[74] Yanyu Li, Pu Zhao, Geng Yuan, Xue Lin, Yanzhi Wang. and Xin Chen. Pruning as- as-search: Efficient neural architecture search via channel pruning and structural reparameterization. arXiv preprint arXiv:2206.01) 198, 2022.
[75] Chris Olah, Nick Cammarata, Ludwig Schubert, Gabriel Goh, Michael Petrov, and Shan Carter. Zoom in: An introduction to circuits. Distill, 2020. doi: 14.3111//istill.00224..00 https://distill. hwtpaalll...uub220oooccuu//ooooo
[76] Samuel Marks, Can Rager, Eric Michaud, Yonatan Belinkov, David Bau, and Aaron Mueller. Sparse feature circuits: Discovering and editing interpretable causal graphs in language models. arXiv preprini arXiv:2403. 19647, 2024.
[77] Jack Lindsey, Wes Gurnee, Emmanuel Ameisen, Brian Chen, Adam Pearce, Nicholas L. Turner, Craig Citro, David Abrahams, Shan Carter, Basil Hosmer, Jonathan Marcus, Michael Sklar, Adly Templeton, Trenton Bricken, Callum McDougall, Hoagy Cunningham, Thomas Henighan, Adam Jermyn, Andy Jones, Andrew Persic, Zhenyi Qi, T. Ben Thompson, Sam Zimmerman, Kelley Rivoire, Thomas Conerly, Chris Olah, and Joshua Batson. On the biology of a large language model. Trunsformer Circuits Thread, 2025. URI. https: ormer //transformer-en
//transf
circuits wwt.riiisspppp/22..55555 http:u//..pp//////ttttrttttteet //t//atttttooonnooo ttraps/bbologe html.
[78] Eric J Michaud, Isaac Liao, Vedang Lad, Ziming Liu, Anish Mudide, Chloe Loughridge, Zifan Carl Guo, Tara Rezaci Kheirkhah, Mateja Vukelić, and Max Tegmark. Opening the ai black box: Distilling machine-learned algorithms into code. Entropy, 26(12):10 1046, 2024.
15

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We clearly state our contributions in our introduction, and these match the contents of the paper.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: We include a limitations section in Section well addressing possible limitations throughout the paper.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Answer: [NA]
16

---

Justification:
We do not include any theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: Yes, we describe our experimental setup in each section. We also have linked to repository containing the code used to produce all our results.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
17

---

Answer: [Yes]
Justification:
We have linked a github the paper.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips cc/ http:/gudddsssdsie/eeeee
wwwisssssssoolllsooeoooo olicc) for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: / //nips.. wtt.cc/ppuuc wtcp//ublccgguuiiceeeeeeeeeeceee
Sttpe//CoeeSSuiiiieeee wwwssssooollii onPolicy) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We include experimental details in the main text.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined other appropriate information about the statistical significance of the experiments?
Answer:
[No]
Justification: In Figure 2, Figure 7, and Figure 10 directly show results across all seeds instead of plotting the mean and variance. In Figure 4 we currently report means without error bars.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
18

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еπхо rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We include compute estimates for our experiments in Appendix C.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: http:://eeuuupppeeeee httpp////pppcccceiiceieeeeeeeeed
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: While we are motivated by some broader concerns around societal impact, the work done here relatively small-scale and exploratory. We discuss positive impacts in the Introduction.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
19

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our work does not pose such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We properly reference the datasets and models we use.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
20

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We do release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmoel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: Our work did not involve crowdsourcing nor work with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question:
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Our work did not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
21

---

Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification:
We did not LLMs in non-standard ways.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, non-standard components.
Please refer to our LLM policy (https fattps://nnt for what should or should not be described.
/ /tppp//////upps.../// /00ps/////ooeeeeeee/22/////////
22

---

A Additional results on CMSP
Here we include some additional results on CMSP. We first explore whether network shape influences learning dynamics when learning composite tasks (w/ corresponding atomic tasks in the training dataset) vs. when learning an atomic task on its own. We find that the effect depth is similar between composite vs. atomic tasks. This suggests that our networks do not necessarily need to be deep in order to take advantage of compositional structure this task. Note that a previous version of this manuscript suggested the opposite-We had observed that deeper networks learn composite CMSP tasks faster than shallow networks, but it turns out that deeper networks also learn standalone atomic tasks faster than shallow networks. We show loss curves of depth-3 width-128, depth-2 width-361, and depth-2 width- 128 networks in Figure
Compositional subtask loss dynamics
Loss on (0,1,2,3)
1.0
0.5
Depth 2, width 128
Depth 2, width 361
Depth 3, width
0.0
10?
10a
10
Training samples
Atomic subtask loss dynamics
1.0
Loss on (0}
Depth 2,
Depth 2, width
Depth 3, width 128
0.0
10º
10
10'
10a
Training samples
Figure 7: Impact of network shape CMSP learning. Top: Forty runs of training with different seeds on CMSP dataset Dip) L DD) U UD U (3) U 11.0..2..3 with m = 4, n 64, k 4. We find that 27/40 runs converge with two hidden layers versus 19/40 runs with single hidden layer with width 128 and 16/40 runs with a single hidden layer and width 361. Single-hidden-layer networks are therefore able to learn composite subtasks in CMSP roughly well as deeper networks. Bottom: Forty runs of training on standalone atomic task m 1, n = 64, 4.
We also show how ablation in Section 3.3, we show ablation scores for each neuron for each subtask Figure 8 and Figure 9, We find that there neurons which have high scores across most subtasks, even subtasks in different "skill trees" (10,1,2) versus (3,4.5)).
vary neurons in unregularized networks. With the setup
In Figure 10 we show sparsity vs accuracy curves from pruning like in Figure 3 (bottom left), across seeds and network sizes. In our unregularized networks, there a lot of variation between runs in how entangled the different subtasks are.
B Additional results on LLMs
B.1
Additional plots
We also show some additional results with LLMs. In Figure 11, show how attribution scores correlate with ablation scores in LLMss. We show results for pruning both neurons and residual stream dimensions.
23

---

102
102
abs change in loss from ablation
abs change loss from ablation
[1]
[2]
[3]
[4]
10a
10a
151
10.
[3. 51
100
100
10
10
0
20
40
60
20
40
60
hidden neuron
hidden neuron
Figure 8: For each hidden neuron in the first hidden layer (left) and the second hidden layer (right) in CMSP network, we show the ablation score of that neuron for each subtask. For some neurons, ablation scores are high across multiple subtasks, even from separate "skill trees""
Ablation scores across layers and subtasks
agg0lablltion score)
laye
aayer
3
Byer
4 layer
5 ayer
10
1.2 layer
layer
10
20
30
40
50
60
hidden neuron
Figure 9: Ablation scores for each subtask and each neuron. This is simply a different way of visualizing the results Figure 8.
In Figure 12, show learning curves across runs when training from scratch on Python documents, when performing distillation, and when performing recovery training after pruning both neurons and residual stream dimensions, varying target sparsities, of Llama-3.2-
In Figure 13, we show recovery curves after pruning neurons and residual stream dimensions from Llama-3.2-18 to varying sparsity levels. We find that pruning random components performs roughly as well as pruning components with the lowest attribution scores.
B.2 Additional experimental details
When training networks from scratch and distilling from scratch (Figure 12 and Figure we train transformers with the Llama 3 architecture [43] of varying size, listed in Table 2.
When prune Llama-3.2-1 combinations shown in Table 3.
prune neurons and residual stream dimensions with the sparsity
C
Compute estimates
We estimate the compute used for each of our experiments.

---

Pruning "pretrained" unregularized MLPs based subtask (0,1,2) ablation scores
Seed O
Seed 1
Seed 2
Seed
Seed 4
1.0
Width 64 Accuracy
0.75
0.5
Width 128
Accuracy
0.75
Width 256
Accuracy
0.75
Width 512
Accuracy
0.75
1.0
Width 1024
Accuracy
0.75
Sparsity
Sparsity
Sparsity
Sparsity
Sparsity
Subtask (0,122) no recovery training
Subtask (0,1,2), 1k recovery steps
Subtask (3,4,5), na recavery training
Subtask (3,4,5), 1k recovery steps
Figure 10: Pruning sparsity accuracy curves for pretrained CMSP networks. We train MLPs of varying width and random seed on CMSP dataset with two skill trees, as in Figure 3. We prune neurons based ablation score on subtask [0,1.2). On some networks, ablating neurons maximum sparsity while preserving performance on subtask (0, 1,2] robustly unlearns subtask (3,4.5). However, for other networks, the subtasks seem more entangled, and we can recover performance on subtask (3,4.5) even after ablating to the maximum extent we can still recover 98% accuracy on subtask (0,1,2) (green line).
For the CMSP training experiments shown in Figure 2 and Figure with 40 different seeds. We ran each job on a GPU, but on a cluster with variety of different node configurations and GPUs. Jobs generally took between 5-60 minutes complete, for between 13-160 total hours of GPU time.
4 configurations each
For the experiments shown in Figure 3 and Figure 10, each run, across choices of seed and 5 choices of width, trained networks on a CMSP task, ran pruning experiments, performed group-lasso training, and pruned again. Each such took between 5-120 minutes GPU configurations, for a total time of 2-50 hours of GPU time.
cluster with varying
For our MNIST experiments shown in Figure 4, we plot total of 54 points, each of which is averaged over 10 training runs. We estimate that each run took around minutes on our cluster, totaling to around 18 hours of GPU time.
25

---

Neurons
Residual Stream Dimensions
10
10
Pearson r = 0.44
Pearson r 0.39
Abs Ablation Score
Abs Ablation Score
10
10-3
10
10
10-4
10-4
10-7
10
10
10 10
10
10
10
10 10:2
Abs Attribution Score
Abs Attribution Score
Figure 11: Comparison of ablation vs. attribution scores for Llama-3.2-1B neurons (left) and residual stream dimensions (right), evaluated on a single batch of Python code documents. For each figure, we fully ablate model components and compare the absolute change in loss (Abs Ablation Score, simply called "ablation score" elsewhere in the paper) with the absolute value the linear estimate this change computed from model gradients (Abs Attribution Score, called "attribution score" elsewhere). Note that the correlation between the log of these scores is 0,80 for neurons and 0.05 for residual stream components.
Train from scratch
Distillation from scratch
Prune + recovery
6 10
51m
Cross-entropy loss (nats)
118m
235m
x 10
256m
71m
296m
x 10
173m
417m
339m
667m
478m
339m
2 10
598m
478m
2.3b
984m
1,5
1,50
3.2b
2.3b
10c
102
103
104
10a
104
10a
104
10a
Figure 12: Training curves of LLMs on Python documents when training from scratch (left), training from scratch but with distillation of a larger pretrained LLM (center), and when training pruned pretrained models (right) to recover performance lost during training.
For our LLM experiments, we ran our experiments on A100-80GB nodes, with a single GPU allocated per experiment. When we trained models from scratch on Python documents, we used 9 configurations with job lengths between 1-3 days. For distillation we had 9 configurations with job lengths between 2-3 days, though some jobs failed. For the group lasso training experiments in Figure 5, we show 3 configurations trained with job lengths of days, and we performed recovery training on these models with job length of day. We estimate that the total time for these jobs was less than 1600 hours of 1100 time, though the full set of experiments we attempted in the work that led this manuscript could over 5000 hours.
26

---

Cross-entropy loss (nats)
10
50% neuron, 20% residual sparsity
neuron, sparsity 90% 90% residual sparsity
2
Altribution pruning
Random pruning
10
102
103
10
Recovery steps
Figure 13: Recovery training curves after pruning neurons and residual stream dimensions, for Llama-3.2-IB on Python code documents. We compare recovery performance when pruning based on attribution scores vs. choosing components randomly. Attribution scores are computed across 1024 documents with max length of 1024 tokens. We find that pruning with attribution scores is better than pruning random components, however this gap is eventually recovered. For instance, at step 5365 when our run with 50% neuron sparsity and 20% residual sparsity first matches the performance of the original model (== 1.3 nats), the performance randomly-prun -pruned model is almost identical at 1.301 nats.
Table 2: Transformer model configurations explored
Hidden size
#Layers
#Heads
Intermediate size
256
4
4
1,024
512
8
2,048
768
12
12
3,072
864
16
16
3,456
1,024
16
16
4,096
1,280
20
20
5,120
1,536
24
24
6,144
1,792
28
28
7,168
2,048
32
32
8,192
Table 3: Neuron and residual sparsity configurations when using attribution pruning on Llama-3.2-1B.
Neuron sparsity
Residual sparsity
0.50
0.50
0.80
0.50
0.90
0.50
0.95
0.50
0.80
0.80
0.90
0.90
27


Reviews:
[{"review_id": "kZi6Ahj0eS", "reviewer": "Reviewer_kBj4", "summary": "This paper studies the challenges in creating _narrow_ AI: \"tool\" AI systems designed to have specific skills as opposed to general-purpose models with broad skills. This is motivated from the viewpoints of information efficiency as well as preservation of human agency. The study is conducted from two broad angles: (1) the efficacy of training a narrowly-skilled model from scratch, and (2) how best one could transfer specific skills (while discarding the rest) from a generalist model into a specialist model.\n\nRegarding (1), an empirical study using the compositional multitask sparse parity (CMSP) task is conducted. The task is designed such that certain difficult subtasks are \"composites\" of easier, \"atomic\" subtasks, inducing the possibility for \"curriculum learning\". It is observed that neural networks indeed learn the compositional task much faster when also trained on the atomic tasks, while training from scratch only on the compositional \"narrow\" subset is sample-inefficient (i.e., requires more samples in total to reach the same performance) -- this indicates that those easier skills cannot be safely skipped when training from scratch. They also find that depth helps learning to solve CMSP, leading to the hypothesis that many deep networks indeed learn hierarchical tasks in the same staggered fashion. \n\nRegarding (2), pruning and distillation are explored as candidate methods for selectively transferring skills from a broader model into a smaller model. This is analysed using the CMSP task above, MNIST digit classification, as well as LLM next-token prediction. They find that the circuitries for different tasks are entangled: this makes pruning for specific skills a non-trivial task. To address this problem, they propose a group lasso regularization that aims to \"clean up\" entire groups of neurons that are irrelevant to the narrow task at hand. The observations suggest that pruning the model whilst applying the proposed regularization, is more efficient in comparison to distillation.", "questions": "Do the authors have any ideas/plans for extending the experimental setup in §3 and §4 to more realistic datasets?", "limitations": "The authors mention limitations in §7, but should mention that CMSP is a toy task (which they do mention in lines 159-162 but should also briefly include in §7 for completeness).", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "### Strengths\n\n1. The questions explored are insightful, interesting and well-motivated. With general-purpose MLLMs taking over the world, this kind of work that explores specialist models instead, is well-timed.\n1. The proposed compositional multitask sparse parity task and the group lasso regularization for pruning can, in themselves, be valuable for future work. Notably, pruning with the sparsity regularization seems to be an effective strategy for narrow skill transfer. \n1. The paper is fairly well-organized given the density of information covered. Sections follow each other logically.\n\n### Weaknesses\n\n1. The curriculum learning claims are based entirely on the toy case study of compositional multitask sparse parity.  While these experiments themselves are quite thorough, they do not guarantee transfer to real-world cases: the authors note this in §3.2. Limitations noted in §7 also apply. \n1. The authors conduct LLM experiments in §5, but in contrast, the vision counterpart of these experiments in §4 is confined to MNIST. Given that vision is just as important as language in modern AI systems, I would be curious to see whether the results transfer to more realistic vision datasets -- absence of these results does impact the paper's significance. \n\nOverall, I believe the strengths sufficiently outweigh the weaknesses: enough for me to recommend borderline acceptance at this stage. I would have considered a higher score if more realistic-dataset experiments had been presented, but I would be happy to reconsider my score during the discussion phase and after looking at insights from other reviewers.", "was_revised": true, "final_justification": "My main concern was that the paper relies too much on toy experiments, and that empirical breadth is not sufficiently covered. In addition to this, I have realised after reading the other reviews that there remains a lack of a clear \"message\" and that the paper is at times a bit disorganized in presenting its main ideas. These reasons prevent me from increasing the final score to 5 at this point. At the same time, I believe the paper already has sufficient merit to be considered for acceptance, hence, my recommendation remains at Borderline Acceptance.", "rebuttal": "Thank you for your thoughtful summary and review! We’re so glad you found the paper insightful, interesting, well-motivated, well-timed, and well-organized. We’ll comment on the weaknesses you identified below:\n\n> The curriculum learning claims are based entirely on the toy case study of compositional multitask sparse parity. While these experiments themselves are quite thorough, they do not guarantee transfer to real-world cases: the authors note this in §3.2. Limitations noted in §7 also apply.\n\nWe agree that the compositional multitask sparse parity task and the group lasso regularization method would be worth studying in more detail in future work, and that our work could be improved by including additional evidence of realistic curriculum learning beyond the toy setting of CMSP. However, we want to emphasize that our CMSP experiments alone will be of substantial interest to the theory community: they provide an extremely simple model organism for studying curriculum learning and distributed representations. For instance, the original exploratory study of the MSP task in Michaud et al. (NeurIPS 2023) then inspired the theoretical work of Fonseca et al. (NeurIPS 2024) and together these helped inspire Ren at al (2025). We are hopeful that the CMSP task we introduce here could inspire similar follow-up efforts in the theory community towards understanding curriculum learning, and that this possibility may make the paper of interest to a particular audience at the conference. We also note that prior work has studied hierarchical learning dynamics in more realistic settings, which we mention on lines 40-41 and 302-303, providing a stronger foundation for the claim that curriculum learning is relevant to real-world cases.\n\n> The authors conduct LLM experiments in §5, but in contrast, the vision counterpart of these experiments in §4 is confined to MNIST. Given that vision is just as important as language in modern AI systems, I would be curious to see whether the results transfer to more realistic vision datasets -- absence of these results does impact the paper's significance.\n\nWe agree that it would strengthen the paper to include more comprehensive vision model experiments and we wish that we had had the time and space to include these! While we didn’t have time to run them in the rebuttal period, we may have time to run them in the discussion period. Do you have any recommendations? One option would be to run experiments similar to our MNIST results but on CIFAR-10 and CIFAR-100.\n\nAnother important observation in the paper was that the group lasso regularization was very effective for targeted unlearning of CMSP tasks. To expand on this, we could extend our LLM and vision experiments to analyze whether group lasso regularization is also an effective unlearning method in these more realistic settings. For instance, we could fine-tune our Python-specific LLMs on a variety of other tasks and compare the networks pruned with group lasso training to more naive pruning baselines. We could also run these experiments in the discussion phase if they would be of interest?\n\nReferences:\n\nMichaud, Eric, et al. \"The quantization model of neural scaling.\" Advances in Neural Information Processing Systems 36 (2023): 28699-28722.\n\nFonseca, Nayara, et al. \"An exactly solvable model for emergence and scaling laws in the multitask sparse parity problem.\" Advances in Neural Information Processing Systems 37 (2024): 39632-39693.\n\nRen, Yunwei, et al. \"Emergence and scaling laws in sgd learning of shallow neural networks.\" arXiv preprint arXiv:2504.19983 (2025)."}, {"review_id": "ZNjJJMEriE", "reviewer": "Reviewer_yXiG", "summary": "The main question that this paper aims to answer is:\n\nwhat if we try to create a specialized (\"narrow\") model by extracting specific skills from a more general-purpose larger model? Which are the main challenges in doing so?\n\nThe authors focus on two specific challenges: \n1) some narrow tasks will require broad training curricula because of hierarchical dependencies between tasks. They explore this issue using a synthetic task called CMSP (compositional multitask sparse parity -- it is a variation of a task that another proposed earlier). \n2) nonlocality of representations: the problem here is it is hard to isolate specific skills through direct pruning because of the distributed way in which neural networks learn representations.\n\nThe paper presents some structured pruning methods (that are not new themselves -- e.g., group-lasso regularization) to extract and transfer narrow skill subsets into smaller networks.", "questions": "* I suggest that the authors read the following paper:\nhttps://arxiv.org/pdf/2406.08658  (Pruning is Optimal for Learning Sparse Features in High-Dimensions)\nIt is a more theoretical paper but I think it is very relevant.\n\n* One of the paper's main claims is that the  distributed nature of representations is a significant challenge for pruning. However the pruning can be done during training or fine-tuning for that more narrow/specific task. Would such an approach challenge this main claim?\n\n* The paper provides some limited results for experiments involving LLMs (Llama), evaluating on Python token prediction. I admit that I found that part of the paper rather speculative and not very convincing. For example, the authors state about Fig-6 \"we tentatively find that pruning substantially outperforms training from scratch and distilling a model from scratch on the data-parameter frontier\" -- or later in the same section: \n\"in our case it does not seem like the relevant features for our task are localized into a set of “Python” vs.“non-Python” neurons, at least that attribution pruning identifies.\"\nIt is important to have such experiments because they bring the paper closer to the \"real world\" -- but it is also necessary for these experiments to be done more deeply/extensively and the analysis of the results to be more conclusive.", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Main strengths\n* The topic (at a high level) is significant because extracting narrow, specialized AI models (through pruning, distillation, etc) efficiently from general-purpose models is a practically important area for scaling up AI systems.\n* The paper’s experimental setup is clearly described. Figures and visualizations are informative and well done.\n* The writing is also clear and good.\n\nMain Weaknesses\n* In my opinion, the paper does not have a very clear result or \"message\". It reads more as an \"exploration report\",  showing some challenges in extracting narrow models from more general-purpose models -- but it is not clear at the end of the day whether these challenges would still exist if the \"extraction process\" was different. The conclusions of the paper are specific to the specific synthetic hierarchical task that is considered and the specific pruning methods that are considered.\n* There is a lack of theoretical or clear analytical results explaining why and how hierarchical curriculum learning and non-locality arise from neural network training. The results rely on empirical observations.\n* The main insights (hierarchical curriculum learning, nonlocality, the use of structured pruning methods) do not come across as highly original. These are well understood concepts in the broader neural network literature.", "was_revised": true, "final_justification": "I have increased my score from borderline reject to borderline accept after reading the rebuttal and the more positive review of the paper by reviewer gjJ8.", "rebuttal": "Thank you for your comments and suggestions! We’re glad you found the paper’s topic significant and practically important, and that you found its experiments, figures, and writing clear and informative. We’re grateful for this opportunity to clarify some points and improve the paper in response to your comments and questions:\n\n> In my opinion, the paper does not have a very clear result or \"message\". It reads more as an \"exploration report\", showing some challenges in extracting narrow models from more general-purpose models -- but it is not clear at the end of the day whether these challenges would still exist if the \"extraction process\" was different. The conclusions of the paper are specific to the specific synthetic hierarchical task that is considered and the specific pruning methods that are considered.\n\nWe agree that the paper has more of an exploratory structure than the average machine learning paper. However, we think that this exploratory structure is appropriate given that the paper focuses on understanding some more conceptual questions about how networks learn and implement their skills that are broadly relevant to understanding the performance of many different methods. The paper makes connections between subfields (like learning dynamics, mechanistic interpretability, pruning, and unlearning) and lays a conceptual groundwork for future work. While this exploratory tone is somewhat nonstandard, we think the paper will still be of interest to folks from multiple communities. We also think that many of our specific results are interesting and novel, including (1) the introduction of the CMSP task as a model organism for understanding curriculum learning, distributed representations, and unlearning, (2) the comparison of which methods are most efficient for creating narrow models, and (3) the use of sparsity regularization training for unlearning. We appreciate the reviewer raising this concern, and we will aim to better emphasize and contextualize these contributions in the final version of the paper.\n\n> There is a lack of theoretical or clear analytical results explaining why and how hierarchical curriculum learning and non-locality arise from neural network training. The results rely on empirical observations.\n\nWe note that there is already a strong theoretical basis for nonlocality arising from neural network training, which we discuss in detail in lines 172-183. While we agree that a mathematical formalism for curriculum learning would be useful, we think that there is substantial value in introducing CSMP and studying it empirically as a first step. For instance, the original exploratory study of the MSP task in Michaud et al. (NeurIPS 2023) then inspired the theoretical work of Fonseca et al. (NeurIPS 2024) and together these helped inspire Ren at al (2025). We are hopeful that the CMSP task we introduce here could inspire similar follow-up efforts in the theory community towards understanding curriculum learning, and that this possibility will make the paper of interest to the theory community regardless.\n\nThe main insights (hierarchical curriculum learning, nonlocality, the use of structured pruning methods) do not come across as highly original. These are well understood concepts in the broader neural network literature.\n\nWe think there is value in drawing connections between these insights, and that many of our specific contributions do provide significant novelty. For instance, the CMSP task is novel, and while sparsity regularization had been used in prior work to help with structured pruning, it had not been studied for its use in unlearning, and prior work had not analyzed whether it is efficient for creating narrow AI systems (which is a significantly different task from standard pruning). We will do a better job of clarifying these contributions in the final paper.\n\n> I suggest that the authors read the following paper: https://arxiv.org/pdf/2406.08658 (Pruning is Optimal for Learning Sparse Features in High-Dimensions) It is a more theoretical paper but I think it is very relevant.\n\nThanks so much for this suggestion! This work is very relevant, and we’ll be sure to cite it in the final paper.\n\n\n> One of the paper's main claims is that the distributed nature of representations is a significant challenge for pruning. However the pruning can be done during training or fine-tuning for that more narrow/specific task. Would such an approach challenge this main claim?\n\nWe expect that the distributed nature of representations will be relevant whenever one does pruning. Is this question specifically asking about the approach of slowly pruning model components without regularization while fine-tuning on a narrow task? If I recall correctly, we did some experiments on this method several months ago, and while it works somewhat, it was pretty sensitive to hyperparameters and not obviously better than first pruning and then finetuning. If you’d like we could revisit these results and discuss them in the discussion period? Either way, we don’t think it challenges the claim that distributed representations make pruning more difficult than it otherwise would be (if one could prune concept-specific neurons without needing to alter the network’s representations at all).\n\n> The paper provides some limited results for experiments involving LLMs (Llama), evaluating on Python token prediction. I admit that I found that part of the paper rather speculative and not very convincing. For example, the authors state about Fig-6 \"we tentatively find that pruning substantially outperforms training from scratch and distilling a model from scratch on the data-parameter frontier\" -- or later in the same section: \"in our case it does not seem like the relevant features for our task are localized into a set of “Python” vs.“non-Python” neurons, at least that attribution pruning identifies.\" It is important to have such experiments because they bring the paper closer to the \"real world\" -- but it is also necessary for these experiments to be done more deeply/extensively and the analysis of the results to be more conclusive.\n\nWhile we were humble in our presentation of these results, we do think that they are informative and impactful in their current form. We think the ‘”Python” vs. “non-Python”’ neuron pruning results are already pretty conclusive, given that the attribution method tracks the true importance of the neurons pretty well (Figure 11, left). In the discussion period, we could provide some further evidence by doing complete ablations of the neurons within a layer and showing that the distribution over effect sizes on Python documents is fairly uniform, if that would be helpful.\n\n\n\nReferences:\n\nMichaud, Eric, et al. \"The quantization model of neural scaling.\" Advances in Neural Information Processing Systems 36 (2023): 28699-28722.\n\nFonseca, Nayara, et al. \"An exactly solvable model for emergence and scaling laws in the multitask sparse parity problem.\" Advances in Neural Information Processing Systems 37 (2024): 39632-39693.\n\nRen, Yunwei, et al. \"Emergence and scaling laws in sgd learning of shallow neural networks.\" arXiv preprint arXiv:2504.19983 (2025)."}, {"review_id": "ENcEU9Dh44", "reviewer": "Reviewer_gjJ8", "summary": "This paper explores the problem of building narrow AI systems, models specialized for a single task or domain, by investigating two core challenges: 1) when can a narrow model be trained from scratch 2) how do you transfer specific skills from a pretrained general model.\n\nFor challenge 1, the authors introduce a synthetic benchmark task, compositional multitask sparse parity (CMSP), designed to exhibit hierarchical dependencies among subtasks. Using CMSP, they show that training on broader data distributions can enable more efficient learning of composite tasks.\n\nFor challenge 2 the paper compares methods for narrowing models with a focus on pruning, distillation and group lasso regularisation - a regularisation method introduced by the authors to encourage the model to learn tasks in a way that better supports pruning. The authors report that pruning + recovery training (continued training on the desired task) often outperforms distillation and training from scratch at preserving desired skills while reducing model capacity. Experiments are conducted on CMSP, MNIST and Python documentation using LLMs.", "questions": "- Have you done any exploration into tasks that have a hierarchical structure in the presentation of the data but not necessarily the task? For example the use of abstractions in image based tasks, where there is an underlying structure that can make the task easier but is not directly tied to the result.\n- By removing the ability to recover performance on a different task, is it possible to degrade performance on your current task if they are closely linked? This seems to me to be a bit of a random proxy to be optimizing for that could possibly lead to degradation of the whole system if done in some settings.\n- As the paper stands, I would raise my score if the figures are better labelled and explained, the language and discussion was made a bit more rigorous and if the limitations section was a more critical reflection of the work actually done in this paper.", "limitations": "The authors did not efficiently evaluate the limitations of their work. They show the \"simple pruning strategy\" can outperform distillation so claiming this as a limitation is a bit hollow. Similarly, the LLM is used as a baseline so while the results may change if the LLM was well tuned the lessons learnt from the LLM experiment should still hold true. The rest of the limitations is just philosophical musings that are more a statement about the subfield of hierarchical learning and not about the work. Instead the authors could have mentioned their lack of task diversity in their synthetic study, lack of evaluation of group lasso regularisation in the LLM experiment and the fact that much of their analysis is done only on synthetic tasks.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n- The paper focuses on narrow AI systems, something that is less common after the recent successes of general-purpose models. This is still an important area of research as we have evidence from biological agents of the benefit of the combination of smaller, focused models and so definitely deserves further exploration.\n- The CMSP benchmark nicely introduces a hierarchical dependency between tasks and can act as a useful tool for future investigations of tasks with this very hierarchical structure.\n- The paper focuses on an interesting problem that is often pushed aside in current works, the overlapping neuron representations of different skills in larger models, which can make pruning challenging and reduces explainability.\n- The paper is laid out well, building both the narrative of their method and the use of their experiments to help the reader understand both the challenges of this problem and the proposed solutions.\n\n**Weaknesses**\n- The paper claims to explore when it is possible to train narrow models from scratch and concludes that \"Our work shows that for some types of tasks, it may be necessary to train on a broad data distribution in order to learn.\" The authors only investigated one type of task structure - highly hierarchical - and found you needed a broad dataset to efficiently learn. This challenge seems poorly investigated and doesn't result in a clear answer.\n- The paper has a heavy focus on results from the synthetic task CMSP. Additionally, they do not look at the effect of their proposed methods on non-hierarchical baselines. The mention two other synthetic tasks that they drew inspiration from which they could have added to their evaluation. While you could argue that the papers that originally introduced these benchmarks already explored the effects of learning these synthetic tasks, it is not clear in the paper that the aim of these investigations was to investigate the learning of narrow tasks which is what this work wishes to explore. The two non-synthetic experiments, MNIST and LLMs on Python documentation, are only briefly mentioned and only offer that pruning works better than distillation without much insight into these findings. Additionally, the group lasso regularisation is not used on the LLM experiment.\n- The paper has a few areas where a speculative reason is provided to explain some phenomena and is then later used as the grounding for a future speculative explanation. For example \"One explanation of these results is that networks compute the parity of the atomic subtask bits in the first hidden layer, and then to learn the composite subtask they compute the parity of these values in the second layer.\" However, this mechanism is never empirically validated. No activation probing, representational analysis, or layer-specific ablation is performed to support this claim. Instead, a comparison between shallow and deeper networks is offered, but this is insufficient to rule out the impact of increased expressivity due to added depth. While a width-matched baseline is included in the appendix, it does not eliminate the core confounding point that deeper networks may simply represent more complex functions by design. This statement is later used to add legitimacy to the broader claim \"This result may also provide some explanation for advantage of depth in deep learning – not only is depth necessary for networks to efficiently approximate certain functions [33], we find here that depth is helpful to efficiently learn tasks with hierarchical structure.\" The use of one unproven explanation as evidence for another raises concerns about the paper’s reasoning rigor in the discussion section.\n- The paper could greatly benefit from another look at the figures and their captions. For example figure 3 has visualizations of 2-hidden-layer MLPs to provide evidence of the use of group lasso regularisation however an explanation for the lines is never given (what do the colours represent? Why is the line thickness different? Is this weight magnitude? Activation size?). I find figure 4 left difficult to parse as the authors have a three dimensional result which they are showing in a 2D graph so gleaning any real patterns requires a lot of thought on the readers side. Figure 5 right has shadow lines but no explanation is given for what this is. It cannot be an STD or error estimate because it is offset from the main lines in some cases and I can't even guess what these shadow lines mean.\n- The limitations section of this paper is poor and lacks an honest self-critique. The authors point out they only look at one, simple pruning strategy and lack of LLM hyperparameter tuning as the only limitations in their work. They then have another paragraph that raises philosophical challenges related to narrow AI and postulate that this entire direction of work could be irrelevant if neural networks learn in a way that would negate this subfield of research. This is not a useful limitations section and ignores some real limitations of the presented work.", "was_revised": true, "final_justification": "I think this is an interesting piece of work as I said in my previous review. Given the commitment to improving the presentation I am happy to raise my score.", "rebuttal": "Thank you for this detailed and thoughtful review! We’re so glad that you found the paper’s topic to be interesting and important, that you found the CMSP task to be useful, and that you found the paper to be laid out well! We appreciate the opportunity to improve both the results and the clarity the paper based on your comments:\n\n> The paper claims to explore when it is possible to train narrow models from scratch and concludes that \"Our work shows that for some types of tasks, it may be necessary to train on a broad data distribution in order to learn.\" The authors only investigated one type of task structure - highly hierarchical - and found you needed a broad dataset to efficiently learn. This challenge seems poorly investigated and doesn't result in a clear answer.\n\nWe think that our results on the hierarchical task, CMSP, are impactful since they are a proof-of-concept that extremely strong curriculum learning effects can exist when the data has a particular structure. We think that this task provides a very nice “model organism” for studying curriculum learning and nonlocal representations that will be of independent interest to the theory community. We do not claim that all realistic tasks display this structure; rather, we use CMSP as an existence proof that some tasks are necessarily hierarchical. We thank the reviewer for pointing this out and plan to revise the paper to more clearly articulate this point.\n\n> The paper has a heavy focus on results from the synthetic task CMSP. Additionally, they do not look at the effect of their proposed methods on non-hierarchical baselines... Additionally, the group lasso regularisation is not used on the LLM experiment.\n\nAh, thanks for this recommendation! To provide a clear baseline, It would have been appropriate for us to also study curriculum learning and our group lasso regularization method on the simpler MSP task, as well as on CMSP, and we could include results on this in the final version if you’d like! We’d guess that MSP won’t have very meaningful interactions between the subtasks (since the relevant task bits are mostly distinct between subtasks), but it is worth checking. We also note that if there were meaningful interactions between subtasks on baselines like MSP, this would demonstrate that these tasks are also hierarchical, rather than invalidating the curriculum learning dynamics seen on CMSP.\n\nWe’re also sorry that our discussion of our results on real-world tasks (MNIST and LLMs) was not clear. On MNIST and LLMs, we did meaningfully study the creation of narrow networks with the same methods we used on CMSP -- pruning without and without group sparsity regularization, distillation, and training from scratch. You said that we do not use group lasso regularization on the LLM experiments, but Figure 5 shows the results of using group sparsity regularization to prune LLM MLP neurons. We will aim to improve the presentation of these results in the final version of the paper -- Figure 5 is indeed a complicated figure, and its caption could be improved.\n\n\n> The paper has a few areas where a speculative reason is provided to explain some phenomena and is then later used as the grounding for a future speculative explanation...\n\nThank you for the opportunity to improve the rigor of this section. In the sentence “One explanation of these results”, we meant to say “One potential explanation of these results could be…”. If you’d like, we could also perform activation probing in the discussion period to test this claim -- this is a great suggestion. Also, the sentence “This result may also provide some explanation for advantage of depth in deep learning – not only is depth necessary for networks to efficiently approximate certain functions [33], we find here that depth is helpful to efficiently learn tasks with hierarchical structure” is based on our empirical result that deeper, parameter-matched networks learn CMSP faster on average than shallow networks. We don’t think it depends on our earlier speculation being true that networks were solving the task in a particular manner. However, we do think it would be valuable to investigate whether this increased training efficiency is a direct result of the hierarchical structure—for instance, we could see if this same effect is present on non-hierarchical baseline tasks, if you would find this helpful. We thank you for suggesting that we further disambiguate these possible explanations.\n\n> The paper could greatly benefit from another look at the figures and their captions. For example figure 3 has visualizations of 2-hidden-layer MLPs to provide evidence of the use of group lasso regularisation however an explanation for the lines is never given (what do the colours represent? Why is the line thickness different? Is this weight magnitude? Activation size?). I find figure 4 left difficult to parse as the authors have a three dimensional result which they are showing in a 2D graph so gleaning any real patterns requires a lot of thought on the readers side. Figure 5 right has shadow lines but no explanation is given for what this is. It cannot be an STD or error estimate because it is offset from the main lines in some cases and I can't even guess what these shadow lines mean.\n\nThank you for drawing our attention to this problem! We’d be happy to clarify this in the final version of the paper. To answer your questions briefly:\nIn Figure 3, the colours represent whether the weights were positive or negative, however that doesn’t influence the takeaway of that figure, which is just that the regularized network is far sparser than the unregularized network. The transparency “alpha” parameter of these lines as well as their width is set based on the magnitude of the weight, which is all we really care about here. Thank you for pointing out that this is not clear enough in the current figure and plan to address it in the caption for the final version. \nIn Figure 5, the shadow lines show the loss over training on each batch, while the primary lines show an exponential moving average of the losses. We will also clarify this in the caption. \n\n> The limitations section of this paper is poor and lacks an honest self-critique. The authors point out they only look at one, simple pruning strategy and lack of LLM hyperparameter tuning as the only limitations in their work. They then have another paragraph that raises philosophical challenges related to narrow AI and postulate that this entire direction of work could be irrelevant if neural networks learn in a way that would negate this subfield of research. This is not a useful limitations section and ignores some real limitations of the presented work.\n\nThank you for pointing this out. We will substantially expand the limitations section in the final version of the paper. Part of the reason why the limitations section is currently short is simply that we were running up into the page limit. In our final version, we will include some discussion about the limited task diversity in our synthetic study and the fact that parts of our analysis are done only on synthetic tasks. Note however that we actually did apply group lasso training to LLMs, as discussed above. \n\nAlso, the more philosophical final paragraph is not really part of the limitations; this will be more clear in the final version of the paper.\n\n> Have you done any exploration into tasks that have a hierarchical structure in the presentation of the data but not necessarily the task? For example the use of abstractions in image based tasks, where there is an underlying structure that can make the task easier but is not directly tied to the result.\nWe have not! That could be interesting in future work!\n\n> By removing the ability to recover performance on a different task, is it possible to degrade performance on your current task if they are closely linked? This seems to me to be a bit of a random proxy to be optimizing for that could possibly lead to degradation of the whole system if done in some settings.\n\nWhen doing group sparsity regularization, the network is being trained to maintain its performance on a given task and so its performance typically does not degrade unless one regularizes too strongly. Note that the group lasso sparsity penalty is applied in conjunction with a standard data loss term (eg, cross entropy loss on the relevant task).\n\n> As the paper stands, I would raise my score if the figures are better labelled and explained, the language and discussion was made a bit more rigorous and if the limitations section was a more critical reflection of the work actually done in this paper.\n\nThank you for the thoughtful engagement with our results and the opportunity to raise our score by making our figures clearer and limitations more rigorous. While we can’t update the manuscript in the discussion period, we believe that it would be quite straightforward to make these improvements to the clarity and rigor of the work that we’ve discussed above. If necessary, we could draft these improvements in the discussion period and present them in text form (e.g. for the new limitations section and some captions) if you’d like? Either way, we agree that the figures could be clearer and the limitations could be more rigorous, and we will be sure to include the aforementioned changes in the final version of the paper."}, {"review_id": "XXDx97ZFhm", "reviewer": "Reviewer_tPS8", "summary": "The submission investigates the creation of narrow, specialized models. The authors argue that hierarchical tasks in some cases require broad pre-training of larger models. Experiments on synthetic data show that such cases can exist. Further, the authors evaluate whether specialized knowledge can be localized in the weights of models, and could be isolated by pruning the other weights. They find that specialized knowledge is entangled in weights, but can be dis-entangled by fine-tuning on the target data with a group-lasso regularization.", "questions": "See strengths and weaknesses.", "limitations": "yes", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "### Strengths\n\n+ The motivation of narrow models is convincing, especially via efficiency. \n\n+ The paper presents an interesting case study on learning dynamics of compositional tasks. While the results are intuitive (atomic tasks learned first, composite tasks learned after), I wonder if the experiments conflates compositionality with 'hardness' of the task and can support the conclusion of hierarchical learning. Presumably, compositional tasks are significantly harder and so will be picked up later. To get clean results, the two need to be disentangled. I'm not sure if CMSP allows for that.\n\n+ Albeit not entirely surprising, the combination of fine-tuning with l1 regularization and subsequent pruning is a reasonable idea, and the experiments support the combination. I'm not sure if similar ideas have been tried before since it seems such a natural combination, but its interesting that localization can be strenghtend post-hoc even on pre-trained models.\n\n+ The submission is generally clear and well written. Experiments are well motivated. I especially appreciate that the authors don't embellish their results and discuss limitations.\n\n### Weaknesses\n\n- The motivation via mechanistic interpretability or verifiability is not convincing to me. Differences between full and narrow models are probably in scale, not in general composition. Particularly in cases like coding task, where models are qualitatively similar. Hence, methods that work on narrow models can also be expected to work on general models.\n\n- The authors make a strong assumption that (some) skills are hierarchical. This is not my field, but I keep wondering. Hierarchical skills exist - but is there actual evidence for it? Is that relevant in practice, or is this a notion carried over from what and how humans learn? Even if it exists, its unclear to what extend the CMSP experiments represent real-world compositionally and how far findings might carry over. The authors touch upon that limitation in line 159. \n\n- While I like Figure 5, I feel it only shows part of a picture. Specifically if random pruning achieves similar performance, it appears that the overall computational load might be lowest to just randomly prune to the desired sparsity and follow up with recovery steps. The group-lasso steps appear to keep the loss low for a low recovery step regime, but afterwards at best match the baseline. \n\n- I keep wondering how the two parts of the paper are connected. The first considers hierarchical skills and how easy / hard they are to learn. The second studies extracting skills from a larger model into a smaller one. The two appear disjunct, albeit they may both be part of a larger pipeline. If the argument the authors wish to make is that one should pre-train broadly to learn complex composite skills faster, and then prune aggressively, I would encourage them to tie the two together more tightly in an overarching narrative. In the second half, the evaluation uses a Llama on code. Maybe running experiments on hierarchical skills in that domain might help to form a coherent narrative?", "was_revised": true, "final_justification": "I appreciate the author's response. It demonstrate awareness of room for improvement. To me, the paper demonstrates very good ideas, but they seem a unfinished. I encourage the authors to improve their submission, regardless of whether it get's accepted this time. For that reason, I keep my score.", "rebuttal": "Thank you for your thoughtful comments! We’re glad you found the paper’s motivation convincing, interesting, and that it was clear and well written. We note that in addition to group lasso regularization, we also investigate a number of other methods for creating narrow networks, such as distillation, attribution pruning, random pruning, and training from scratch as a baseline. We’ll respond to your comments below:\n\n> The motivation via mechanistic interpretability or verifiability is not convincing to me. Differences between full and narrow models are probably in scale, not in general composition. Particularly in cases like coding task, where models are qualitatively similar. Hence, methods that work on narrow models can also be expected to work on general models.\n\nOur hope is just that more narrow models will have fewer mechanisms for us to enumerate and explain! It could be worth including some additional discussion of this point at the end of the paper, along with our more philosophical point about how to understand the computation that models are doing in the first place.\n\n> The authors make a strong assumption that (some) skills are hierarchical. This is not my field, but I keep wondering. Hierarchical skills exist - but is there actual evidence for it? Is that relevant in practice, or is this a notion carried over from what and how humans learn? Even if it exists, its unclear to what extend the CMSP experiments represent real-world compositionally and how far findings might carry over. The authors touch upon that limitation in line 159.\n\nYes it is unclear right now to what degree curriculum effects like those we observed on CMSP influence LLM learning, for instance. But we think the CMSP task is an interesting existence proof that certain datasets can have extremely strong effects like this!\n\n> While I like Figure 5, I feel it only shows part of a picture. Specifically if random pruning achieves similar performance, it appears that the overall computational load might be lowest to just randomly prune to the desired sparsity and follow up with recovery steps. The group-lasso steps appear to keep the loss low for a low recovery step regime, but afterwards at best match the baseline.\n\nYes this is accurate! We could have made this point more clearly, but we did mention that: “We find that despite the loss increasing substantially after naively pruning, we can also quickly recover that lost performance, and overall this strategy seems to be more efficient than using group lasso training. We therefore next compare naive pruning + recovery training against distillation and training networks from scratch”. It’s possible that group lasso regularization would more effectively unlearn LLM skills vs. random pruning, but we haven’t yet tested this (but could in the discussion period next week, if you’d find that helpful!).\n\n> I keep wondering how the two parts of the paper are connected. The first considers hierarchical skills and how easy / hard they are to learn. The second studies extracting skills from a larger model into a smaller one. The two appear disjunct, albeit they may both be part of a larger pipeline. If the argument the authors wish to make is that one should pre-train broadly to learn complex composite skills faster, and then prune aggressively, I would encourage them to tie the two together more tightly in an overarching narrative. In the second half, the evaluation uses a Llama on code. Maybe running experiments on hierarchical skills in that domain might help to form a coherent narrative?\n\nAnother idea we had of tying the two parts together would be to evaluate whether group lasso sparsity training leads to more effective unlearning in LLMs, like we showed in CMSP. For instance, we could fine-tune our Python-specific LLMs on a variety of other tasks and compare the networks pruned with group lasso training to vs. more naive pruning baselines. We could also run these experiments in the discussion phase if they would be of interest?"}]

Author rebuttal:

Meta-review: This paper explores the challenges of creating small, specialized AI models, finding that some narrow skills paradoxically require training on broad data due to hierarchical learning dependencies. When transferring capabilities from large to small models, the authors show that pruning can outperform distillation even though skills are not neatly localized, and they investigate using regularization to better align skills with prunable components. The paper presents experiments on a synthetic task (compositional multitask sparse parity), MNIST, and language modeling on Python documents.

This is a well-written, creative paper, and the experiments including the novel synthetic task are well-chosen. It frames known techniques such as pruning and distillation in an original problem setting that has the potential to inspire interesting future work. Multiple reviewers noted that the study was still somewhat inconclusive, but the consensus among reviewers and AC was that there were enough interesting findings to merit publication.
Decision: Accept (poster)