NEURAL NETWORKS FOR LEARNING COUNTERFAC-TUAL G-INVARIANCES FROM SINGLE ENVIRONMENTS

Abstract

Despite -or maybe because of-their astonishing capacity to fit data, neural networks are believed to have difficulties extrapolating beyond training data distribution. This work shows that, for extrapolations based on finite transformation groups, a model's inability to extrapolate is unrelated to its capacity. Rather, the shortcoming is inherited from a learning hypothesis: Examples not explicitly observed with infinitely many training examples have underspecified outcomes in the learner's model. In order to endow neural networks with the ability to extrapolate over group transformations, we introduce a learning framework counterfactually-guided by the learning hypothesis that any group invariance to (known) transformation groups is mandatory even without evidence, unless the learner deems it inconsistent with the training data. Unlike existing invariance-driven methods for (counterfactual) extrapolations, this framework allows extrapolations from a single environment. Finally, we introduce sequence and image extrapolation tasks that validate our framework and showcase the shortcomings of traditional approaches.

1. INTRODUCTION

Neural networks are widely praised for their ability to interpolate the training data. However, in some applications, they have also been shown to be unable to learn patterns that can provably extrapolate out-of-distribution (beyond the training data distribution) (Arjovsky et al., 2019; D'Amour et al., 2020; de Haan et al., 2019; Geirhos et al., 2020; McCoy et al., 2019; Schölkopf, 2019) . Recent counterfactual-based learning frameworks for extrapolation tasks -such as ICM and IRM (Arjovsky et al., 2019; Besserve et al., 2018; Johansson et al., 2016; Louizos et al., 2017; Peters et al., 2017; Schölkopf, 2019; Krueger et al., 2020) detailed in Section 2-assume the learner is given data from multiple environmental conditions (say environments E1 and E2) and is expected to learn patterns that work well over an unseen environment E3. In particular, the key idea behind IRM is to force the neural network to learn an internal representation of the input data that is invariant to environmental changes between E1 and E2, and, hence, hopefully also invariant to E3, which may not be true for nonlinear classifiers (Rosenfeld et al., 2020) . While successful for a class of extrapolation tasks, these frameworks require multiple environments in the training data. But, are we asking the impossible? Can humans even perform single-environment extrapolation? Young children, unlike monkeys and baboons, assume that a conditional stimulus F given another stimulus D extrapolates to a symmetric relation D given F without ever seeing any such examples (Sidman et al., 1982) . E.g., if given D, action F produces a treat, the child assumes that given F, action D also produces a treat. Young children differ from primates in their ability to use symmetries to build conceptual relations beyond visual patterns (Sidman and Tailby, 1982; Westphal-Fitch et al., 2012) , allowing extrapolations from intelligent reasoning. However, forcing symmetries against data evidence is undesirable, since symmetries can provide valuable evidence when they are broken. Unfortunately, single-environment extrapolations have not been addressed in the literature. The challenge comes from a learning framework where examples not explicitly observed with infinitely many independent training examples are underspecified in the learner's statistical model, which is shared by both objective (frequentist) and subjective (Bayesian) learner's frameworks. For instance, consider a supervised learning task where the training data contains infinitely many sequences x (tr) =(A,B) associated with label y (tr) = C, but no examples of a sequence x (tr) =(B,A). If given a test example x (te) =(B,A), the hypothesis considers it to be out of distribution and the prediction P (Y (te) = C|X (te) = (B,A)) is undefined, since P (X (tr) = (B,A)) = 0. This happens regardless of a prior over P (X (tr) ). This unseen-is-underspecified learning hypothesis is not guaranteed to push neural networks to assume symmetric extrapolations without evidence. Contributions. Since symmetries are intrinsically tied to human single-environment extrapolation capabilities, this work explores a learning framework that modifies the learner's hypothesis space to allow symmetric extrapolation (over known groups) without evidence, while not losing valuable antisymmetric information if observed to predict the target variable in the training data. Formally, a symmetry is an invariance to transformations of a group, known as a G-invariance. In Theorem 1 we show that the counterfactual invariances needed for symmetry extrapolation -denoted Counterfactual G-invariances (CG-invariances)-are stronger than traditional G-invariances. Theorem 2, then, introduces a condition in the structural causal model where G-invariances of linear automorphism groups are safe to use as CG-invariances. With that, Theorem 3 defines a partial order over the appropriate invariant subspaces that we use to learn the correct G-invariances from a single environment without evidence, while retaining the ability to be sensitive to antisymmetries shown to be relevant in the training data. Finally, we introduce sequence and image counterfactual extrapolation tasks with experiments that validate the theoretical results and showcase the advantages of our approach.

2. RELATED WORK

Counterfactual inference and invariances. Recent efforts have brought counterfactual inference to machine learning models. Independent causal mechanism (ICM) and Invariant Risk Minimization (IRM) methods (Arjovsky et al., 2019; Besserve et al., 2018; Johansson et al., 2016; Parascandolo et al., 2018; Schölkopf, 2019) , Causal Discovery from Change (CDC) methods (Tian and Pearl, 2001) , and representation disentanglement methods (Bengio et al., 2020; Goudet et al., 2017) broadly look for representations, classifiers, or mechanism descriptions, that are invariant across multiple environments observed in the training data or inferred from the training data (Creager et al., 2020) . They rely on multiple environment samples in order to reason over new environments. To the best of our knowledge there is no clear effort for extrapolations from a single environment. The key similarity between the ICM framework and our framework is the assumption of independently sampled mechanisms (the transformations) and causes. Domain adaptation and domain generalization. Domain adaptation and domain generalization (e.g. (Long et al., 2017; Muandet et al., 2013; Quionero-Candela et al., 2009; Rojas-Carulla et al., 2018; Shimodaira, 2000; Zhang et al., 2015) and others) ask questions about specific -observed or known-changes in the data distribution rather than counterfactual questions. A key difference is that counterfactual inference accounts for hypothetical interventions, not known ones. 2020) considers learning image invariances from the training data, however does not consider extrapolation tasks. Moreover, it does not provide a concrete theoretical proof of invariance, relying on experimental results over interpolation tasks for validation. Another parallel work (Zhou et al., 2021) uses metalearning to learn symmetries that are shared across several tasks (or environments). The works of van der Wilk et al. (2018) and Anselmi et al. (2019) focus on learning invariances from training data for better generalization error of the training distribution. However, none of these works consider the extrapolation task. In contrast, our framework formally considers counterfactual extrapolation, for which we provide both theoretical and experimental results.



Forced G-invariances. Forcing a G-invariance may contradict the training data, where the target variable is actually influenced by the transformation of the input. For instance, handwritten digits are not invariant to 180 o rotations, since digits 6 and 9 would get confused. Data augmentation is a type of forced G-invariance(Chen et al., 2020; Lyle et al., 2020)  and hence, will fail to extrapolate. Other works forcing G-invariances that will also fail include (not an extensive list): Zaheer et al. (2017) and Murphy et al. (2019a;b) for permutation groups over set and graph inputs; Cohen and Welling (2016), Cohen et al. (2019) for dihedral and spherical transformation groups over images. Learning invariances from training data. The parallel work of Benton et al. (

