Factorization Regret mediates compositional generalization in latent space

TL;DR

This paper introduces a framework using factorization regret to analyze and improve compositional generalization in latent space, demonstrated through a novel POMDP environment and architecture.

Contribution

It develops a variational inference framework and a new RCC architecture to enable learning variable interactions, enhancing generalization in latent space.

Findings

Factorization Regret explains accuracy gaps in RNNs.

RCCs improve compositional generalization and offline learning.

Identifies a failure mode where confidence decouples from accuracy.

Abstract

Are there still barriers to generalization once all of the relevant variables are known? We address this question via a framework that casts compositional generalization as a variational inference problem over latent variables with parametric interactions. To explore this framework, we develop the Cognitive Gridworld, a stationary Partially Observable Markov Decision Process (POMDP) in which observations are generated jointly by multiple latent variables, yet feedback is provided only for a single goal variable. This setting allows us to describe Factorization Regret: an information-theoretic quantity that measures the contribution of latent variable interactions to task performance. Using this metric, we first analyze Recurrent Neural Networks (RNNs) that are explicitly provided with the interactions and find that Factorization Regret explains the accuracy gap between Echo State and Fully Trained networks. Additionally, our analysis uncovers a theoretically predicted failure mode, where confidence becomes decoupled from accuracy. These results suggest that utilizing the interactions between relevant variables is a non-trivial capability. We then address a harder regime where the interactions themselves must be learned by an embedding model. Learning how variables interact while learning how to infer their values is a variational inference problem. We approach this dilemma via Representation Classification Chains (RCCs), a novel architecture which disentangles variable inference and parameter estimation. We demonstrate that, by learning how variables interact, RCCs facilitate compositional generalization to novel combinations of relevant variables and offline learning in novel action spaces. Together, these results establish a theoretically grounded setting for researching, developing and evaluating goal-directed generalist agents.