[Paper Review] Learning in Factored Domains with Information-Constrained Visual Representations

Tyler Malloy, Miao Liu+4 more|arXiv (Cornell University)|Mar 30, 2023

Cell Image Analysis TechniquesBiochemistry, Genetics and Molecular Biology3 citations

TL;DR

This paper proposes a factored reinforcement learning model using a modified $eta$-VAE to learn disentangled, information-constrained visual representations that enable fast, sample-efficient learning in visual contextual bandit tasks. By generating and evaluating hypotheses about reward structure from compressed latent features, the model achieves rapid convergence—selecting the optimal action after just two experiences—demonstrating a trade-off between learning speed and reconstruction accuracy.

ABSTRACT

Humans learn quickly even in tasks that contain complex visual information. This is due in part to the efficient formation of compressed representations of visual information, allowing for better generalization and robustness. However, compressed representations alone are insufficient for explaining the high speed of human learning. Reinforcement learning (RL) models that seek to replicate this impressive efficiency may do so through the use of factored representations of tasks. These informationally simplistic representations of tasks are similarly motivated as the use of compressed representations of visual information. Recent studies have connected biological visual perception to disentangled and compressed representations. This raises the question of how humans learn to efficiently represent visual information in a manner useful for learning tasks. In this paper we present a model of human factored representation learning based on an altered form of a $β$-Variational Auto-encoder used in a visual learning task. Modelling results demonstrate a trade-off in the informational complexity of model latent dimension spaces, between the speed of learning and the accuracy of reconstructions.

Motivation & Objective

• To investigate how compressed, disentangled visual representations support rapid learning in complex visual tasks.
• To model human-like sample efficiency in reinforcement learning by leveraging factored representations of state and reward.
• To examine the trade-off between reconstruction accuracy and learning speed in information-constrained latent spaces.
• To evaluate whether disentangled representations can support effective hypothesis generation for reward prediction.
• To explore a behaviorally grounded definition of disentanglement based on utility and task performance.

Proposed method

• A modified $eta$-VAE is trained to reconstruct visual inputs while also predicting rewards from latent features, with a joint loss function balancing reconstruction, KL regularization, and reward prediction accuracy.
• The model uses a factored representation of the reward function, decomposing it into independent components based on latent features corresponding to distinct visual attributes (e.g., glasses, hats).
• Hypotheses about future rewards and state transitions are generated and iteratively re-evaluated using the latent representation, enabling efficient exploration and learning.
• The model is pre-trained on a large dataset of celebrity faces (220K images) before fine-tuning in a contextual bandit setting with four stimulus categories (glasses, hats, both, neither).
• Reward prediction is modeled as a linear sum of individual attribute contributions, with temporal credit assignment via a value function over next-state latent features.
• The hypothesis generation mechanism assumes deterministic, additive reward structures, allowing for efficient learning in a low-sample regime.

Experimental results

Research questions

• RQ1 How do information-constrained visual representations affect the speed and accuracy of reinforcement learning in visual tasks?
• RQ2 Can disentangled latent representations support fast hypothesis generation about reward structure in a factored MDP setting?
• RQ3 What is the trade-off between reconstruction fidelity and learning efficiency in compressed latent spaces?
• RQ4 How do latent representations evolve to reflect utility-based equivalence in visual stimuli?
• RQ5 Can a behaviorally defined notion of disentanglement explain faster learning in human-like agents?

Key findings

• Models with smaller latent dimensionality achieved faster learning, selecting the higher-reward bandit arm after only two experiences, demonstrating high sample efficiency.
• Smaller latent spaces showed lower reconstruction accuracy due to a tighter information bottleneck, confirming the trade-off between speed and fidelity.
• Latent representations of high-utility stimuli (e.g., both glasses and hats) became increasingly distinct from zero-utility stimuli over time, indicating utility-based representation learning.
• The model's latent space exhibited a form of utility-based equivalence, where stimuli with similar rewards developed similar representations.
• The hypothesis generation mechanism enabled rapid convergence by focusing on salient, disentangled features linked to reward outcomes.
• Pre-training on a large face dataset enabled effective transfer to the contextual bandit task, supporting generalization from compressed visual features.