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.

Learning Task Informed Abstractions

Task-Induced Representation Learning

Abstract Reasoning with Distracting Features

Learning Markov State Abstractions for Deep Reinforcement Learning

Don’t Touch What Matters: Task-Aware Lipschitz Data Augmentation for Visual Reinforcement Learning

Temporal Prediction Model with Context-Aware Data Augmentation for Robust Visual Reinforcement Learning

Don't Touch What Matters: Task-Aware Lipschitz Data Augmentation for Visual Reinforcement Learning

Visual Imitation Learning of Task-Oriented Object Grasping and Rearrangement

Stabilizing Visual Reinforcement Learning Via Asymmetric Interactive Cooperation

Improving Target-driven Visual Navigation with Attention on 3D Spatial Relationships

Learning with Language-Guided State Abstractions

Biologically-Motivated Learning Model for Instructed Visual Processing

Multi-task Hierarchical Adversarial Inverse Reinforcement Learning

Theia: Distilling Diverse Vision Foundation Models for Robot Learning

Following Instructions by Imagining and Reaching Visual Goals

Learning task-agnostic representation via toddler-inspired learning

Make the Pertinent Salient: Task-Relevant Reconstruction for Visual Control with Distractions

Robust Visual Imitation Learning with Inverse Dynamics Representations

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

Learning to Discover Task-Relevant Features for Interpretable Reinforcement Learning

Task Aware Dreamer for Task Generalization in Reinforcement Learning