Flexible task abstractions emerge in linear networks with fast and bounded units

Thank you for your thoughtful comments. We are glad to hear that you find the theoretical analysis of the underlying mechanism insightful. Moreover, we are encouraged that you value our aim to test our framework against behavioral experiments as well as to demonstrate a basic usefulness beyond synthetic tasks.

Q1: How many training instances were in the spiky loss part, in both the synthetic task and the MNIST task? E.g. Does the model achieve single-shot adaptation in the task block?

The reviewer brings up the important question about how quickly is this model able to adapt to previous tasks and new tasks. As we run our simulations in continuous time (gradient flow), each gradient is calculated over a large number of samples (200 for both cases). We made this choice to permit theoretical analysis where a sample average is required. Figure 1F shows the number of time steps needed until adaptation for the synthetic plot, and we supply a similar version for fashionMNIST in Fig. R2.

To investigate the sample efficiency of the model on the synthetic task, we now provide a simulation in Fig. R3 with coarsely discretized time where only a single sample is used in every update. While this makes the simulation more noisy in a single timestep, its qualitative features persist: In late blocks, the loss drops significantly on the first sample, and reaches its minimum after a few samples (few-shot adaptation). More generally, the persistence of the fast adaptation in the equivalent model in Fig. 3 (where noisy input is absent) suggests fast adaptation is indeed facilitated by the gates.

We will include both figures in the appendix of the final revision.

Would the gates adapt to a new task replacing an old task quickly?

If there are some shared components between the old task(s) and new task(s), the gates would adapt to appropriately modulate the parts of the network that correspond to the new task. For example, we illustrate this in our generalization experiments where we switch the network to a new set of tasks (A+B, B+C, A+C) composed of different combinations of the tasks previously encountered (A, B, C). In these cases, the network adapts the gates to turn on the parts of the network corresponding to the new tasks, and turns off the rest of the network, allowing it to generalize its learned knowledge rapidly.

Must gating-like properties be at the 2nd layer in the deep monolithic network? What happens when you insert regularization and faster learning in layer 1 units? In general, for larger-scale applications, is it better to have gating-like mechanisms in later parts of the model?

Because the deep monolithic model is linear, there is no meaningful difference between gating occurring in the first versus the second layer. The main distinction is that rather than gating the output of computations already performed in the network, inducing gating-like behavior in the first layer would instead gate off certain computations in advance (by zeroing their input) and only perform computations on specific paths within the network. If regularization and a faster learning rate are instead applied to layer 1 units, we would expect to see similar gating-like behavior emerge in the first layer rather than the second. Indeed, the question of whether there is an advantage to gating earlier or later in a network is an interesting open question which may only arise once additional depth and nonlinearity are added to a model.

What happens when non-gating weights are also equipped with some amount of neuron-like properties?

Thank you for addressing the minimal conditions needed for gating. We can distinguish two settings:

First, what happens if we equip the weights with a fast timescale? We find that with very fast weight learning rates, specialization eventually disappears: The weight representations will be fast enough to quickly adapt after a block change, not requiring any gating (Fig. 5). This however is an unnatural assumption, as synaptic plasticity is slow.

Second, we can ask whether an architecture with two layers of weights, but without any scalar gates will show gating-like behavior. By regularizing and increasing the learning rate of one layer, we indeed observe the fully-connected weight layer to exhibit gating-like behavior, while the other layer learns to specialize to the corresponding tasks, emulating the same behavior we see in our gated model. (Fig. 6).

Are there neuroscientific reasons to maintain separate groups of task weights and gate weights?

The reviewer brings up an important question as to how the model structure maps to brain structures. There are multiple gating circuits in the brain. One is the prefrontal cortex (PFC) gating activity in earlier cortices. We think of the weights as the synaptic strengths in motor and sensory cortices as they learn a task, while the gating to switch between distinct computations is projected top-down from PFC. In addition, cognitive flexibility and switching task representations in PFC occurs through rapid changes in neural activity rather than synaptic updates. As such, we see them implemented in different regions, and also through different substrates (i.e. synaptic plasticity vs. neural activity changes).

One comment is that because the paper is very dense, some details are omitted (e.g. in figures) and need to be inferred.

We thank the reviewer for this valuable comment. We reviewed the manuscript and found some important details omitted especially regarding the details of the experiments run in each figure. The final version of the manuscript now addresses this.

Thank you for your review and thoughtful comments. We are delighted that you appreciate the simplicity of our model and the thorough experiments we provide.

The idea is to have subnetworks which can learn each task, and then learn a high level router to the subnetworks.

We thank the reviewer for the accurate summary regarding the routing function of the gating variables in our setup. In contrast to other methods that learn gating of task representations, we update weights and gates simultaneously through gradient descent, reflecting real-world conditions in which learning needs to be segmented and structured at the same time. Thus, they are fundamentally treated on equal footing. Critically, our suggested regularization (alleviating overspecification by favoring c to be of order 1) does not strictly encourage one task to be active at a time, since mixed solutions are given the same penalty.

We are excited about this model because it solves several computational problems with one mechanism: It is able to discover boundaries between tasks, retrieve task representations for previously learned tasks, and move representations for new tasks away from previous ones.

it seems like parts of the architecture, such as the number of subnetworks and the size of the subnetworks are tuned to match those given in the system.

So far, we have indeed only considered the case where the number of subnetworks P equals the number of tasks M (P=M=2) for simplicity. We fully agree that analysing P>M (i.e., multiple paths) improves the understanding of the model. To this end, we conduct a new experiment in the attached figure R4. The system still learns to solve both tasks and adapts flexibly to context switches. In this setting however, students only partially specialize, since they are now underconstrained. This disappears when regularization of student weight magnitude is imposed, which can physiologically be interpreted as a “representational cost”: Over time, we would expect a biological system to prune the additional components.

We note that, under certain conditions, the flexible regime can also emerge in the fully-connected architecture which is inherently free of these assumptions (Fig. 6).

On this basis, we adjust the presentation of our base model to the more general case P>M.

The paper structure seems different from what is typically done for neurips
Many of the experiment details are also in the appendix, which is weird because without the details, it's somewhat hard to understand the results presented in the main text.

Thank you for pointing this out! We had deferred our related work section to the appendix for space constraints, but see how it is valuable in the main text, so that we will incorporate a shorter version there. We will move necessary experimental details back to the main text as well.

$why MNIST vs. a weird permutation of MNIST$ is it because the tasks have to be orthogonal?

This leads to my main question and reservation, which is that the model might only work if the input tasks are orthogonal?

We thank the reviewer for this important comment. We ran several simulations here to identify how orthogonality might impact the model. Prior related work (Lee et al., 2024) provided an extensive investigation of orthogonality of representations in linear networks. For theoretical simplicity, we considered only orthogonal tasks. We ran three new simulations in the PDF attached to this rebuttal.

First, we plot how parameterizing the overlap between teachers affects student specialization (Fig. R1A). Ideally, we expect the model to maintain a shared representation according to the similarity between the tasks while separating the non-similar parts. Indeed, we see such a graded specialization in our model proportional to the overlap between the tasks. .

Second, we train our model on a set of three compositional tasks, with pairwise overlap (e.g., A+B, B+C, A+C), making them non-orthogonal. Our model learns to specialize to the underlying shared components and appropriately gate them to solve each of the three non-orthogonal tasks (Figure R1B).

Third, we move slightly from MNIST to fashionMNIST to operate on natural images, and use two permutations with varying degrees of orthogonality. We use permutations because we are strictly interested in the setting of cognitive flexibility where the same input can be mapped to different responses. The new permutations are based on delineating (orthogonal) upper-vs.-lower body or (correlated) warm-to-cold weather. We find that the system works well in both cases, with specialization occurring only marginally later in the non-orthogonal setting (Figure R2).

In summary, the studied phenomenon is robust to correlations between tasks.

My current score is assuming that this limitation is true, but if it were not true, I would give a higher score (maybe an 8).

We appreciate the reviewer stating their concern and clearly stating how it influenced their decision. We hope that our responses above and the additional simulations addressed the concern about orthogonality fully.

I would explain the regularization constraints are in words besides just presenting the equation.

We discuss two kinds of regularization:

1. Non-negativity of gates. It is motivated by the fact that gating behavior can only amplify or attenuate a variable, but not change its sign.
2. Alleviating overspecification. Due to the linear nature of our architecture, the output scale can be changed both by amplifying weights or gates. However, gates should not need to account for learning the scale of the task. The regularization encourages gates to be of order 1. This has the effect whereby upscaling one gate above equilibrium will drive other gates to be downscaled.

We will add a short explanation to mathematical expressions in-place and emphasize the more extensive treatment in Appendix B.

Thank you for your helpful comments. We appreciate that you found the conducted experiments relevant to our claims.

Investigating these general, simple changes as drivers of flexible task abstractions is creative and exciting.

We are glad that this contrast came across well: Despite the simplicity of the proposed mechanism, the effect on inducing flexible task abstractions is remarkable.

Paper setup is very confusing. It's okay to have a nontraditional structure, but it ends up being a prose-heavy laundry list of various properties of the method and results.

Thank you for this valuable feedback on the organization of our manuscript. We agree that it failed to guide the reader through.

On a high level, our paper has this structure:

Hypothesis: Are neuronal gates updated through gradient descent sufficient to separate tasks into dedicated representations? First, we answer through simulations whether and how implicit context switches are learned (Fig. 2). Second, we analytically address why task representations arise (Fig. 4). Third, we apply the model to cognitive science (Fig. 7). The remaining figures are not essential to this narrative, but function as controls for the necessary conditions for flexible gating.

To address this point, we now strictly keep to this structure:

1. We reorganized the order of the mathematical analysis and grouped subsections therein. We now first precisely define flexible adaptation, and then discuss the main driver for its emergence.

2. We moved the Related works section back to the main text to better state our contribution.

3. We moved figure 3 showing how the learned abstractions generalize to the supplementary, as it does not directly support the main narrative. We clearly motivate the remaining figures as controls.

Contribution

The reviewer raises concerns about the contribution of our work.

We aim to contribute to theoretical cognitive science by developing a minimal neural network model of cognitive flexibility. We build on a line of recent theoretical NeurIPS papers studying the learning dynamics in linear networks and relating them to cognition, starting with Saxe et al. (2013, Exact solutions to the nonlinear dynamics of learning). Later studies examined how gating alleviates interference, but their gating was static and handed to the network (Saxe et al., 2022, Dynamics of abstraction in gated networks). We generalize this line of work with an analytically tractable model of how appropriate gating emerges dynamically.

Our model contributes to cognitive science:

1. We offer, to our knowledge, the first neural network model that benefits from data distribution shifts and struggles in the shuffled data regime, similar to humans.

2. We provide a direct comparison to humans where their task switching behavior accelerates as they practice the tasks involved. This came from both our simulations and our theoretical analysis, and we provide a mechanistic explanation for the phenomenon.

3. Not only does the model infer tasks and retrieve the suitable task abstraction, but similar to humans, further learning or credit assignment is gated by the inferred task. I.e. only the parameters for the inferred task are changed.

We believe that such a model can make additional behavioral predictions and provide a neural basis for explaining additional experimental findings, such as Heald et al’s influential study on sensorimotor repertoires (2021, Nature).

A proven (at least, reliable) model of some biological brain phenomenon is valuable. However, this paper has only shown vague similarity, not much comparative predictive power of this approach vs an unconstrained approach on neural data

We aim to present a minimal hypothesis of how task abstractions in (biological) neural networks might emerge during training.

The reviewer highlights the importance of connecting directly to neural data. Unfortunately, almost no full recordings from animals during the training phase on multiple tasks exists. Recordings are mostly from well-trained animals switching between the tasks (discussed by Bouchacourt et al (2022), Fast rule switching and slow rule updating). Neuroscientists point to a technical difficulty where animals take weeks to months to learn two opposing tasks while recording electrodes slip and drift within days.

Still, technical improvements especially with calcium field imaging might make recordings during learning possible soon: In addition to behavioral signatures, our model predicts gradual differentiation of both task-specific activity vectors in neural space modulated by context that arise over the course of learning. We expect this simple and interpretable signatures to be robustly observed across tasks and species, an advantage over purely Bayesian models and more expressive machine learning architectures.

A machine learning result can also have impact, but the paper doesn't clearly argue for that.

The reviewer identified that the simulated model might be extensible to an ML method for larger non-linear datasets. We view constructing an ML system that incorporates these principles to improve machine learning results as a clearly separate project, which we are considering taking on in the future after having laid the theoretical groundwork with this paper.

The mechanisms appear to be inspired by catastrophic forgetting but then are applied to the related but distinct problem of generalizable task adaptations.

The reviewer highlights that results on the generalizability of our learned task abstractions belong to a different problem and need comparison to different controls such as few-shot learning or meta-learning models. We moved these results to the appendix and clarified they only serve to show that the abstractions are functional and can be recomposed through gradient descent.

We will clarify our contributions in the introduction section in the final revision.