Compositional Risk Minimization

Divyat Mahajan, Mohammad Pezeshki, Charles Arnal, Ioannis Mitliagkas, Kartik Ahuja, Pascal Vincent

TL;DR: Provable method for extrapolating classifiers to novel combinations of attributes (a.k.a. compositional generalization)

Abstract:

Compositional generalization is a crucial step towards developing data-efficient intelligent machines that generalize in human-like ways.
In this work, we tackle a challenging form of distribution shift, termed compositional shift, where some attribute combinations are completely absent at training but present in the test distribution. This shift tests the model's ability to generalize compositionally to novel attribute combinations in discriminative tasks. We model the data with flexible additive energy distributions, where each energy term represents an attribute, and derive a simple alternative to empirical risk minimization termed compositional risk minimization (CRM). We first train an additive energy classifier to predict the multiple attributes and then adjust this classifier to tackle compositional shifts. We provide an extensive theoretical analysis of CRM, where we show that our proposal extrapolates to special affine hulls of seen attribute combinations. Empirical evaluations on benchmark datasets confirms the improved robustness of CRM compared to other methods from the literature designed to tackle various forms of subpopulation shifts.

Lay Summary: Modern AI systems often struggle to generalize to novel combinations of familiar attributes, a phenomenon known as compositional shift. We propose Compositional Risk Minimization (CRM), a method that separately models the influence of individual attributes on the data distribution and combines their effects additively. By enabling the model to reason about individual parts and how they work together, CRM supports more robust generalization to unseen attribute compositions. We provide both theoretical guarantees and empirical results to validate our approach. This work advances reliable, data-efficient learning and brings models closer to human-like generalization in open-world settings.

Link To Code: https://github.com/facebookresearch/compositional-risk-minimization

Primary Area: Deep Learning->Theory

Keywords: Compositional Generalization, Out of Distribution Generalization, Provable Extrapolation, Energy Based Models

Submission Number: 7418