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

arXiv: https://arxiv.org/abs/2407.17356 · HTML: https://mlsanity.com/paper/2407.17356/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 28 of its 59 references that are in the map.
- First posted on arXiv 24 Jul 2024.

## Builds on (most cited first, 28 of 28)

- [On Calibration of Modern Neural Networks](https://mlsanity.com/paper/1706.04599/) (2017)
- [Overcoming catastrophic forgetting in neural networks](https://mlsanity.com/paper/1612.00796/) (2016)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms](https://mlsanity.com/paper/1708.07747/) (2017)
- [Parameter-Efficient Transfer Learning for NLP](https://mlsanity.com/paper/1902.00751/) (2019)
- [Auto-Encoding Variational Bayes](https://mlsanity.com/paper/1312.6114/) (2013)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [Prefix-Tuning: Optimizing Continuous Prompts for Generation](https://mlsanity.com/paper/2101.00190/) (2021)
- [The Power of Scale for Parameter-Efficient Prompt Tuning](https://mlsanity.com/paper/2104.08691/) (2021)
- [A Baseline for Detecting Misclassified and Out-of-Distribution Examples in Neural Networks](https://mlsanity.com/paper/1610.02136/) (2016)
- [A Review of Uncertainty Quantification in Deep Learning: Techniques, Applications and Challenges](https://mlsanity.com/paper/2011.06225/) (2020)
- [Do Deep Generative Models Know What They Don't Know?](https://mlsanity.com/paper/1810.09136/) (2018)
- [Neural Module Networks](https://mlsanity.com/paper/1511.02799/) (2015)
- [Supermasks in Superposition](https://mlsanity.com/paper/2006.14769/) (2020)
- [Recurrent Independent Mechanisms](https://mlsanity.com/paper/1909.10893/) (2019)
- [Likelihood Regret: An Out-of-Distribution Detection Score For Variational Auto-encoder](https://mlsanity.com/paper/2003.02977/) (2020)
- [The Goldilocks Principle: Reading Children's Books with Explicit Memory Representations](https://mlsanity.com/paper/1511.02301/) (2015)
- [Alleviating catastrophic forgetting using context-dependent gating and synaptic stabilization](https://mlsanity.com/paper/1802.01569/) (2018)
- [Iterative Amortized Inference](https://mlsanity.com/paper/1807.09356/) (2018)
- [Modular Networks: Learning to Decompose Neural Computation](https://mlsanity.com/paper/1811.05249/) (2018)
- [Call for Papers -- The BabyLM Challenge: Sample-efficient pretraining on a developmentally plausible corpus](https://mlsanity.com/paper/2301.11796/) (2023)
- [Choreographer: Learning and Adapting Skills in Imagination](https://mlsanity.com/paper/2211.13350/) (2022)
- [Understanding the Limitations of Conditional Generative Models](https://mlsanity.com/paper/1906.01171/) (2019)
- [Optimizing Reusable Knowledge for Continual Learning via Metalearning](https://mlsanity.com/paper/2106.05390/) (2021)
- [Complexity of Word Collocation Networks: A Preliminary Structural Analysis](https://mlsanity.com/paper/1310.5111/) (2013)
- [Learning, Planning, and Control in a Monolithic Neural Event Inference Architecture](https://mlsanity.com/paper/1809.07412/) (2018)
- [Thalamus: a brain-inspired algorithm for biologically-plausible continual learning and disentangled representations](https://mlsanity.com/paper/2205.11713/) (2022)
- [Efficient Expansion and Gradient Based Task Inference for Replay Free Incremental Learning](https://mlsanity.com/paper/2312.01188/) (2023)

## Built on by (most cited first, 4 of 4)

- [Do different prompting methods yield a common task representation in language models?](https://mlsanity.com/paper/2505.12075/) (2025)
- [Meta-Dynamical State Space Models for Integrative Neural Data Analysis](https://mlsanity.com/paper/2410.05454/) (2024)
- [Flexible task abstractions emerge in linear networks with fast and bounded units](https://mlsanity.com/paper/2411.03840/) (2024)
- [Predicting Causal Effects from Natural Language Queries using Structured Representations](https://mlsanity.com/paper/2605.29631/) (2026)

## Abstract (arXiv)

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.