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.

What problem does this paper attempt to address?

The problem that this paper attempts to solve is: how to enable neural networks to achieve more efficient learning, better generalization ability and flexible task adaptability through task - abstract representations like the human brain. Specifically, the paper explores how to effectively infer these abstract representations without explicitly providing task abstractions, and how to recombine these abstractions to adapt to new problems. To achieve this goal, the authors propose the Gradient - Based Inference (GBI) method, which is an optimization method from the perspective of variational inference, capable of efficiently inferring current task requirements in neural networks and recombining abstractions to adapt to new situations by iteratively optimizing the task - representation layer. ### Main Contributions 1. **Proposing the GBI method**: The GBI method utilizes the back - propagation of gradients to the task - representation layer to efficiently infer current task requirements and recombine task abstractions through further iterative optimization to adapt to new tasks. 2. **Verifying the advantages of GBI**: Through multiple experiments (including toy datasets, image classification and language models), the advantages of GBI in improving learning efficiency, generalization ability and reducing forgetting are demonstrated. 3. **Unique advantages**: GBI not only performs well in learning and generalization, but also has the ability to retain uncertainty estimation information and detect abnormal samples. ### Experimental Results 1. **Toy Dataset Experiments**: - **Data Efficiency**: GBI - LSTM learns tasks more quickly than the traditional LSTM. - **Generalization Ability**: GBI - LSTM performs better on data points outside the training range and does not experience catastrophic forgetting. - **Bayesian Properties**: GBI - LSTM can be used to approximate Bayesian quantities, such as posterior probabilities and likelihood functions. 2. **Image Classification Experiments**: - **Performance Comparison**: GBI performs excellently in image classification tasks, outperforming other gradient - based inference methods, and can achieve high accuracy with only one forward propagation when the computational budget is limited. - **Uncertainty Estimation and OOD Detection**: GBI shows obvious advantages in uncertainty and abnormal sample detection, especially after data normalization, its performance is better than the existing state - of - the - art methods. 3. **Language Model Experiments**: - **Domain Generality**: GBI also performs well in language modeling tasks, demonstrating its potential as a general - purpose method. - **Task Inference and Data Efficiency**: GBI also shows good task - inference ability and data efficiency in language tasks. ### Conclusion By proposing the GBI method, the paper successfully addresses the challenges of neural networks in task - abstract representation, making them approach or even exceed human cognitive abilities in terms of learning efficiency, generalization ability and task adaptability. The GBI method not only performs well in multiple tasks, but also has unique advantages, such as uncertainty estimation and abnormal sample detection, providing a new direction for future neural network research.

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

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

Neural representation of abstract task structure during generalization

Reward-predictive representations generalize across tasks in reinforcement learning

Emergent neural dynamics and geometry for generalization in a transitive inference task

A Generalizable Approach to Learning Optimizers

Gradient Aligned Domain Generalization with a Mutual Teaching Teacher-student Network for Intelligent Fault Diagnosis

Gradient Adversarial Training of Neural Networks

Representations and generalization in artificial and brain neural networks

Gradient-adjusted Incremental Target Propagation Provides Effective Credit Assignment in Deep Neural Networks

Compositional generalization through abstract representations in human and artificial neural networks

Human-Like Geometric Abstraction in Large Pre-trained Neural Networks

Instilling Inductive Biases with Subnetworks

Fast Task-Aware Architecture Inference

Improved Techniques for Training Adaptive Deep Networks

Evolutionary Training and Abstraction Yields Algorithmic Generalization of Neural Computers

Neuralizer: General Neuroimage Analysis without Re-Training

Neural Prioritisation of Past Solutions Supports Generalisation

Do What Nature Did to Us: Evolving Plastic Recurrent Neural Networks for Task Generalization

Inductive Gradient Adjustment For Spectral Bias In Implicit Neural Representations

Disentangling Trainability and Generalization in Deep Neural Networks