Gradient-based inference of abstract task representations for generalization in neural networks

Gradient based inference improves interpretability during training

We thank the reviewer for this comment. The reviewer notes that our framework may have an additional benefit we have not considered. Our method involves neural networks inferring a low-dim task abstraction prior to performing a task, observing what low-dimensional task abstraction the network infers may indeed contribute to the interpretability of neural computations.

Please include a diagram of the GBI-LSTM architecture for the toy and language tasks. Specifically which synapses are modified during training and inference time i.e. which is the task abstraction layer/weights z? (Pg 5, line 242)

We agree. A diagram would more visually and clearly define what is being optimized. We now include this in Fig 3 as the first subpanel. Specifically, Z is a vector of units that feeds as input to the neural network through a with a standard set of weights (i.e., a weight matrix with dimensions [dim(Z), neural network hidden units size]). Importantly, the projection from Z to the network is optimized during training as part of the neural network parameters. During inference we only optimize the activations of the Z units (i.e., their neural activation (firing rates)).

The idea to optimize only the input representation weights instead of the entire model is not novel. This idea dates long back to the idea of learning schemas and adjusting new information to fit the prior learned template (Lampinen, McClelland 2020 PNAS; Kumar et al. 2024 arXiv 2106.03580).

Thank you for pointing this out. Upon reflection, we recognize that our previous descriptions may have led to a misunderstanding. Unlike the methods cited, which optimize input representation weights or low-rank perturbations to adapt to new tasks while treating the neural network as a fixed reservoir, our approach differs in key ways. Specifically, we do not optimize the task abstraction input weights to the network. Instead, we optimize the low-dimensional activations of the Z units (e.g. 2 units in the toy task, and 10 units in image generation experiments). We focus on understanding how neural networks can be trained to respond to low-dimensional task abstractions, and then infer them by optimizing the task abstraction units directly during testing. This distinction sets our work apart from prior methods that rely on gain modulation or low-rank updates to preserve prior knowledge.

giving task category as input should significantly reduce the training complexity of needing to infer the task (Kumar et al. 2022 Cerebral Cortex).

We found Kumar et al. 2022 and 2024 to be very relevant work with similar motivation to disentangle learning a computation from forming a representation of the computation itself. We now cite these works as models with a distinct task representation layer. Such methods are expected to reduce the complexity of training and create adaptable models that can be adapted simply by changing their task representation input.

Why was the difference in training loss not as apparent? Did the authors perform hyper parameter sweeps for the learning rate and number of units? It is easy to choose a set of hyper parameters where the distinction between LSTM and GBI is artificially similar.

We agree with these observations. The difference in training is modest in our case because the tasks are of high complexity (e.g. learning wikipedia dataset) while we provide coarse task abstractions (e.g. dataset identifier). We expect that models that discover task abstractions from data to have more complex task abstractions that further break down the computation space and lead to larger differences in training loss. There have been a few examples of models that do this of late (Hummos, ICLR 2023; Butz et al. 2019, Sandbrink et al., NeurIPS 2024).

Hummos, A. Thalamus: a brain-inspired algorithm for biologically-plausible continual learning and disentangled representations. (ICLR, 2023).

Butz et al., Learning, planning, and control in a monolithic neural event inference architecture. Neural Networks 117, 135–144 (2019).

Sandbrink et al., Neural networks with fast and bounded units learn flexible task abstractions. (NeurIPS 2024, spotlight)

I appreciate the authors' efforts in addressing my concerns. Nevertheless, I believe it is necessary to maintain the current score. I recommend conducting additional experiments with various architectures to determine if there are significant performance improvements. Moreover, the submission would benefit from analysis and ablation studies on representational motifs to better understand their impact on generalization.

Clear and relevant objective. The aim of the work is clearly defined.

Thanks for the encouraging remarks.

Although I appreciate that the authors examined their method on various scenarios, I am not sure if the complexity of these tasks is enough.

Thank you for highlighting this important point regarding task complexity. We agree that the tasks examined in our study are relatively simple, and the reviewer's concern about scaling to larger vision or language datasets is valid. Our exploration of scaling to the CIFAR-100 dataset indeed revealed valuable insights about the requirements for meaningful scaling.

In what follows we expand on two main points. A promising class of models have emerged of late with neural models capable of forming their own task abstractions directly from data. First, we offer a foundational contribution to these models by exploring the properties of gradient-based inference, which they heavily use. Second, for meaningful scaling up of our results, future work will have to rely on those models to provide richer task abstractions, beyond the human-provided labels we use in this work.

A promising direction is an emerging class of models that relies on gradient-based Expectations Maximization dynamics to identify tasks in their training data and label them with internally generated task abstractions (Hummos, ICLR 2023; Butz et al. 2019, Sandbrink et al., NeurIPS 2024). These models simply optimize \theta (the neural network parameters) and $Z$ (the task abstraction layer) through gradient descent, with $Z$ having a faster learning rate. This straightforward setup can dynamically form task abstractions in $Z$, allowing the network parameters to organize into modules specific to each task. Such models have demonstrated advantages in mitigating catastrophic forgetting, enhancing adaptability, and improving generalization.

Hummos, A. Thalamus: a brain-inspired algorithm for biologically-plausible continual learning and disentangled representations. (ICLR, 2023).

Butz et al., Learning, planning, and control in a monolithic neural event inference architecture. Neural Networks 117, 135–144 (2019).

Sandbrink et al., Neural networks with fast and bounded units learn flexible task abstractions. (NeurIPS 2024, spotlight)

However, these internally generated task abstractions pose challenges for evaluation as they can drift significantly during training. It becomes difficult to assess the accuracy of gradient descent as an inference mechanism—how effectively can it retrieve previously learned tasks, handle uncertainty in the $Z$ space, or detect out-of-distribution data?

Our study addresses the limitation by using human-provided labels as task abstractions, such as image class in the image generation experiment. The gradient-based EM methods use iterative optimization with hundreds of optimization passes, while we here found that one-step gradients might be sufficient. In addition, we provide estimates of how accurate gradient based inference might be.

Our first point here is: this work offers a foundation for gradient-based EM methods. In fact, one of those papers already cited an earlier version of this work, though we will not specify further to maintain anonymity.

However, our efforts to scale up to CIFAR-100 revealed the limitations of these human-provided labels. In this case, the model was given the image class as a task abstraction to guide image reconstruction. Despite this, backpropagation largely ignored the additional task information and relied more heavily on visual features from the encoder. This suggests that image class labels do not significantly reduce variance in the pixel space, limiting their utility to the model. This is reflected in low accuracy for GBI in inferring image class. We concluded that low-complexity, human-provided task abstractions are insufficient for describing more complex datasets.

Our second point, to meaningfully scale up the framework in this work will have to rely on gradient-based EM methods for richer task abstractions.

Reflecting on this, we recognize that our original manuscript could have framed these motivations more clearly, which would better highlight the unique contributions of this work. We now describe these considerations in the introduction to motivate the study of gradient based inference. We also add additional discussion as we introduce CIFAR100 results and explain what it will take to scale up.

We hope that the possibility of a contribution to answer foundational questions for this novel class of models might offset the limited complexity of the datasets this work tackles.

The training methods used in the experiments (sequence prediction by an RNN and data reconstruction by an autoencoder with its latent variables concatenated with contextual information) are not exactly the same as the one mentioned in Section 2 (MLE of the likelihood function (1)). It is not mentioned in the paper how the extensions to those variants do (or do not) affect the argument regarding the proposal in Section 2.

We appreciate the reviewer for raising this important conceptual point. This took a bit of reflection. To ensure we are aligned on the interpretation of the differences between the theoretical methods and implemented experiments, we welcome further clarification if needed.

Our understanding is as follows: the methods section primarily considers a likelihood function of the form L=f(X,Z). In contrast, the RNN and autoencoder experiments incorporate additional inputs into this computational graph. Specifically, the RNN uses hidden states evolving from previous inputs, while the autoencoder leverages latent encodings produced by the encoder. Notably, these additional inputs are themselves trainable, potentially complicating the theoretical analysis.

The reviewer has rightly highlighted the need to consider how these extensions might influence the behavior of the model compared to the theoretically motivated framework. One way to reconcile this is to view the encoder and RNN as transformations of the input X, effectively modifying its distribution. While this induces non-stationarity in the distribution of X during training, evidence (e.g., Wu et al., 2020 https://arxiv.org/abs/2004.09189) suggests that the encoder often converges faster than the decoder. This faster convergence could reduce the transient non-stationarity from the perspective of the decoder, rendering the extensions to the computational graph less impactful on the overall argument.

Another way to reconcile is to consider the latent variables from the encoder, and the hidden units activations from the RNN, as part of the model parameters. Our framework calls for updating the parameters of the model during training, and we can show that doing so, also updates the encoder latent and RNN hidden activations according to their gradients. Thus we can treat them as part of the model parameters being optimized during training, and our theoretical description of the methods should hold. To briefly state this symbolically for the RNN case:

The hidden state of the RNN is updated as $h_t = W h_{t-1}$. During training, the weights $W$ are updated via gradient descent:

$

\Delta W = -\eta \frac{\partial L}{\partial W} = -\eta \frac{\partial L}{\partial h_t} (h_{t-1})^T

$

where $\eta$ is the learning rate. After the weight update, the new hidden state is given by:

$

h_t^{\text{new}} = (W + \Delta W) h_{t-1} = h_t - \eta \frac{\partial L}{\partial h_t} |h_{t-1}|_2^2

$

with $\|h_{t-1}\|_2^2 = h_{t-1}^T h_{t-1}$

This shows that updating the weights $W$ during training indirectly updates the hidden state $h_t$ as well, effectively performing a gradient descent step on $h_t$ with respect to its gradient $\frac{\partial L}{\partial h_t}$, scaled by $\|h_{t-1}\|_2^2$. This aligns with our framework, where we treat the RNN hidden activations $h_t$ as part of the model parameters being optimized during training.

We are still writing a brief overview of these considerations to add to the paper. Thanks for the helpful comment.

the work brings a fresh perspective and could potentially bring new insights to the field.

We thank the reviewer for the encouraging comments.

The reviewer made several well thought out suggestions to improve the paper, leading to several changes and additions as we detail below.

motivation for gradient-based approach

Thank you for this comment. We agree that our initial submission did not motivate the advantages of using gradients for task inference. Gradients provide two key benefits. First, is avoiding the alignment problem: task abstractions in neural systems influence computations via the parameters they modulate. In other methods, inferred task abstractions can become misaligned as network parameters are updated during training. In contrast, gradients are computed based on the current state of the network, ensuring alignment between the task abstractions and the underlying computations. We have clarified this point in the introduction to better motivate the use of gradients as a solution.

Second, a novel set of models appeared recently that train models by both updating the weights, and also updating a task abstraction layer with a higher learning rate (Hummos, ICLR 2023; Butz et al. 2019, Sandbrink et al., NeurIPS 2024). This simple setup can surprisingly discover tasks in the unlabeled stream of data, represent them as distinct units, and switch task abstractions appropriately to solve previous tasks and compose solutions to new ones. These models highlighted the need for a principled study of inference through gradients, but because the task abstractions are internally generated, and drift during training, it is difficult to assess how well gradient-based inference works in this setting. Our paper, instead, uses human provided labels enabling us to quantify the accuracy of gradient descent in task abstraction space. We now make this connection in the introduction of the revised manuscript.

Hummos, A. Thalamus: a brain-inspired algorithm for biologically-plausible continual learning and disentangled representations. (ICLR, 2023).

Butz et al., Learning, planning, and control in a monolithic neural event inference architecture. Neural Networks 117, 135–144 (2019).

Sandbrink et al., Neural networks with fast and bounded units learn flexible task abstractions. (NeurIPS 2024, spotlight)

The section on the one-step gradient update and maximal entropy initialisation could benefit from a clearer, more intuitive explanation. To improve clarity, the authors could add a visual schematic that illustrates the process step-by-step, and how a single gradient update shifts this initial state towards a more task-specific representation.

Thank you for this thoughtful comment. We created the schematic and it does offer a more intuitive account of how gradients interact with the task abstraction space. This is currently in the new figure 1. Though we are still refining the plot. Thanks for this helpful comment.

Improving coherence: ..A bit more introduction and summarisation at the beginning and end of sections, focused on tying each section to the core ideas would improve the flow of the paper.

Thank you for pointing out the need to better align our experiments with the paper’s motivations. We have made several changes to address this:

1. Clarified the Benefits of GBI: We now group the expected benefits into two categories: (i) during training—faster learning and reduced forgetting, and (ii) during testing—accurate task abstraction inference and recomposition for generalization.
2. Updated Section Titles: Section titles now explicitly state the key findings and align with the claims in the introduction.
3. Added Introductions and Summaries: Each experimental subsection begins with an overview of its goals tying the findings back to the paper’s central ideas.

limited use of intuitive examples: incorporating a few straightforward examples or analogies could provide readers with a more accessible understanding of the contributions being presented.

Very valuable improvement to the paper. We now use an example inspired from Daniel Wolpert’s work on how motor feedback can be also seen as a mechanism to infer other people’s motives during social interactions. We added this to the introduction, and we summarize the example here:

During upbringing, we may observe situations where people’s feelings were labeled as sadness or as anxiety. As we interact with others we may try to infer their emotional states relying on subtle cues. We incrementally adjust our conclusions with every cue, moving closer to one emotion or the other, as make predictions and receive feedback during the interaction. If, by the end of the interaction, both emotions seem equally likely, then either the situation is uncertain or that the person is experiencing an emotion outside of those two.