On the creation of narrow AI: hierarchy and nonlocality of neural network skills

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:

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.

Our 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.

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.

Yes 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!

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.

Yes 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!).

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?

Another 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?

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:

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.

We 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.

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.

Ah, 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.

We’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.

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...

Thank 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.

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.

Thank 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: In 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. In 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.

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.

Thank 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.

Also, the more philosophical final paragraph is not really part of the limitations; this will be more clear in the final version of the paper.

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. We have not! That could be interesting in future work!

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.

When 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).

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.

Thank 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.

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:

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.

We 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.

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.

We 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.

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.

We 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.

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.

Thanks so much for this suggestion! This work is very relevant, and we’ll be sure to cite it in the final paper.

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?

We 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).

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.

While 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.

References:

Michaud, Eric, et al. "The quantization model of neural scaling." Advances in Neural Information Processing Systems 36 (2023): 28699-28722.

Fonseca, 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.

Ren, Yunwei, et al. "Emergence and scaling laws in sgd learning of shallow neural networks." arXiv preprint arXiv:2504.19983 (2025).

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:

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.

We 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.

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.

We 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.

Another 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?

References:

Michaud, Eric, et al. "The quantization model of neural scaling." Advances in Neural Information Processing Systems 36 (2023): 28699-28722.

Fonseca, 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.

Ren, Yunwei, et al. "Emergence and scaling laws in sgd learning of shallow neural networks." arXiv preprint arXiv:2504.19983 (2025).