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

Authors: Ali Hummos, Felipe del Rio, Mien Brabeeba Wang, Julio Hurtado, Cristian Buc Calderon, Guangyu Robert Yang
Primary area: transfer learning, meta learning, and lifelong learning
Keywords: Cognitive science, cognitive control, cognitive abstractions, task representations, context-dependent models, variational expectation-maximization

Abstract:
Humans and many animals show remarkably adaptive behavior and can respond differently to the same input depending on their internal goals. The brain not only represents the intermediate abstractions needed to perform a computation but also actively maintains a representation of the computation itself (task abstraction). Such separation of the computation and its abstraction is associated with faster learning, flexible decision-making, and broad generalization capacity. We investigate if such benefits might extend to neural networks trained with task abstractions. For such benefits to emerge, one needs a task inference mechanism that possesses two crucial abilities: First, the ability to infer abstract task representations when no longer explicitly provided (task inference), and second, manipulate task representations to adapt to novel problems (task recomposition). To tackle this, we cast task inference as an optimization problem from a variational inference perspective and ground our approach in an expectation-maximization framework. We show that gradients backpropagated through a neural network to a task representation layer are an efficient heuristic to infer current task demands, a process we refer to as gradient-based inference (GBI). Further iterative optimization of the task representation layer allows for recomposing abstractions to adapt to novel situations. Using a toy example, a novel image classifier, and a language model, we demonstrate that GBI provides higher learning efficiency and generalization to novel tasks and limits forgetting. Moreover, we show that GBI has unique advantages such as preserving information for uncertainty estimation and detecting out-of-distribution samples.

Decision: reject
Reviewer scores: 5, 5, 6, 5

Reviews and discussion:
<thread>
creation date:
1730191549804

thread_type:
official_review

<thread_object>
creator: reviewer_vl85

<content>
summary:
The authors propose a method called GBI to train a recurrent and convolutional neural network to infer the task category during test time. The inference is driven by gradient updates only at the task category layer, which can either be done iteratively or by approximating the maximal entropy point, while keeping the rest of the weights fixed. The authors demonstrated the benefits of this method across a variety of toy, image generation and language generation tasks in terms of a lower training loss. Additionally, the authors argue that once the model learns different task representations, only the task category layer needs to be optimized to decrease test or generalization error.

strengths:
-- included anonymous code for replication
-- demonstrated generality of claim that using gradients to infer task category improves performance across 3 modalities, toy, vision and language tasks.
-- One step gradient update by symbolically estimating the optimal point seems to be novel. 
-- Gradient based inference improves interpretability during training (e.g. bayesian inference like estimation) and to identify OOD samples.

weaknesses:
-- 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)
-- 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). 
-- the claim that GBI-LSTM shows no signs of forgetting compared to the LSTM is not strong. The baseline MSE performance is 0.24, which the GBI maintains for new datasets, but the deviation by LSTM does not seem to be significant (Table 1). Was a statistical test done to compare LSTM and GBI performance? 
-- Why learning to infer task category improves learning and generalization in these tasks is unclear. Perhaps the authors can perform low dimensional analysis to show how the network learns to represent different datasets into non-overlapping subspaces and during inference, the network's activity converges towards a specific prior learned subspace or learns to compose them (Lin et al. 2024 arXiv 2309.04504)?
-- The authors argue that GBI improves generalization loss in language prediction task. Although the baseline LSTM shows consistent loss of 6.8 (I assume all model weights are fixed), the GBI loss starts off higher of around 6.95 and decreases to 6.6 over 100 optimization steps. Does the loss continue to decrease with longer optimization steps? If not, is a generalization loss of 6.6 significant compared to 6.8? 
-- Given that the models are an LSTM and not a large model, I thin it is reasonable to expect at least 30 seed runs instead of 4 as in Fig. 5D to increase the confidence in results, especially when the difference afforded by LSTM and GBI is small. 
-- Since the authors used LSTM for toy dataset, and a CNN for image, it would have been a solid contribution if the authors demonstrated GBI using a simple transformer architecture for language prediction instead of LSTM.

questions:
-- Is the one-step gradient update method novel? Or was it developed prior?
-- giving task category as input should significantly reduce the training complexity of needing to infer the task (Kumar et al. 2022 Cerebral Cortex). 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. 
-- Why was LSTM chosen instead of GRU or a Vanilla RNN? Training an LSTM on the simple toy dataset might be an overkill. 
-- why is the ratio of shared units between LSTM and GBI different in the 0th binned training block (Fig. 2C)? They should aggregated such that the initialized activation is the same for comparison.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730201242096

thread_type:
official_review

<thread_object>
creator: reviewer_4ta1

<content>
summary:
Inspired by cognitive science, the authors identify the lack of an efficient module for inferring the current task in current models. According to the authors, previous approaches to task inference do not meet the following two assumptions:
1. lack of efficient detection if the task has been repeated
2. lack of recomposing mechanism of previously learned tasks.
The authors propose an approach grounded in variational inference, using an expectation-maximization-like framework, and they demonstrate GBI’s effectiveness across synthetic, image classification, and language modeling tasks.

strengths:
1. Clear and relevant objective. The aim of the work is clearly defined.
2. The authors provide code.
3. Experimental validation on various scenarios, from synthetic datasets to complex tasks (image classification and language modeling).

weaknesses:
Although I appreciate that the authors examined their method on various scenarios, I am not sure if the complexity of these tasks is enough. For example, in OOD detection, adding a comparison of CIFAR10 vs SVHN datasets would be valuable.

questions:
I am very curious about catastrophic forgetting matters. The authors gently mentioned this feature of their method, but it was only evaluated on the toy dataset; why? Could you provide more experiments in this area?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730617857821

thread_type:
official_review

<thread_object>
creator: reviewer_mwgp

<content>
summary:
This paper proposes a method of incorporating a task obstruction process to improve the performance of neural networks. It considers the framework in which a likelihood function is maximized with respect to the network weights with task abstraction data included as inputs of contextual information in the training phase and the abstraction data is estimated via gradient descent in the test phase. It proposes a method of approximating the iterative optimization of task obstruction with a one-step gradient update. The proposal is based on insights from neuroscience about the roles played by intermediate abstraction in performing tasks and techniques in variational inference to handle contextual information with latent variables.

strengths:
- A simple and effective method of approximating the iterative optimization of task obstruction is proposed. It is achieved by taking the maximum entropy point as the initial value and updating it with a one-step gradient update. The validity of this approximation is checked experimentally (Fig. 3G, H).

- Experiments show the superiority of the proposed method to the canonical classifier in OOD detection(Table 3, Fig 4A, B, C).

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

- The purpose of the experiments comparing LSTM and GBI-LSTM (Fig. 2, Fg.3D,E, 5A,B,D) is not clear. It seems to me that they are just confirming the impact of the input of task labels for the sequence prediction problems with multiple tasks. It is preferable if this point is clarified in line with the motivations or expectations in Section 1 or Section 2.

- Experiments imply that the canonical classifier may have better accuracy than the proposed methods (Table 2), but no reasoning for this is provided. It does not throw away the value of the superiority in OOD detection. Rather, the trade-off relation between them is worth a remark if it is confirmed to exist.

- The following references are not found in the submission: (L357)"supplementary material". (L396)"Table S8"

questions:
- Please consider the possibility of the clarifications mentioned in the weaknesses section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730691603225

thread_type:
official_review

<thread_object>
creator: reviewer_v6yp

<content>
summary:
The authors introduce Gradient-Based Inference (GBI) of abstract task representations, a method that enables neural networks to infer and adapt task representations dynamically, promoting faster learning and better generalisation. Inspired by human adaptability—where task abstractions allow flexible responses to the same input depending on internal goals—their approach enables neural networks to infer and adapt task representations on the fly.

They frame the setting as an optimisation problem through variational inference, and their GBI uses backpropagated gradients to infer and adjust task representations in a neural network. Experiments in a range of domains including image classification and language modelling demonstrate benefits in learning efficiency, generalisation, and reduced forgetting, as well as its performance in uncertainty estimation and out-of-distribution detection.

strengths:
- originality and contributions: This approach brings insights from cognitive science on human task learning and generalisation to deep learning. Their gradient-based inference (GBI) method introduced in the paper is novel, and seems to be an innovative application of gradients in task inference and adaptation beyond traditional optimization. Furthermore, by positioning GBI as a model capable of estimating uncertainty and detecting out-of-distribution samples, the work brings a fresh perspective and could potentially bring new insights to the field. 
 - significance: the problem setting is an important one: enabling artificial agents to flexibly to situations depending on their varying goals. The connections made to human and animal cognition and learning provide a solid foundation for this setting, grounding the approach in well-established principles of adaptive behaviour and task representation in human learning.
 - The paper demonstrates the effectiveness across varied tasks highlighting its potential as a versatile domain-agnostic approach. Overall their results demonstrate some promising advantages in learning efficiency, generalisation, and uncertainty estimation.
 - The inclusion of the code (via link to an anonymous repo) is appreciated for reproducibility and clarifying implementation details (however the repository would benefit from better organisation, see weaknesses).

weaknesses:
- motivation for gradient-based approach: the motivation for adopting a gradient-based approach could be clearer, as the paper does not sufficiently explain why gradients offer an advantage for task inference and recomposition over alternative methods. While gradient-based updates are obviously often employed in optimisation, it is less intuitive why they would be particularly effective in inferring abstract task representations or detecting out-of-distribution samples. A more thorough discussion grounding and comparing the gradient-based approach with other task inference methods, and meta-learning perspectives could help to clarify the benefits here.
 - 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.
 - Improving coherence: the paper would benefit from better coherence as the flow between sections feels disjointed. This makes it challenging to follow the core narrative. While each section presents important concepts mostly clearly, there is often a lack of clear transitions that tie ideas together leaving the reader to infer the connections. For example, the jump from theoretical discussions of gradient-based inference to experimental details seems abrupt, and then later suffers from limited explanation of how each experiment directly relates back to the proposed framework. 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. 
 - 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. For instance, using a simple scenario (like inferring a task from partial information in an everyday context) could illustrate how the model operates based on gradient information to infer task information and provide useful uncertainty estimates.
 - lack of comparison to important uncertainty estimation methods: while the paper claims advantages in OOD detection and uncertainty estimation, the results provided are not compared against established methods like ensemble approaches.
 - poorly organised code: to improve accessibility, the authors could streamline the repository, making it quicker and easier to connect model implementations with the methods and experiments described in the paper.

questions:
- can the motivation for the gradient-based inference be described more clearly and intuitively? particularly, relating back to the problem setting and the limitations of existing methods.
 - how do the uncertainty estimates and OOD detection capabilities provided by the model compare to other approaches like ensembles or Bayesian neural networks?
 - given that the authors ground the problem setting and draw inspiration from human cognition and learning, do the authors feel their approach has biological plausibility at some level?

Overall, the discussed limitations hold it back from a higher score, however, the paper's originality and promising empirical results make it a interesting contribution. With a clearer exposition and better grounding/comparisons to other approaches, I would raise my score.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732500984359

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response to all reviewers

comment:
We thank the reviewers for their thoughtful comments with ways to improve the presentation of the paper. We were relieved that the reviewers found the method “novel” (reviewers vL85, v6Yp), “innovative” (reviewer v6Yp), and “simple and effective” (reviewer MWGp). Reviewers also found the approach to have a “clear and relevant objective” (reviewer 4ta1) and “with a solid foundation.. in well-established principles” (v6Yp). Further, reviewer vL85 identified a new connection: the method might improve “interpretability” of neural networks.

In response to the reviewers’ comments, we made several key improvements to the manuscript:

1. **We now clearly state the goal of each experiment**: The results section has been restructured to clearly map each experiment to the benefits of task abstractions, including faster learning, reduced forgetting, and improved generalization and OOD detection. Figure 2 updated with two more panels to concretely show these benefits.

2. **New Schematic (Figure 1)**: We replaced the initial figure with a detailed schematic illustrating the framework, including what is optimized during training and inference.

3. **A new discuss on the motivation for using gradients for inference**: The introduction now explicitly motivates the use of gradients for task inference. Using gradients avoids misalignment between task abstractions and the network parameters they control during learning. 

4. **Compare GBI to ensemble and Bayesian networks on OOD detection**: We ran additional experiments to compare OOD detection AUROC and implemented ensemble methods and Bayesian networks. GBI still shows better performance in comparison. 

6. **Code Improvements and Additional Experiments**: We are actively organizing the code base, and also running further experiments on OOD detection on CIFAR 10. We will update this forum with progress on those fronts. 

We hope our updates and discussions below address all comments raised and we are looking forward to participating in the discussion phase.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732892345393

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
Dear Reviewers,

We appreciate your thoughtful feedback and the recognition that the paper introduces a simple and novel method to build neural networks that are responsive to task abstractions. It is clear that there is general agreement about the conceptual contributions of this work. However, there are differing perspectives on where the most impactful extensions of this work should lie.

Suggestions included testing the framework on other architectures (eg, transformer), scaling to larger image datasets, and analyzing how generalization occurs in the current smaller networks. While we understand the value of each of these directions, we focused our efforts on extending the framework in ways we felt were most conceptually significant and aligned with our goals. Specifically, we demonstrated the method on:

1.	A toy yet intuitive sequence prediction task, exposing a link to Bayesian inference.
2.	An image generation task with a larger number of task abstractions.
3.	A language modeling task, which we viewed as a critical milestone due to the challenges posed by discrete tokens, which raised a serious question about their compatibility with gradient descent continuous dynamics.

We acknowledge that there may be differences in opinion about the most valuable paths forward, and we respect the diverse perspectives brought to this discussion.

We conclude by highlighting what seems to be a shared recognition among reviewers: the method is simple, novel, and represents a conceptual advance in training neural networks to be responsive to task abstractions. Crucially, while such networks can depend on abstractions provided by humans or another network, our method also allows them to infer abstractions directly from data when needed.

Thank you again for your time, effort, and thoughtful engagement with our work.

Sincerely,

The Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes to use Gradient-Based Inference (GBI) for learning abstract task representations, which allows neural networks to infer and adapt task representations dynamically. Inspired by human adaptability—where task abstractions allow flexible responses to the same input depending on internal goals—this approach enables neural networks to infer and adapt task representations on the fly. Experiments in a range of domains including image classification and language modelling demonstrate benefits in learning efficiency, generalisation, and reduced forgetting, as well as its performance in uncertainty estimation and out-of-distribution detection.

Strengths: The GBI method is novel and the problem it addresses is important. The connection to cognitive science and human cognition provides a solid grounding for this work.  The results show promising advantages in learning efficiency and generalization. The paper includes code which is appreciated and improves reproducibility.

Weaknesses:  The main concerns that remaining after the discussion relate to the experimental evaluation. The presented experiments are performed on simple datasets which serve as an intuitively appealing proof-of-concept, but do not provide strong empirical evidence for the advantages of the proposed method under realistic conditions and compared  to strong baselines.

This paper clearly has a lot of potential, but as it is, I consider it a borderline paper and lean towards rejection.