Factorization Regret mediates compositional generalization in latent space

Abstract

Are there still barriers to generalization once all of the relevant variables are known? We consider the challenge of generalizing to a novel combination of task-relevant latent variables. 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 with parametric interactions. 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. Estimating how variables interact while simultaneously inferring their values is a variational inference problem. To explicitly disentangle variable inference from parameter estimation, we develop Representation Classification Chains (RCCs), a novel architecture for learning how latent variables interact. RCCs are capable of learning from sparse and partial reward, provided only for a subset of goal variables. Finally, we demonstrate the usefulness of RCCs in enabling generalization to novel combinations of latent variables through offline learning. In summary, we present a theoretically grounded setting for research, development and evaluation of goal-directed general intelligence.