Original: L EARNING EXPLANATIONS THAT ARE HARD TO VARY
Revised: L EARNING EXPLANATIONS THAT ARE HARD TO VARY

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: In this paper, we investigate the principle that good explanations are hard to vary in the context of deep learning.
Revised: In this paper, we investigate the principle that good explanations are hard to vary in the context of deep learning.

Original: We show that averaging gradients across examples – akin to a logical OR ( _ ) of patterns – can favor memorization and ‘patchwork’ solutions that sew together different strategies, instead of identifying invariances.
Revised: We show that averaging gradients across examples – akin to a logical OR ( _ ) of patterns – can favor memorization and ‘patchwork’ solutions that sew together different strategies, instead of identifying invariances.

Original: To inspect this, we ﬁrst formalize a notion of consistency for minima of the loss surface, which measures to what extent a minimum appears only when examples are pooled.
Revised: To inspect this, we ﬁrst formalize a notion of consistency for minima of the loss surface, which measures to what extent a minimum appears only when examples are pooled.

Original: We then propose and experimentally validate a simple alternative algorithm based on a logical AND ( ^ ), that focuses on invariances and prevents memorization in a set of real-world tasks.
Revised: We then propose and experimentally validate a simple alternative algorithm based on a logical AND ( ^ ), that focuses on invariances and prevents memorization in a set of real-world tasks.

Original: Finally, using a synthetic dataset with a clear distinction between invariant and spurious mechanisms, we dissect learning signals and compare this approach to well-established regularizers.
Revised: Finally, using a synthetic dataset with a clear distinction between invariant and spurious mechanisms, we dissect learning signals and compare this approach to well-established regularizers [Table]

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: +

Original: [Equation]

Original: [Table]
Revised: [Table]

Original: /M/rpDS49jMZJSlhJBaLBqliFLPZ86wvNQpSE0u40NLeysSIay7IRlS0IXjLL6+SZrXiuRXv/rJcu8njKMApnMEFeHAFNbiDOjRAgIJneIU359F5cd6dj0XrmpPPnMAfOJ8/zuuPzQ==</latexit>
Revised: /M/rpDS49jMZJSlhJBaLBqliFLPZ86wvNQpSE0u40NLeysSIay7IRlS0IXjLL6+SZrXiuRXv/rJcu8njKMApnMEFeHAFNbiDOjRAgIJneIU359F5cd6dj0XrmpPPnMAfOJ8/zuuPzQ==</latexit>

Original: [Figure]
Revised: [Figure]

Original: Consider the top of Figure 1, which shows a view from above of the loss surface obtained as we vary a two dimensional parameter vector θ “ p θ 1 , θ 2 q , for a ﬁctional dataset containing two observations x A and x B .
Revised: Consider the top of Figure 1, which shows a view from above of the loss surface obtained as we vary a two dimensional parameter vector θ “ p θ 1 , θ 2 q , for a ﬁctional dataset containing two observations x A and x B .

Original: Note the two global minima at coordinates p θ 1 , θ 2 q “ p´ 0 . 5 , ´ 0 . 5 q and p 0 . 8 q .
Revised: Note the two global minima on the top-right

Original: Depending on the initial values of θ — marked as white circles — gradient descent converges to one of the two minima.
Revised: Depending on the initial values of θ — marked as white circles — gradient descent converges to one of the two minima.

Original: Judging solely by the value of the loss function, which is zero in both cases, the two minima look equally good.
Revised: Judging solely by the value of the loss function, which is zero in both cases, the two minima look equally good.

Original: However, looking at the loss surfaces for x A and x B separately, as shown below, a crucial difference between those two minima appears: Starting from the same initial parameter conﬁgurations and following the gradient of the loss, ∇ θ L p θ, x i q , the probability of ﬁnding the same p 0 . 8 q -minimum in either case is zero.
Revised: However, looking at the loss surfaces for x A and x B separately, as shown below, a crucial difference between those two minima appears: Starting from the same initial parameter conﬁgurations and following the gradient of the loss, ∇ θ L p θ, x i q , the probability of ﬁnding the same minimum on the top-right in either case is zero.

Original: In contrast, the minimum in the lower-left corner has a signiﬁcant overlap across the two loss surfaces, so gradient descent can converge to it even if training on x A (or x B ) only.
Revised: In contrast, the minimum in the lower-left corner has a signiﬁcant overlap across the two loss surfaces, so gradient descent can converge to it even if training on x A (or x B ) only.

Original: Note that after averaging there is no way to tell what the two loss surfaces looked like: Are we destroying information that is potentially important?
Revised: Note that after averaging there is no way to tell what the two loss surfaces looked like: Are we destroying information that is potentially important?

Original: In this paper, we argue that the answer is yes.
Revised: In this paper, we argue that the answer is yes.

Original: In particular, we hypothesize that if the goal is to ﬁnd invariant mechanisms in the data, these can be identiﬁed by ﬁnding explanations (e.g. model parameters) that are hard to vary across examples.
Revised: In particular, we hypothesize that if the goal is to ﬁnd invariant mechanisms in the data, these can be identiﬁed by ﬁnding explanations (e.g. model parameters) that are hard to vary across examples.

Original: A notion of invariance implies something that stays the same, as something else changes.
Revised: A notion of invariance implies something that stays the same, as something else changes.

Original: We assume that data comes from different environments :
Revised: We assume that data comes from different environments :

Original: An invariant mechanism is shared across all, generalizes out of distribution (o.o.d.), but might be hard to model; each environment also has spurious explanations that are easy to spot (‘shortcuts’), but do not generalize o.o.d.
Revised: An invariant mechanism is shared across all, generalizes out of distribution (o.o.d.), but might be hard to model; each environment also has spurious explanations that are easy to spot (‘shortcuts’), but do not generalize o.o.d.

Original: We formalize a notion of consistency , which characterizes to what extent a minimum of the loss surface appears only when data from different environments are pooled.
Revised: We formalize a notion of consistency , which characterizes to what extent a minimum of the loss surface appears only when data from different environments are pooled.

Original: Minima with low consistency are ‘patchwork’ solutions, which (we hypothesize) sew together different strategies and should not be expected to generalize to new environments.
Revised: Minima with low consistency are ‘patchwork’ solutions, which (we hypothesize) sew together different strategies and should not be expected to generalize to new environments.

Original: An intuitive description of this principle was proposed by physicist David Deutsch: “good explanations are hard to vary” (Deutsch, 2011).
Revised: An intuitive description of this principle was proposed by physicist David Deutsch: “good explanations are hard to vary” (Deutsch, 2011).

Original: Using the notion of consistency, we deﬁne Invariant Learning Consistency (ILC), a measure of the expected consistency of the solution found by a learning algorithm on a given hypothesis class.
Revised: Using the notion of consistency, we deﬁne Invariant Learning Consistency (ILC), a measure of the expected consistency of the solution found by a learning algorithm on a given hypothesis class.

Original: The
Revised: The

Original: ILC can be improved by changing the hypothesis class or the learning algorithm, and in the last part of the paper we focus on the latter.
Revised: ILC can be improved by changing the hypothesis class or the learning algorithm, and in the last part of the paper we focus on the latter.

Original: We then analyse why current practices in deep learning provide little incentive for networks to learn invariances, and show that standard training is instead set up with the explicit objective of greedily maximizing speed of learning, i.e., progress on the training loss.
Revised: We then analyse why current practices in deep learning provide little incentive for networks to learn invariances, and show that standard training is instead set up with the explicit objective of greedily maximizing speed of learning, i.e., progress on the training loss.

Original: When learning “as fast as possible” is not the main objective, we show we can trade-off some “learning speed” for prioritizing learning the invariances.
Revised: When learning “as fast as possible” is not the main objective, we show we can trade-off some “learning speed” for prioritizing learning the invariances.

Original: A practical instantiation of ILC leads to o.o.d.
Revised: A practical instantiation of ILC leads to o.o.d.

Original: generalization on a challenging synthetic task where several established regularizers fail to generalize; moreover, following the memorization task from Zhang et al.
Revised: generalization on a challenging synthetic task where several established regularizers fail to generalize; moreover, following the memorization task from Zhang et al.

Original: , ILC prevents convergence on CIFAR-10 with random labels, as no shared mechanism is present, and similarly when a portion of training labels is incorrect.
Revised: , ILC prevents convergence on CIFAR-10 with random labels, as no shared mechanism is present, and similarly when a portion of training labels is incorrect.

Original: Lastly, we set up a behavioural cloning task based on the game CoinRun (Cobbe et al., 2019b), and observe better generalization on new unseen levels.
Revised: Lastly, we set up a behavioural cloning task based on the game CoinRun (Cobbe et al., 2019b), and observe better generalization on new unseen levels.

Original: [Figure]
Revised: [Figure]

Original: [Equation]
Revised: [Equation]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: An example.
Revised: An example.

Original: Take these two second-hand books of chess puzzles.
Revised: Take these two second-hand books of chess puzzles.

Original: We can learn the two independent shortcuts (red arrows for the left book OR hand-written solutions on the right), or actually learn to play chess (the invariant mechanism).
Revised: We can learn the two independent shortcuts (red arrows for the left book OR hand-written solutions on the right), or actually learn to play chess (the invariant mechanism).

Original: While both strategies solve other problems from the same books (i.i.d.), only the latter generalises to new chess puzzle books (o.o.d.).
Revised: While both strategies solve other problems from the same books (i.i.d.), only the latter generalises to new chess puzzle books (o.o.d.).

Original: How to distinguish the two?
Revised: How to distinguish the two?

Original: We would not have learned about the red arrows had we trained on the book on the right, and vice versa with the hand-written notes.
Revised: We would not have learned about the red arrows had we trained on the book on the right, and vice versa with the hand-written notes.

Original:  E XPLANATIONS THAT ARE HARD TO VARY
Revised:  E XPLANATIONS THAT ARE HARD TO VARY

Original: We consider datasets t D e u e P E , with | E | “ d , and D e “ p x ei , y ei q , i e “ 1 , . . . , n e .
Revised: We consider datasets t D e u e P E , with | E | “ d , and D e “ p x ei , y ei q , i e “ 1 , . . . , n e .

Original: Here x ei P X Ď R m is the vector containing the observed inputs, and y ei P Y Ď R p the targets.
Revised: Here x ei P X Ď R m is the vector containing the observed inputs, and y ei P Y Ď R p the targets.

Original: The superscript e P E indexes some aspect of the data collection process, and can be interpreted as an environment label.
Revised: The superscript e P E indexes some aspect of the data collection process, and can be interpreted as an environment label.

Original: Our objective is to infer a function f : X Ñ Y — which we call mechanism — assigning a target y ei to each input x ei ; as explained in the introduction, we assume that such function is shared across all environments.
Revised: Our objective is to infer a function f : X Ñ Y — which we call mechanism — assigning a target y ei to each input x ei ; as explained in the introduction, we assume that such function is shared across all environments.

Original: For estimation purposes, f may be parametrized by a neural network with continuous activations; for weights θ P Θ Ď R n , we denote the neural network output at x P X as f θ p x q .
Revised: For estimation purposes, f may be parametrized by a neural network with continuous activations; for weights θ P Θ Ď R n , we denote the neural network output at x P X as f θ p x q .

Original: Gradient-based optimization.
Revised: Gradient-based optimization.

Original: To ﬁnd an appropriate model f θ , standard optimizers rely on gradients from a pooled loss function L : R n Ñ R .
Revised: To ﬁnd an appropriate model f θ , standard optimizers rely on gradients from a pooled loss function L : R n Ñ R .

Original: This function measures the average performance of the neural network when predicting data labels, across all environments: L p θ q : “ 1 | E | ř e P E
Revised: This function measures the average performance of the neural network when predicting data labels, across all environments: L p θ q : “ 1 | E | ř e P E

Original: L
Revised: L

Original: e
Revised: e

Original: p θ q , with L e p θ q : “ 1 | D e | ř p x ei ,y ei qP D e (cid:96)
Revised: p θ q , with L e p θ q : “ 1 | D e | ř p x ei ,y ei qP D e (cid:96)

Original: p
Revised: p

Original: f p x ei ; θ q , y ei q ; where (cid:96) :
Revised: f p x ei ; θ q , y ei q ; where (cid:96) :

Original: R p
Revised: R p

Original: ˆ R p Ñ r 0 , `8q is usually chosen to be the L 2 loss or the cross-entropy loss.
Revised: ˆ R p Ñ r 0 , `8q is usually chosen to be the L 2 loss or the cross-entropy loss.

Original: The parameter updates according to gradient descent (GD) are given by θ k ` 1 GD “ θ k GD ´
Revised: The parameter updates according to gradient descent (GD) are given by θ k ` 1 GD “ θ k GD ´

Original: η ∇L p θ k GD q , where η ą 0 is the learning rate.
Revised: η ∇L p θ k GD q , where η ą 0 is the learning rate.

Original: Under some standard assumptions (Lee et al., 2016), p θ k GD q k ě 0 converges to a local minimizer of L , with probability one.
Revised: Under some standard assumptions (Lee et al., 2016), p θ k GD q k ě 0 converges to a local minimizer of L , with probability one.

Original: When do we not learn invariances?
Revised: When do we not learn invariances?

Original: We start by describing what might prevent learning invariances in standard gradient-based optimization.
Revised: We start by describing what might prevent learning invariances in standard gradient-based optimization.

Original: (i) Training stops once the loss is low enough.
Revised: (i) Training stops once the loss is low enough.

Original: If optimization learned spurious patterns by the time it converged, invariances will not be learned anymore.
Revised: If optimization learned spurious patterns by the time it converged, invariances will not be learned anymore.

Original: This depends on the rate at which different patterns are learned.
Revised: This depends on the rate at which different patterns are learned.

Original: The rates at which invariant patterns emerge (and vice-versa, the spurious patterns do not) can be improved by e.g.: (a) careful architecture design, e
Revised: The rates at which invariant patterns emerge (and vice-versa, the spurious patterns do not) can be improved by e.g.: (a) careful architecture design, e

Original: g. as done by hardcoding spatial equivariance in convolutional networks; (b) ﬁne-tuning models pre-trained on large amounts of data, where strong features already emerged and can be readily selected.
Revised: g. as done by hardcoding spatial equivariance in convolutional networks; (b) ﬁne-tuning models pre-trained on large amounts of data, where strong features already emerged and can be readily selected.

Original: (ii) Learning signals: everything looks relevant for a dataset of size 1.
Revised: (ii) Learning signals: everything looks relevant for a dataset of size 1.

Original: Due to the summation in the deﬁnition of the pooled loss L , gradients for each example are computed independently.
Revised: Due to the summation in the deﬁnition of the pooled loss L , gradients for each example are computed independently.

Original: Informally, each signal is identical to the one for an equivalent dataset of size 1, where every pattern appears relevant to the task.
Revised: Informally, each signal is identical to the one for an equivalent dataset of size 1, where every pattern appears relevant to the task.

Original: To ﬁnd invariant patterns across examples, if we compute our training signals on each of them independently, we have to rely on the way these are aggregated.
Revised: To ﬁnd invariant patterns across examples, if we compute our training signals on each of them independently, we have to rely on the way these are aggregated.

Original: (iii) Aggregating gradients: averaging maximizes learning speed.
Revised: (iii) Aggregating gradients: averaging maximizes learning speed.

Original: The default method to pool gradients is the arithmetic mean .
Revised: The default method to pool gradients is the arithmetic mean .

Original: GD applied to L is designed to minimize the pooled loss by prioritizing descent speed .
Revised: GD applied to L is designed to minimize the pooled loss by prioritizing descent speed .

Original: Indeed, a step of GD is equivalent to ﬁnding a tight 3 quadratic upper bound ˆ L to L , and then jumping to the minimizer of this approximation (Nocedal and Wright, 2006).
Revised: Indeed, a step of GD is equivalent to ﬁnding a tight 3 quadratic upper bound ˆ L to L , and then jumping to the minimizer of this approximation (Nocedal and Wright, 2006).

Original: While speed is often desirable, by construction GD ignores one potentially crucial piece of information: The gradient ∇L is the result of averaging signals ∇L e , which correspond to the patterns visible from each environment at this stage of optimization.
Revised: While speed is often desirable, by construction GD ignores one potentially crucial piece of information: The gradient ∇L is the result of averaging signals ∇L e , which correspond to the patterns visible from each environment at this stage of optimization.

Original: In other words, GD with average gradients greedily maximizes for learning speed, but in some situations we would like to trade some convergence speed for invariance.
Revised: In other words, GD with average gradients greedily maximizes for learning speed, but in some situations we would like to trade some convergence speed for invariance.

Original: [Figure]
Revised: [Figure]

Original: For instance, instead of performing an arithmetic mean between gradients (logical OR), we might want to look towards a logical AND, which can be characterized as a geometric mean.
Revised: For instance, instead of performing an arithmetic mean between gradients (logical OR), we might want to look towards a logical AND, which can be characterized as a geometric mean.

Original: Fig.
Revised: Fig.

Original: 1 shows how a sum can be seen as a logical OR: the two orthogonal gradients from data A and data B at (0.5,0.5) point to different directions, yet both are kept in the combined gradient.
Revised: 1 shows how a sum can be seen as a logical OR: the two orthogonal gradients from data A and data B at (0.5,0.5) point to different directions, yet both are kept in the combined gradient.

Original: In Sec.
Revised: In Sec.

Original: 2.3 we elaborate on this idea and on implementing a logical AND between gradients.
Revised: 2.3 we elaborate on this idea and on implementing a logical AND between gradients.

Original: Before presenting this discussion, we take some time to better motivate the need for invariant learning consistency and to construct a precise mathematical deﬁnition of consistency.
Revised: Before presenting this discussion, we take some time to better motivate the need for invariant learning consistency and to construct a precise mathematical deﬁnition of consistency.

Original: 2. F ORMAL DEFINITION OF ILC
Revised: 2. F ORMAL DEFINITION OF ILC

Original: Let Θ ˚ A be the set of convergence points of algorithm A when trained using all environments (pooled data): that is, Θ ˚ A “ t θ ˚ P Θ | D θ 0 P R n s.t.
Revised: Let Θ ˚ A be the set of convergence points of algorithm A when trained using all environments (pooled data): that is, Θ ˚ A “ t θ ˚ P Θ | D θ 0 P R n s.t.

Original: A 8 p θ 0 , E q “ θ ˚ u .
Revised: A 8 p θ 0 , E q “ θ ˚ u .

Original: For instance, if A is gradient descent, the result of Lee et al.
Revised: For instance, if A is gradient descent, the result of Lee et al.

Original: (2016) implies that Θ ˚ A is the set of local minimizers of the pooled loss L .
Revised: (2016) implies that Θ ˚ A is the set of local minimizers of the pooled loss L .

Original: To each θ ˚ P Θ ˚ A , we want to associate a consistency score, quantifying the concept “ good θ ˚ are hard to vary ”.
Revised: To each θ ˚ P Θ ˚ A , we want to associate a consistency score, quantifying the concept “ good θ ˚ are hard to vary ”.

Original: In other words
Revised: In other words

Original: [Table]
Revised: [Table]

Original: we would like the score to capture the consistency of the loss landscape around θ ˚ across the different environments.
Revised: we would like the score to capture the consistency of the loss landscape around θ ˚ across the different environments.

Original: For example, in Fig.
Revised: For example, in Fig.

Original: 1 the loss landscape near the bottom-left minimizer is consistent across environments, while the top-right minimizer is not.
Revised: 1 the loss landscape near the bottom-left minimizer is consistent across environments, while the top-right minimizer is not.

Original: [Table]
Revised: [Table]

Original: MbUmuXWQ==</latexit> Let us characterize the landscape around θ ˚ from the per- spective of a ﬁxed environment e P E .
Revised: MbUmuXWQ==</latexit> Let us characterize the landscape around θ ˚ from the per- spective of a ﬁxed environment e P E .

Original: We deﬁne the set
Revised: We deﬁne the set

Original: N (cid:15)e,θ ˚ to be the largest path-connected region of space containing both θ ˚ and the set t θ P Θ s.t.
Revised: N (cid:15)e,θ ˚ to be the largest path-connected region of space containing both θ ˚ and the set t θ P Θ s.t.

Original: | L e p θ q ´ L e p θ ˚ q| ď (cid:15) u , with (cid:15) ą 0 .
Revised: | L e p θ q ´ L e p θ ˚ q| ď (cid:15) u , with (cid:15) ą 0 .

Original: In other words, if θ P N (cid:15)e,θ ˚ then there exist a path-connected region in parameter space including θ ˚ and θ where each parameter also is in N (cid:15)e,θ ˚ and its loss on environment e is comparable.
Revised: In other words, if θ P N (cid:15)e,θ ˚ then there exist a path-connected region in parameter space including θ ˚ and θ where each parameter also is in N (cid:15)e,θ ˚ and its loss on environment e is comparable.

Original: From the perspective of environment e , all these points are equivalent to θ ˚ .
Revised: From the perspective of environment e , all these points are equivalent to θ ˚ .

Original: We would like to evaluate the elements of this set with respect to a different environment e 1 ‰ e .
Revised: We would like to evaluate the elements of this set with respect to a different environment e 1 ‰ e .

Original: We will say that e 1 is consistent with e in θ ˚ if max θ P N (cid:15)e,θ ˚ |
Revised: We will say that e 1 is consistent with e in θ ˚ if max θ P N (cid:15)e,θ ˚ |

Original: L e 1 p θ q ´ L e p θ q ˇˇ is small.
Revised: L e 1 p θ q ´ L e p θ q ˇˇ is small.

Original: Repeating this reasoning for all environment pairs, we arrive at the following inconsistency score :
Revised: Repeating this reasoning for all environment pairs, we arrive at the following inconsistency score :

Original: [Equation]
Revised: [Equation]

Original: This consistency is our formalization of the principle “ good explanations are hard to vary ”.
Revised: This consistency is our formalization of the principle “ good explanations are hard to vary ”.

Original: Finally, we can write down an invariant learning consistency score for A :
Revised: Finally, we can write down an invariant learning consistency score for A :

Original: [Equation]
Revised: [Equation]

Original: That is, the learning consistency of an algorithm measures the expected consistency across environments of the minimizer it converges to on the pooled data.
Revised: That is, the learning consistency of an algorithm measures the expected consistency across environments of the minimizer it converges to on the pooled data.

Original: Example: low consistency of a classic patchwork solution.
Revised: Example: low consistency of a classic patchwork solution.

Original: One-hidden-layer networks with sigmoid activations and enough neurons can approximate any function f ˚ : r 0 , 1 s Ñ R (Cybenko, 1989).
Revised: One-hidden-layer networks with sigmoid activations and enough neurons can approximate any function f ˚ : r 0 , 1 s Ñ R (Cybenko, 1989).

Original: In appendix A.1 we show how the construction used to obtain the weights leads to a maximally inconsistent solution according to I (cid:15) p θ ˚ q , which would not be expected to generalize o.o.d.
Revised: In appendix A.1 we show how the construction used to obtain the weights leads to a maximally inconsistent solution according to I (cid:15) p θ ˚ q , which would not be expected to generalize o.o.d.

Original: 2. ILC AS A LOGICAL AND BETWEEN LANDSCAPES
Revised: 2. ILC AS A LOGICAL AND BETWEEN LANDSCAPES

Original: H A , H B H A
Revised: H A , H B H A

Original: + B
Revised: + B

Original: H A ∧ B
Revised: H A ∧ B

Original: Here we draw a connection between our deﬁnition of inconsistency and the local geometric properties of the loss landscapes.
Revised: Here we draw a connection between our deﬁnition of inconsistency and the local geometric properties of the loss landscapes.

Original: For the sake of clarity, we consider two environments ( A and B ) and assume θ ˚ to be a local minimizer (with zero loss) for both environments.
Revised: For the sake of clarity, we consider two environments ( A and B ) and assume θ ˚ to be a local minimizer (with zero loss) for both environments.

Original: Using a
Revised: Using a

Original: Taylor approximation 5 , we get L p θ q « 12 p θ ´ θ ˚ q J H
Revised: Taylor approximation 5 , we get L p θ q « 12 p θ ´ θ ˚ q J H

Original: A ` B p θ ´ θ ˚ q for } θ ´ θ ˚ } « 0 , where H A ` B “ p H A ` H B q { 2 is the arithmetic mean of the Hessians
Revised: A ` B p θ ´ θ ˚ q for } θ ´ θ ˚ } « 0 , where H A ` B “ p H A ` H B q { 2 is the arithmetic mean of the Hessians

Original: H A : “ ∇ 2 L A p θ ˚ q and H B : “ ∇ 2 L A p θ ˚ q .
Revised: H A : “ ∇ 2 L A p θ ˚ q and H B : “ ∇ 2 L A p θ ˚ q .

Original: H A ` B does not capture the possibly conﬂicting geometries of landscape A or B : It performs a “logical OR” on the dominant eigendirections.
Revised: H A ` B does not capture the possibly conﬂicting geometries of landscape A or B : It performs a “logical OR” on the dominant eigendirections.

Original: In contrast, the geometric mean , or Karcher mean, H A ^ B (Ando et al., 2004) is affected by the inconsistencies between landscapes: It performs a “logical AND”.
Revised: In contrast, the geometric mean , or Karcher mean, H A ^ B (Ando et al., 2004) is affected by the inconsistencies between landscapes: It performs a “logical AND”.

Original: In appendix A.2, we give a formal deﬁnition of H A ^ B , and show that for diagonal Hessians, I (cid:15) p θ ˚ q ď 2 (cid:15) p det p H A
Revised: In appendix A.2, we give a formal deﬁnition of H A ^ B , and show that for diagonal Hessians, I (cid:15) p θ ˚ q ď 2 (cid:15) p det p H A

Original: ` B q det p H A ^ B q q 2 .
Revised: ` B q det p H A ^ B q q 2 .

Original: As for the geometric mean of positive numbers, 0 ď det p H A ^ B q ď det p H A ` B q ; thus, inconsistency is lowest when shapes of A and B are similar – exactly as in the bottom-left minimizer of Fig. 1.
Revised: As for the geometric mean of positive numbers, 0 ď det p H A ^ B q ď det p H A ` B q ; thus, inconsistency is lowest when shapes of A and B are similar – exactly as in the bottom-left minimizer of Fig. 1.

Original: From Hessians to gradients.
Revised: From Hessians to gradients.

Original: We just saw that the consistency of θ ˚ is linked to the geometric mean of the Hessians t H e p θ ˚ qu e P E .
Revised: We just saw that the consistency of θ ˚ is linked to the geometric mean of the Hessians t H e p θ ˚ qu e P E .

Original: Under the simplifying assumption that each H e is diagonal 6 and all eigenvalues λ ei are positive, their geometric mean is
Revised: Under the simplifying assumption that each H e is diagonal 6 and all eigenvalues λ ei are positive, their geometric mean is

Original: [Equation]
Revised: [Equation]

Original: The curvature of the corresponding loss in the i -th eigendirection depends on how consistent the curvatures of each environment are in that direction.
Revised: The curvature of the corresponding loss in the i -th eigendirection depends on how consistent the curvatures of each environment are in that direction.

Original: Consider now optimizing from a point θ k ; gradient descent reads θ
Revised: Consider now optimizing from a point θ k ; gradient descent reads θ

Original: [Equation]
Revised: [Equation]

Original: , where H
Revised: , where H

Original: [Equation]
Revised: [Equation]

Original: e P E λ en q .
Revised: e P E λ en q .

Original: For η small enough 7 , we have | θ k ` 

Revised: For η small enough 7 , we have | θ k `

Original: i
Revised: i

Original: [Equation]
Revised: [Equation]

Original: As noted, this choice maximises the speed of convergence to θ ˚ , but does not take into account whether this minimizer is consistent.
Revised: As noted, this choice maximises the speed of convergence to θ ˚ , but does not take into account whether this minimizer is consistent.

Original: We can reduce the speed of convergence on directions where landscapes have different curvatures – which would lead to a high inconsistency – by following the gradients from the geometric mean of the landscapes, as opposed to the arithmetic mean.
Revised: We can reduce the speed of convergence on directions where landscapes have different curvatures – which would lead to a high inconsistency – by following the gradients from the geometric mean of the landscapes, as opposed to the arithmetic mean.

Original: I.e, we substitute the full gradient ∇L p θ q “ H ` p θ k ´ θ ˚ q with ∇L ^ p θ q “ H ^ p θ k ´ θ ˚ q .
Revised: I.e, we substitute the full gradient ∇L p θ q “ H ` p θ k ´ θ ˚ q with ∇L ^ p θ q “ H ^ p θ k ´ θ ˚ q .

Original: We also have  8 ∇L ^ p θ q “ p ś e P E ∇L e p θ qq 1 { | E | : to reduce the speed of convergence in directions with inconsistency, we can take the element-wise geometric mean of gradients from different environments (see also Fig.
Revised: Also, we have that 8 ∇L ^ p θ q “ p ś e P E [Equation] e p θ qq 1 { | E | : to reduce the speed of convergence in directions with inconsistency, we can take the element-wise geometric mean of gradients from different environments (see also Fig.

Original: 11 in  appendix).
Revised: 11 in the appendix).

Original:  M ASKING GRADIENTS WITH A LOGICAL AND
Revised: 2. M ASKING GRADIENTS WITH A LOGICAL AND

Original: The element-wise geometric mean of gradients, instead of the arithmetic mean, increases consistency in the convex quadratic case.
Revised: The element-wise geometric mean of gradients, instead of the arithmetic mean, increases consistency in the convex quadratic case.

Original: This alternative gradient descent has a few practical limitations:
Revised: However, there are a few practical limitations:

Original: (i)
Revised: (i)

Original: The geometric mean is only deﬁned when all the signs are consistent.
Revised: The geometric mean is only deﬁned when all the signs are consistent.

Original: It is  to be deﬁned how sign inconsistencies, which can occur in non-convex settings, should be dealt with.
Revised: It is still to be deﬁned how sign inconsistencies, which can occur in non-convex settings, should be dealt with.

Original: (ii)
Revised: (ii)

Original: It provides little ﬂexibility for ‘partial’ agreement: Even a single zero gradient component in one environment stops optimization in that direction.
Revised: It provides little ﬂexibility for ‘partial’ agreement: Even a single zero gradient component in one environment stops optimization in that direction.

Original: (iii) For numerical stability, it needs to be computed in log domain (computationally expensive).
Revised: (iii) For numerical stability, it needs to be computed in log domain (more computationally expensive).

Original: (iv) Adaptive step-size schemes (e
Revised: (iv) Adaptive step-size schemes (e

Original: g. Adam (Kingma and Ba, 2014)) rescale the signal componentwise for local curvature adaptation.
Revised: g. Adam (Kingma and Ba, 2014)) rescale the signal componentwise for local curvature adaptation.

Original: The exact magnitude of the geometric mean would be ignored and most of the difference from arithmetic averaging will come from the zero-ed components
Revised: The exact magnitude of the geometric mean would be ignored and most of the difference from arithmetic averaging will come from the zero-ed components.

Original: (i) can be overcome by treating different signs as zeros, resulting in a geometric mean of 0 if there is any sign disagreement across environments for a gradient component.
Revised: (i) can be overcome by treating different signs as zeros, resulting in a geometric mean of 0 if there is any sign disagreement across environments for a gradient component.

Original: For (ii) we can allow for some disagreement (with a hyperparameter), by not masking out if there is a large percentage of environments with gradients in that direction.
Revised: For (ii) we can allow for some disagreement (with a hyperparameter), by not masking out if there is a large percentage of environments with gradients in that direction.

Original: (iii) and (iv) can be addressed together: Since the ﬁnal magnitude will be rescaled except for masked components, i.e. where the geometric mean is 0, we can use the average gradients (fast to compute) and mask out the components based on the sign agreement (computable avoiding the log domain).
Revised: (iii) and (iv) can be addressed together: Since the ﬁnal magnitude will be rescaled except for masked components, i.e. where the geometric mean is 0, we can use the average gradients (fast to compute) and mask out the components based on the sign agreement (computable avoiding the log domain).

Original: The AND-mask.
Revised: The AND-mask.

Original: We translate the reasoning we just presented to a practical algorithm that we will refer to as the AND-mask .
Revised: We translate the reasoning we just presented to a practical algorithm that we will refer to as the AND-mask .

Original: In its most simple implementation, we zero out those gradient components with respect to weights that have inconsistent signs across environments.
Revised: In its most simple implementation, we zero out those gradient components with respect to weights that have inconsistent signs across environments.

Original: Formally, the masked gradients at iteration k are m t p θ k qd ∇L p θ k q , where m t p θ k q vanishes for any component where there are less than t P t d { 2 , d { 2 ` 1 , . . .
Revised: Formally, the masked gradients at iteration k are m t p θ k qd ∇L p θ k q , where m t p θ k q vanishes for any component where there are less than t P t d { 2 , d { 2 ` 1 , . . .

Original: , d u agreeing gradient signs across environments ( d is the number of environments in the batch), and is equal to one otherwise.
Revised: , d u agreeing gradient signs across environments ( d is the number of environments in the batch), and is equal to one otherwise.

Original: For convenience, our implementation of the AND-mask uses a threshold τ P r 0 , 1 s as hyper-parameter instead of t , such that t “ d 2 p τ ` 1 q .
Revised: For convenience, our implementation of the AND-mask uses a threshold τ P r 0 , 1 s as hyper-parameter instead of t , such that t “ d 2 p τ ` 1 q .

Original: Mathematically, for every component r m τ s j of m τ , r m τ s j “ 1 r τd ď |
Revised: Mathematically, for every component r m τ s j of m τ , r m τ s j “ 1 r τd ď |

Original: [Equation]
Revised: [Equation]

Original: Computing the AND-mask has the same time and space complexity of standard gradient descent, i.e., linear in the number of examples that we average.
Revised: Computing the AND-mask has the same time and space complexity of standard gradient descent, i.e., linear in the number of examples that we average.

Original: Due to its simplicity and computational efﬁciency, this is the algorithm that we will use in the experiment section.
Revised: Due to its simplicity and computational efﬁciency, this is the algorithm that we will use in the experiment section.

Original: As a ﬁrst result, we show that following the AND-masked gradient leads to convergence in the directions made visible by the
Revised: As a ﬁrst result, we show that following the AND-masked gradient leads to convergence in the directions made visible by the

Original: AND-mask.
Revised: AND-mask.

Original: The proof is presented in appendix A.3.
Revised: The proof is presented in appendix A.3.

Original: Proposition 1 .
Revised: Proposition 1 .

Original: Let L have L -Lipschitz gradients and consider a learning rate η ď 1 { L .
Revised: Let L have L -Lipschitz gradients and consider a learning rate η ď 1 { L .

Original: After k iterations, AND-masked GD visits at least once a point θ where } m t p θ q d ∇L p θ q} 

Revised: After k iterations, AND-masked GD visits at least once a point θ where } m t p θ q d ∇L p θ q}

Original: ď O p 1 { k q .
Revised: ď O p 1 { k q .

Original: [Figure]
Revised: [Figure]

Original: Behaviour in the face of randomness.
Revised: Behaviour in the face of randomness.

Original: Here we put the AND mask through a theoretical test:
Revised: Here we put the AND mask through a theoretical test:

Original: For gradients coming from different environments that are inconsistent (or even random), how fast does the AND mask reduce the magnitude of the step taken in parameter space, compared to standard GD?
Revised: For gradients coming from different environments that are inconsistent (or even random), how fast does the AND mask reduce the magnitude of the step taken in parameter space, compared to standard GD?

Original: In case of inconsistency, the AND mask should quickly make the gradient steps more conservative.
Revised: In case of inconsistency, the AND mask should quickly make the gradient steps more conservative.

Original: To assess this property, we consider a ﬁxed set of n parameters θ and gradients ∇L e drawn independently from a multivariate Gaussian with zero mean and unit covariance.
Revised: To assess this property, we consider a ﬁxed set of n parameters θ and gradients ∇L e drawn independently from a multivariate Gaussian with zero mean and unit covariance.

Original: Proposition 2 .
Revised: Proposition 2 .

Original: Consider the setting we just outlined, with L “ p 1 { d q ř de “ 1 L e .
Revised: Consider the setting we just outlined, with L “ p 1 { d q ř de “ 1 L e .

Original: While E } ∇L p θ q} 2 “
Revised: While E } ∇L p θ q} 2 “

Original: O p n { d q , we have that @
Revised: O p n { d q , we have that @

Original: t P t d { 2 ` 1 , . . .
Revised: t P t d { 2 ` 1 , . . .

Original: , d u , D c
Revised: , d u , D c

Original: P p 1 , 2 s such that E } m t p θ q d ∇L p θ q} 2 ď
Revised: P p 1 , 2 s such that E } m t p θ q d ∇L p θ q} 2 ď

Original: O p n { c d q .
Revised: O p n { c d q .

Original: The proof is presented in appendix A.4, and an illustration with numerical veriﬁcation in Fig.
Revised: The proof is presented in Appendix A.4, and an illustration with numerical veriﬁcation in Fig.

Original: 4 (the magnitudes of masked gradients ( • ) for more than 100 examples were always zero in the numerical veriﬁcation).
Revised: 4 (the magnitudes of masked gradients (•) for more than 100 examples were always zero in the numerical veriﬁcation).

Original: Intuitively, in the presence of purely random patterns, the AND-mask has a desirable property: it decreases the strength of these signals exponentially fast, as opposed to linearly.
Revised: Intuitively, in the presence of purely random patterns, the AND-mask has a desirable property: it decreases the strength of these signals exponentially fast, as opposed to linearly.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: Real-world datasets are generated by (causal) generative processes which share mechanisms (Pearl, 2009).
Revised: Real-world datasets are generated by (causal) generative processes which share mechanisms (Pearl, 2009).

Original: However, mechanisms and spurious signals are often entangled, making it hard to assess what part of the learning signal is due to either.
Revised: However, mechanisms and spurious signals are often entangled, making it hard to assess what part of the learning signal is due to either.

Original: As the goal of this paper is to dissect these two components to understand how they ultimately contribute to the learning process, we create a simple synthetic dataset that allows us to control the complexity, intensity, and number of shortcuts in the data.
Revised: As the goal of this paper is to dissect these two components to understand how they ultimately contribute to the learning process, we create a simple synthetic dataset that allows us to control the complexity, intensity, and number of shortcuts in the data.

Original: After that, we evaluate whether spurious signals can be detected even in high-dimensional networks and datasets by testing the AND-mask on a memorization task similar to the one proposed in Zhang et al.
Revised: After that, we evaluate whether spurious signals can be detected even in high-dimensional networks and datasets by testing the AND-mask on a memorization task similar to the one proposed in Zhang et al.

Original: (2016), and on a behavioral cloning task using the game CoinRun (Cobbe et al., 2018).
Revised: (2016), and on a behavioral cloning task using the game CoinRun (Cobbe et al., 2018).

Original: 3. T HE SYNTHETIC MEMORIZATION DATASET
Revised: 3. T HE SYNTHETIC MEMORIZATION DATASET

Original: We introduce a binary classiﬁcation task.
Revised: We introduce a binary classiﬁcation task.

Original: The input dimensionality is d “ d M ` d S .
Revised: The input dimensionality is d “ d M ` d S .

Original:  p p y | x d M q is the same across all environments (i.e. the mechanism ). p p y | x d S , e q is not the same across all environments (the shortcuts ).
Revised: While p p y | x d M q is the same across all environments (i.e. the mechanism ), p p y | x d S , e q is not the same across all environments (the shortcuts ).

Original: While the mechanism is shared, it needs a highly non-linear decision boundary to classify the data.
Revised: While the mechanism is shared, it needs a highly non-linear decision boundary to classify the data.

Original: The shortcuts are not shared across environments, but provide a simple way to classify the data, even when pooling all the environments together.
Revised: The shortcuts are not shared across environments, but provide a simple way to classify the data, even when pooling all the environments together.

Original: See Figure 5 for a concrete example with d M and d S equal to 2, and two environments ( A and B ).
Revised: See Figure 5 for a concrete example with d M and d S equal to 2, and two environments ( A and B ).

Original: The spirals (on d M ) are invariant but hard to model.
Revised: The spirals (on d M ) are invariant but hard to model.

Original: The shortcuts (on d S ) are simple blobs but different in every environment: in A , linearly separable through a vertical decision boundary, in B with a horizontal one.
Revised: The shortcuts (on d S ) are simple blobs but different in every environment: in A , linearly separable through a vertical decision boundary, in B with a horizontal one.

Original: If the two environments are pooled, a new diagonal decision boundary emerges on the shortcut dimensions as the most ‘natural’ one.
Revised: If the two environments are pooled, a new diagonal decision boundary emerges on the shortcut dimensions as the most ‘natural’ one.

Original: While this perfectly classiﬁes data in both environments A and B , critically it would have not been found by training on either partition A or B alone .
Revised: While this perfectly classiﬁes data in both environments A and B , critically it would have not been found by training on either partition A or B alone .

Original: The out-of-distribution (o.o.d.)
Revised: The out-of-distribution (o.o.d.)

Original: test data has the same mechanism but random shortcuts.
Revised: test data has the same mechanism but random shortcuts.

Original: Therefore, any method relying
Revised: Therefore, any method relying

Original: [Figure]
Revised: [Table]

Revised: +tmJrr0cy7TzKBk80VRJohJyPRvEnKFzIixJZQpbm8lbEgVZcamU7EheIsvL5PHs7rn1r2781rjuoijDEdwDKfgwQU04Aaa0AIGA3iGV3hzhPPivDsf89aSU8wcwh84nz8ZKY2t</latexit> performance.

Original: Details about the , and training curves are reported in appendix B
Revised: Details about the dataset, baselines, and training curves are reported in appendix B

Original: Despite the apparent simplicity of this dataset, note that it is challenging to ﬁnd the invariant mechanism.
Revised: Despite the apparent simplicity of this dataset, note that it is challenging to ﬁnd the invariant mechanism.

Original: In high dimensions, even with tens of pooled environments, the shortcuts allow for a simple classiﬁcation rule under almost every classical deﬁnition of ‘simple’: the boundary is linear , it has a large margin , it can be expressed with small weights , it is fast to learn , robust to input noise, and has perfect accuracy and no i.i.d.
Revised: In high dimensions, even with tens of pooled environments, the shortcuts allow for a simple classiﬁcation rule under almost every classical deﬁnition of ‘simple’: the boundary is linear , it has a large margin , it can be expressed with small weights , it is fast to learn , robust to input noise, and has perfect accuracy and no i.i.d.

Original: generalization gap .
Revised: generalization gap .

Original: Finding the complex decision boundary of the spirals, instead, is a ﬁddly process and arguably a much slower path towards small loss.
Revised: Finding the complex decision boundary of the spirals, instead, is a ﬁddly process and arguably a much slower path towards small loss.

Original: Baselines.
Revised: Baselines.

Original: We evaluate several domain-agnostic baselines (all multilayer perceptrons) with some of the most common regularizers used in deep learning — Dropout, L1, L2, Batch normalization.
Revised: We evaluate several domain-agnostic baselines (all multilayer perceptrons) with some of the most common regularizers used in deep learning — Dropout, L1, L2, Batch normalization.

Original: We also consider methods that explicitly make use of the environment labels, namely: (i) Domain
Revised: We also consider methods that explicitly make use of the environment labels, namely: (i) Domain

Original: Adversarial Neural Networks (DANN) (Ganin et al., 2016), a method speciﬁcally designed to address domain adaptation by obfuscating domain information with an adversarial classiﬁer; (ii) Invariant
Revised: Adversarial Neural Networks (DANN) (Ganin et al., 2016), a method speciﬁcally designed to address domain adaptation by obfuscating domain information with an adversarial classiﬁer; (ii) Invariant

Original: Risk Minimization (IRM) (Arjovsky et al., 2019), discussed in detail in appendix B. The AND-mask is trained with the same conﬁgurations in Table 1.
Revised: Risk Minimization (IRM) (Arjovsky et al., 2019), discussed in detail in appendix B. The AND-mask is trained with the same conﬁgurations in Table 1.

Original: [Figure]
Revised: [Figure]

Original: Results.
Revised: Results.

Original: Fig.
Revised: Fig.

Original: 6 shows training and test accuracy.
Revised: 6 shows training and test accuracy.

Original: DANN fails because it can align the representation-layer distributions from different environments using only shortcuts, such that they become indistinguishable to the domaindiscriminating classiﬁer.
Revised: DANN fails because it can align the representation-layer distributions from different environments using only shortcuts, such that they become indistinguishable to the domaindiscriminating classiﬁer.

Original: The AND-mask was the only method to achieve perfect test accuracy, by ﬁtting the spirals instead of the shortcuts.
Revised: The AND-mask was the only method to achieve perfect test accuracy, by ﬁtting the spirals instead of the shortcuts.

Original: In particular, the combination of the AND-mask with L1 or L2 regularization gave the most robust results overall, as they help suppress neurons that at initialization are tuned towards the shortcuts.
Revised: In particular, the combination of the AND-mask with L1 or L2 regularization gave the most robust results overall, as they help suppress neurons that at initialization are tuned towards the shortcuts.

Original: Correlations between average, memorization and generalization gradients.
Revised: Correlations between average, memorization and generalization gradients.

Original: We can intervene on the data-generating processof the dataset in order to examine the learning signal coming from the mechanisms and from the shortcuts.
Revised: Due to the synthetic nature of the dataset, we can intervene on its data-generating process in order to examine the learning signals coming from the mechanisms and from the shortcuts.

Original: We isolate the two and measure their contribution to the average gradients for random initializations.
Revised: We isolate the two and measure their contribution to the average gradients, as we vary the agreement threshold of the mask.

Revised: More precisely, we look at the gradients computed with respect to the weights of a randomly initialized network for different sets of data: (i) The original data, with mechanisms and shortcuts.

Original:  Figure 7 shows the correlation between the components of the original average gradient  and the shortcut gradients (dashed line), and between the original average gradients and the mechanism gradients (solid line).While the signal from the mechanisms is present in the original average gradients (i.e. ρ « 0 . 4 for τ “ 0 ), its magnitude is smaller and it is ‘drowned’ by the memorization signal.
Revised: (ii) Randomly permuting the dataset over the mechanisms dimensions, thus leaving the “memorization” signal of the shortcuts. (iii) Randomly permuting over the shortcuts dimensions, isolating the “generalization” signal of the mechanisms alone. Figure 7 shows the correlation between the components of the original average gradient (i) and the shortcut gradients ((ii), dashed line), and between the original average gradients and the mechanism gradients ((iii), solid line).

Original: [Figure]
Revised: [Figure]

Revised: While the signal from the mechanisms is present in the original average gradients (i.e. ρ « 0 . 4 for τ “ 0 ), its magnitude is smaller and it is ‘drowned’ by the memorization signal.

Original: Instead, increasing the threshold of the AND-mask (right side) suppresses memorization gradients due to the shortcuts, and for τ « 1 most of the gradient components remaining contain signal from the mechanism.
Revised: Instead, increasing the threshold of the AND-mask (right side) suppresses memorization gradients due to the shortcuts, and for τ « 1 most of the gradient components remaining contain signal from the mechanism.

Original: On the left side, we test the other side of our hypothesis: An XORmask zeroes out consistent gradients, preserves those with different signs, and results in a sharper decrease of the correlation with the mechanism gradients.
Revised: On the left side, we test the other side of our hypothesis: An XORmask zeroes out consistent gradients, preserves those with different signs, and results in a sharper decrease of the correlation with the mechanism gradients.

Original: 3. E XPERIMENTS ON CIFAR- 
Revised: 3. E XPERIMENTS ON CIFAR-

Original: Memorization in a vision task.
Revised: Memorization in a vision task.

Original: Zhang et al.
Revised: Zhang et al.

Original: (2016) showed that neural networks trained with standard regularizers — like L2 and Dropout — can still memorize large training datasets with shufﬂed labels, i.e. reaching « 100% training accuracy.
Revised: (2016) showed that neural networks trained with standard regularizers — like L2 and Dropout — can still memorize large training datasets with shufﬂed labels, i.e. reaching « 100% training accuracy.

Original: Their experiments raised signiﬁcant questions about the generalization properties of neural networks and the role of regularizers in constraining the hypothesis class.
Revised: Their experiments raised signiﬁcant questions about the generalization properties of neural networks and the role of regularizers in constraining the hypothesis class.

Original: Our hypothesis is that ILC — for example implemented as the AND-mask — should prevent memorization on a similar task with the shufﬂed labels, as gradients will tend to largely ‘disagree’ in the absence of a shared mechanism.
Revised: Our hypothesis is that ILC — for example implemented as the AND-mask — should prevent memorization on a similar task with the shufﬂed labels, as gradients will tend to largely ‘disagree’ in the absence of a shared mechanism.

Original: However, when the labels are not shufﬂed ,
Revised: However, when the labels are not shufﬂed ,

Original: ILC should have a much weaker effect, as real shared mechanisms are still present in the data.
Revised: ILC should have a much weaker effect, as real shared mechanisms are still present in the data.

Original: To test our hypothesis, we ran an experiment that closely resembles the one in (Zhang et al., 2016) on
Revised: To test our hypothesis, we ran an experiment that closely resembles the one in (Zhang et al., 2016) on

Original: CIFAR-10.
Revised: CIFAR-10.

Original: We trained a ResNeton CIFAR-10 with random labels , with and without the AND-mask.
Revised: We trained a ResNet on CIFAR-10 with random labels , with and without the AND-mask.

Original: In all experiments we used batch size 80, and treated each example as its own “environment”.
Revised: In all experiments we used batch size 80, and treated each example as its own “environment”.

Original: Recall that standard gradient averaging is equivalent to an AND-mask with threshold 0.
Revised: Recall that standard gradient averaging is equivalent to an AND-mask with threshold 0.

Original: As shown in Figure 8, the ResNet with standard average gradients memorized the data, while slightly increasing the threshold for the AND-mask quickly prevented memorization (dark blue line).
Revised: As shown in Figure 8, the ResNet with standard average gradients memorized the data, while slightly increasing the threshold for the AND-mask quickly prevented memorization (dark blue line).

Original: In contrast, training the same networks on the dataset with the original labels resulted in both of them converging and generalizing to the test set, conﬁrming that the mask did not signiﬁcantly affect the generalization error with a general underlying mechanism in the data.
Revised: In contrast, training the same networks on the dataset with the original labels resulted in both of them converging and generalizing to the test set, conﬁrming that the mask did not signiﬁcantly affect the generalization error with a general underlying mechanism in the data.

Original: [Figure]
Revised: [Figure]

Original: Note that there is no standard notion of environments in CIFAR10, which is why we treated every example as coming from its own environment.
Revised: Note that there is no standard notion of environments in CIFAR10, which is why we treated every example as coming from its own environment.

Original: This assumption is not unreasonable, as every image in the dataset was literally collected in a different physical environment.
Revised: This assumption is not unreasonable, as every image in the dataset was literally collected in a different physical environment.

Original: If anything, it is the standard i.i.d.
Revised: If anything, it is the standard i.i.d.

Original: assumption that hides this variety behind a notion of a single distribution encompassing all environments.
Revised: assumption that hides this variety behind a notion of a single distribution encompassing all environments.

Original: The results of this experiment further support this interpretation, and can serve as evidence that — in some cases — we might be able to identify invariances even without an explicit partition into environments, as this can be already identiﬁed at the level of individual examples.
Revised: The results of this experiment further support this interpretation, and can serve as evidence that — in some cases — we might be able to identify invariances even without an explicit partition into environments, as this can be already identiﬁed at the level of individual examples.

Original: [Figure]
Revised: [Figure]

Original: Label noise.
Revised: Label noise.

Original: Following up on this experiment, we test how the AND-mask performs in the presence of label noise, i.e. when a portion of the labels in the training set are randomly shufﬂed (25% here).
Revised: Following up on this experiment, we test how the AND-mask performs in the presence of label noise, i.e. when a portion of the labels in the training set are randomly shufﬂed (25% here).

Original: According to our hypothesis, gradients computed on examples with random labels should disagree and get masked out by the AND-mask, while signal from correctly labeled data should contribute to update the model.
Revised: According to our hypothesis, gradients computed on examples with random labels should disagree and get masked out by the AND-mask, while signal from correctly labeled data should contribute to update the model.

Original: As shown in Figure 9, the performance on the incorrectly labeled portion of the dataset is well below chance for the AND-mask (as it predicts correctly despite the wrong labels), while the baseline again memorizes the incorrect labels.
Revised: As shown in Figure 9, the performance on the incorrectly labeled portion of the dataset is well below chance for the AND-mask (as it predicts correctly despite the wrong labels), while the baseline again memorizes the incorrect labels.

Original: On the test set (with untouched labels), the baseline peaks early then decreases as the model overﬁts, while the AND-mask slowly but steadily improves.
Revised: On the test set (with untouched labels), the baseline peaks early then decreases as the model overﬁts, while the AND-mask slowly but steadily improves.

Original: 3. B EHAVIORAL C LONING ON C OIN R UN
Revised: 3. B EHAVIORAL C LONING ON C OIN R UN

Original: CoinRun (Cobbe et al., 2019b) is a game introduced to test how RL agents generalize to novel situations.
Revised: CoinRun (Cobbe et al., 2019b) is a game introduced to test how RL agents generalize to novel situations.

Original: The agent needs to collect coins, jumping on top of walls and boxes and avoiding enemies.
Revised: The agent needs to collect coins, jumping on top of walls and boxes and avoiding enemies.

Original: 9 Each level is procedurally generated — i.e. it has a different combination of sprites, background, and layout — but the physics and goals are invariant.
Revised: 9 Each level is procedurally generated — i.e. it has a different combination of sprites, background, and layout — but the physics and goals are invariant.

Original: Cobbe et al.
Revised: Cobbe et al.

Original: (2019b) showed that state-of-the-art RL algorithms fail to model these invariant mechanisms, performing poorly on new levels unless trained on thousands of them.
Revised: (2019b) showed that state-of-the-art RL algorithms fail to model these invariant mechanisms, performing poorly on new levels unless trained on thousands of them.

Original: To test our hypothesis, we set up a behavioral cloning task using CoinRun.
Revised: To test our hypothesis, we set up a behavioral cloning task using CoinRun.

Original: 10 We start by pre-training a strong policy π ˚ using standard PPO (Schulman et al., 2017) for 400M steps on the full distribution of levels.
Revised: 10 We start by pre-training a strong policy π ˚ using standard PPO (Schulman et al., 2017) for 400M steps on the full distribution of levels.

Original: We then generate a dataset of pairs p s, π ˚ p a | s qq from the on-policy distribution.
Revised: We then generate a dataset of pairs p s, π ˚ p a | s qq from the on-policy distribution.

Original: The training data consists of 1000 states from each of 64 levels, while test data comes from 2000 levels.
Revised: The training data consists of 1000 states from each of 64 levels, while test data comes from 2000 levels.

Original: A ResNet-18 ˆ π θ is then trained to minimize the loss D KL p π ˚
Revised: A ResNet-18 ˆ π θ is then trained to minimize the loss D KL p π ˚

Original: ||
Revised: ||

Original: ˆ π θ q on the training set.
Revised: ˆ π θ q on the training set.

Original: We compare the generalization performance of regular Adam to a version that uses the AND-mask.
Revised: We compare the generalization performance of regular Adam to a version that uses the AND-mask.

Original: For each method we ran an automatic hyperparameter optimization study using Tree-structured Parzen Estimation (Bergstra et al., 2013) of 1024 trials. 
Revised: For each method we ran an automatic hyperparameter optimization study using Tree-structured Parzen Estimation (Bergstra et al., 2013) of 1024 trials. [Table]

Original: Despite the theoretical computational efﬁciency of computing the AND-mask as presented in Section 2.3 (i.e., linear time and memory in the size of the mini-batch, just like classic SGD), current deep learning frameworks 
Revised: Despite the theoretical computational efﬁciency of computing the ANDmask as presented in Section 2.3 (i.e., linear time and memory in the size of the mini-batch, just like classic SGD), current deep learning frameworks like

Original: Adam Adam+ AND-mask 1.

Original: K L - d i v e r g e n ce on t e s t 1e

Original: Behavioral Cloning on CoinRun like PyTorch (Paszke et al., 2017) have optimized routines that sum gradients across examples in a mini-batch before it is possible to efﬁciently compute the AND-mask.
Revised: PyTorch (Paszke et al., 2017) have optimized routines that sum gradients across examples in a mini-batch before it is possible to efﬁciently compute the AND-mask.

Original: We therefore test the AND-mask in a slightly different way.
Revised: We therefore test the AND-mask in a slightly different way.

Original: In training, in each iteration we sample a batch of data from a randomly chosen level out of the 64 available (and cycle through them all once per epoch).
Revised: In training, in each iteration we sample a batch of data from a randomly chosen level out of the 64 available (and cycle through them all once per epoch).

Original: We then apply the AND-mask ‘temporally’, only allowing gradients that are consistent across time (and therefore across levels).
Revised: We then apply the AND-mask ‘temporally’, only allowing gradients that are consistent across time (and therefore across levels).

Original: See Algorithm 1 in appendix B.6 for a detailed description of this alternative formulation of the AND-mask.
Revised: See Algorithm 1 in appendix B.6 for a detailed description of this alternative formulation of the AND-mask.

Original: The ﬁgure shows the minimum test loss for the 10 best runs, supporting the hypothesis that the AND-mask helps identify invariant mechanisms across different levels.
Revised: The ﬁgure shows the minimum test loss for the 10 best runs, supporting the hypothesis that the AND-mask helps identify invariant mechanisms across different levels.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Generalization and covariate shift.
Revised: Generalization and covariate shift.

Original: The classic formulation of statistical learning theory (Vapnik, 1995) concerns learning from independent and identically distributed samples.
Revised: The classic formulation of statistical learning theory (Vapnik, 1995) concerns learning from independent and identically distributed samples.

Original: The case where the distribution of the covariates at test time differs from the one observed during training is termed covariate shift (Sugiyama et al., 2007; Quionero-Candela et al., 2009; Sugiyama and Kawanabe, 2012).
Revised: The case where the distribution of the covariates at test time differs from the one observed during training is termed covariate shift (Sugiyama et al., 2007; Quionero-Candela et al., 2009; Sugiyama and Kawanabe, 2012).

Original: Standard solutions involve re-weighting of the training examples, but require the additional assumption of overlapping supports for train and test distributions.
Revised: Standard solutions involve re-weighting of the training examples, but require the additional assumption of overlapping supports for train and test distributions.

Original: Causal models and invariances.
Revised: Causal models and invariances.

Original: A fundamental motivation for our work comes from causality (Pearl, 2009; Peters et al., 2017), based on the idea that statistical dependencies are epiphenomena of an underlying causal model.
Revised: A fundamental motivation for our work comes from causality (Pearl, 2009; Peters et al., 2017), based on the idea that statistical dependencies are epiphenomena of an underlying causal model.

Original: The causal description identiﬁes stable elements – e.g. physical mechanisms – connecting causes and effects, which are expected to remain invariant under interventions or changing external conditions (Haavelmo, 1943; Schölkopf et al., 2012)).
Revised: The causal description identiﬁes stable elements – e.g. physical mechanisms – connecting causes and effects, which are expected to remain invariant under interventions or changing external conditions (Haavelmo, 1943; Schölkopf et al., 2012)).

Original: This motivates our notion of invariant mechanisms, and inspired related notions which have been proposed for robust regression (Rojas-Carulla et al., 2018; Heinze-Deml et al., 2018; Arjovsky et al., 2019).
Revised: This motivates our notion of invariant mechanisms, and inspired related notions which have been proposed for robust regression (Rojas-Carulla et al., 2018; Heinze-Deml et al., 2018; Arjovsky et al., 2019; Hermann and Lampinen, 2020; Ahuja et al., 2020; Krueger et al., 2020).

Original: We discuss  these works in more detail in appendix C.1.
Revised: We discuss some of these works in more detail in appendix C.1.

Original: Domain generalization.
Revised: Domain generalization.

Original: ILC can be used in a setting of domain generalization (Muandet et al., 2013), but it is not limited to it: as demonstrated in the experiments in Section 3.2, the AND-mask can be applied even if domain labels are not available.
Revised: ILC can be used in a setting of domain generalization (Muandet et al., 2013), but it is not limited to it: as demonstrated in the experiments in Section 3.2, the AND-mask can be applied even if domain labels are not available.

Original: In contrast, by treating every example as a single domain, methods relying on domain classiﬁers (like DANN Ganin et al.
Revised: In contrast, by treating every example as a single domain, methods relying on domain classiﬁers (like DANN Ganin et al.

Original: (2016) or Balaji et al.
Revised: (2016) or Balaji et al.

Original: (2018)) would require as many output units as there are training examples (i.e. 50’000 for CIFAR-10).
Revised: (2018)) would require as many output units as there are training examples (i.e. 50’000 for CIFAR-10).

Original: Gradient agreement.
Revised: Gradient agreement.

Original: Looking at gradient agreement to learn meaningful representations in neural networks has been explored in (Du et al., 2018;
Revised: Looking at gradient agreement to learn meaningful representations in neural networks has been explored in (Du et al., 2018;

Original: Eshratifar et al., 2018; Fort et al., 2019; Zhang et al., 2019b).
Revised: Eshratifar et al., 2018; Fort et al., 2019; Zhang et al., 2019b).

Original: These approaches mainly rely on a measure of cosine similarity between gradients, which we did not consider here for two main reasons: (i)
Revised: These approaches mainly rely on a measure of cosine similarity between gradients, which we did not consider here for two main reasons: (i)

Original: It is a ‘global’ property of the gradients, and it would not allow us to extract precise information about different patterns in the network; (ii) It is unclear how to extend it beyond pairs of vectors, and for pairwise interactions its computational cost scales quadratic in the number of examples used.
Revised: It is a ‘global’ property of the gradients, and it would not allow us to extract precise information about different patterns in the network; (ii) It is unclear how to extend it beyond pairs of vectors, and for pairwise interactions its computational cost scales quadratic in the number of examples used.

Original:  C ONCLUSIONS
Revised:  C ONCLUSIONS

Original: Generalizing out of distribution is one of the most signiﬁcant open challenges in machine learning, and relying on invariances across environments or examples may be key in certain contexts.
Revised: Generalizing out of distribution is one of the most signiﬁcant open challenges in machine learning, and relying on invariances across environments or examples may be key in certain contexts.

Original: In this paper we analyzed how neural networks trained by averaging gradients across examples might converge to solutions that ignore the invariances, especially if these are harder to learn than spurious patterns.
Revised: In this paper we analyzed how neural networks trained by averaging gradients across examples might converge to solutions that ignore the invariances, especially if these are harder to learn than spurious patterns.

Original: We argued that if learning signals are collected on one example at the time — as it is the case for gradients, e.g., computed with backpropagation — the way these signals are aggregated can play a signiﬁcant role in the patterns that will ultimately be expressed: Averaging gradients in particular can be too permissive, acting as a logical OR of a collection of distinct patterns, and lead to a ‘patchwork’ solution.
Revised: We argued that if learning signals are collected on one example at the time — as it is the case for gradients, e.g., computed with backpropagation — the way these signals are aggregated can play a signiﬁcant role in the patterns that will ultimately be expressed: Averaging gradients in particular can be too permissive, acting as a logical OR of a collection of distinct patterns, and lead to a ‘patchwork’ solution.

Original: We introduced and formalized the concept of Invariant Learning Consistency, and showed how to learn invariances even in the face of alternative explanations that — although spurious — fulﬁll most characteristics of a good solution.
Revised: We introduced and formalized the concept of Invariant Learning Consistency, and showed how to learn invariances even in the face of alternative explanations that — although spurious — fulﬁll most characteristics of a good solution.

Original: The AND-mask is but one of multiple possible ways to improve consistency, and it is unlikely to be a practical algorithm for all applications.
Revised: The AND-mask is but one of multiple possible ways to improve consistency, and it is unlikely to be a practical algorithm for all applications.

Original: However, we believe this should not distract from the general idea which we are trying to put forward — namely, that it is worthwhile to study learning of explanations that are hard to vary , with the longer term goal of advancing our understanding of learning, memorization and generalization.
Revised: However, we believe this should not distract from the general idea which we are trying to put forward — namely, that it is worthwhile to study learning of explanations that are hard to vary , with the longer term goal of advancing our understanding of learning, memorization and generalization.

Original: A A PPENDIX TO S ECTIONA. A CLASSIC EXAMPLE OF A PATCHWORK SOLUTION
Revised: A A PPENDIX TO S ECTIONA. A CLASSIC EXAMPLE OF A PATCHWORK SOLUTION

Original: Consider a neural network with one hidden layer consisting of two neurons and sigmoidal activations:
Revised: Consider a neural network with one hidden layer consisting of two neurons and sigmoidal activations:

Original: [Equation]
Revised: [Equation]

Original: We want to learn the continuous function f ˚ : r 0 , 1 s Ñ r 0 , 2 s deﬁned as
Revised: We want to learn the continuous function f ˚ : r 0 , 1 s Ñ r 0 , 2 s deﬁned as

Original: [Equation]
Revised: [Equation]

Original: To perform this task, we have access to (noiseless) data from two environments:
Revised: To perform this task, we have access to (noiseless) data from two environments:

Original: [Equation]
Revised: [Equation]

Original: There is a simple constructive way, provided by the universal function approximation theorem Cybenko (1989) to ﬁt this function 11 using f θ up to an arbitrarily small mean squared error L A ` B p θ ˚ q .
Revised: There is a simple constructive way, provided by the universal function approximation theorem Cybenko (1989) to ﬁt this function 11 using f θ up to an arbitrarily small mean squared error L A ` B p θ ˚ q .

Original: Leaving out the details of such a construction (Cybenko (1989) for details), the reader can check on the left panel of Figure 10 that θ ˚ “ p 100 , ´ 50 , 100 , ´ 75 , 1 , 1 q provides a good ﬁt for both environments A and B — both L A p θ ˚ q and L B p θ ˚ q are small.
Revised: Leaving out the details of such a construction (Cybenko (1989) for details), the reader can check on the left panel of Figure 10 that θ ˚ “ p 100 , ´ 50 , 100 , ´ 75 , 1 , 1 q provides a good ﬁt for both environments A and B — both L A p θ ˚ q and L B p θ ˚ q are small.

Original: θ “ θ ˚ “ p 100 , ´ 45 , 100 , ´ 75 , 1 , 1 q 0.
Revised: θ “ θ ˚ “ p 100 , ´ 45 , 100 , ´ 75 , 1 , 1 q 0.

Original: θ “ ˜ θ ˚ “ p 100 , ´ 45 , 100 , ´ 75 , 1 , ´ 0 .
Revised: θ “ ˜ θ ˚ “ p 100 , ´ 45 , 100 , ´ 75 , 1 , ´ 0 .

Original: 5 q
Revised: 5 q

Original: However, it is easy to realize that θ ˚ — while being a solution which can be returned by gradient descent using the pooled data A+B — is not consistent (formal deﬁnition given in the main paper in
Revised: However, it is easy to realize that θ ˚ — while being a solution which can be returned by gradient descent using the pooled data A+B — is not consistent (formal deﬁnition given in the main paper in

Original: Section 2).
Revised: Section 2).

Original: Indeed, it is possible to modify ˜ θ ˚ such that the loss in environment A remains almost unchanged, while the loss in environment B gets larger.
Revised: Indeed, it is possible to modify ˜ θ ˚ such that the loss in environment A remains almost unchanged, while the loss in environment B gets larger.

Original: In particular, on the right panel of Figure 10, we show that ˜ θ ˚ “ p 100 , ´ 50 , 100 , ´ 75 , 1 , ´ 0 .
Revised: In particular, on the right panel of Figure 10, we show that ˜ θ ˚ “ p 100 , ´ 50 , 100 , ´ 75 , 1 , ´ 0 .

Original: 5 q is such that L A p θ ˚ q ď L A p ˜ θ ˚ q ` (cid:15) (with (cid:15) very small) but L B p θ ˚ q !
Revised: 5 q is such that L A p θ ˚ q ď L A p ˜ θ ˚ q ` (cid:15) (with (cid:15) very small) but L B p θ ˚ q !

Original: L B p ˜ θ ˚ q .
Revised: L B p ˜ θ ˚ q .

Original: According to our deﬁnition in Equation 1 (see main paper), we have
Revised: According to our deﬁnition in Equation 1 (see main paper), we have

Original: I (cid:15) p θ ˚ q ď | L B p θ ˚ q ´ L B p ˜ θ ˚ q| — that is a large number (low consistency).
Revised: I (cid:15) p θ ˚ q ď | L B p θ ˚ q ´ L B p ˜ θ ˚ q| — that is a large number (low consistency).

Original: Remark 1 (Connection to out of distribution generalization) .
Revised: Remark 1 (Connection to out of distribution generalization) .

Original: The main point of this analysis was to show an example of where our measure of consistency behaves according to expectations:
Revised: The main point of this analysis was to show an example of where our measure of consistency behaves according to expectations:

Original: A typical implementation of the universal approximation theorem — which one would not expect to generalize out of distribution, due to its ‘patchwork’ behavior — leads indeed to a very low consistency score.
Revised: A typical implementation of the universal approximation theorem — which one would not expect to generalize out of distribution, due to its ‘patchwork’ behavior — leads indeed to a very low consistency score.

Original: A. S ECTION 2.2: C ONSISTENCY AS ARITHMETIC / GEOMETRIC MEAN OF LANDSCAPES 
Revised: A. S ECTION 2.2: C ONSISTENCY AS ARITHMETIC / GEOMETRIC MEAN OF LANDSCAPES Geometric mean of matrices.

Original: Geometric mean of matrices.

Original: Given an n -tuple of d ˆ d positive deﬁnite matrices p
Revised: Given an n -tuple of d ˆ d positive deﬁnite matrices p

Original: A j q nj “ 1 , the geometric (Karcher) mean Ando et al.
Revised: A j q nj “ 1 , the geometric (Karcher) mean Ando et al.

Original: (2004) is the unique positive deﬁnite solution X to the equation ř mi
Revised: (2004) is the unique positive deﬁnite solution X to the equation ř mi

Original: “ 1 log p A ´ 

Revised: “ 1 log p A ´

Original: i X q “ 0 , where log is the matrix logarithm.
Revised: i X q “ 0 , where log is the matrix logarithm.

Original: This matrix average has many desirable properties, which make it relevant to signal processing and medical imaging.
Revised: This matrix average has many desirable properties, which make it relevant to signal processing and medical imaging.

Original: The Karcher mean can also be written as arg min X P S
Revised: The Karcher mean can also be written as arg min X P S

Revised: [Table]

Original: `` p d q f p X q “ 12 m ř mi “ 1 d p A i , X q 2 , where d is the Riemannian distance in the manifold of SPD matrices S
Revised: ` p d q f p X q “ 12 m ř mi “ 1 d p A i , X q 2 , where d is the Riemannian distance in the manifold of SPD matrices S

Original: `` p d q .
Revised: `` p d q .

Revised: [Table]

Revised: [Figure]

Original: Link between consistency and geometric means.
Revised: Link between consistency and geometric means.

Original: Here we show how the consistency score introduced in Equation 1 can be linked (in a simpliﬁed setting) to a comparison between the arithmetic and geometric means of the Hessians approximating the landscapes of two separate environments A and B .
Revised: Here we show how the consistency score introduced in Equation 1 can be linked (in a simpliﬁed setting) to a comparison between the arithmetic and geometric means of the Hessians approximating the landscapes of two separate environments A and B .

Original: At the local minimizer θ ˚ “ 0 , we assume that L A “ L B “ 0 and consider the local quadratic approximations L A p θ q “ 12 θ J H
Revised: At the local minimizer θ ˚ “ [Equation] we assume that L A “ L B “ 0 and consider the local quadratic approximations L A p θ q “ 12 θ J H

Original: A θ and L B p θ q “ 12 θ J H B θ .
Revised: A θ and L B p θ q “ 12 θ J H B θ .

Original: Here, we make the additional simplifying assumption that H A and H B are diagonal (or, more broadly, co-diagonalizable):
Revised: Here [Equation] we make the additional simplifying assumption that H A and H B are diagonal (or, more broadly, co-diagonalizable):

Original: H A “ diag p λ 
Revised: [Equation] λ Bn q

Original: A 1 , ¨ ¨ ¨ , λ An q , H B “ diag p λ B 1 , ¨ ¨ ¨ , λ Bn q , with λ
Revised: [Equation]

Original:  Ai ě 0 and λ Bi ě 0 for all i “ 1 , . . . , n .
Revised: with λ Ai ě 0 and λ Bi ě 0 for all i “ 1 , . . . , n .

Original: The arithmetic and geometric means (noted as H A ` B and H A ^ B ) of these matrices are deﬁned in this simpliﬁed setting as follows:
Revised: The arithmetic and geometric means (noted as H A ` B and H A ^ B ) of these matrices are deﬁned in this simpliﬁed setting as follows:

Original: [Equation]
Revised: [Equation]

Original: As motivated in the main paper and in Figure 12, one can link the consistency of two landscapes to a comparison between the geometric and arithmetic means of the corresponding Hessians.
Revised: As motivated in the main paper and in Figure [Equation] one can link the consistency of two landscapes to a comparison between the geometric and arithmetic means of the corresponding Hessians.

Original: Proposition 3 .
Revised: Proposition 3 .

Original: In the setting we just described, the consistency score in Equation 1 can be estimated as follows:
Revised: In the setting we just described, the consistency score in Equation 1 can be estimated as follows:

Original: [Equation]
Revised: [Equation]

Revised: [Figure]

Original: Before showing the proof, we note that the proposition gives a lower bound on the consistency.
Revised: Before showing the proof, we note that the proposition gives a lower bound on the consistency.

Original: That is, it provides a pessimistic estimate.
Revised: That is, it provides a pessimistic estimate.

Original: Yet, as we motivated, this estimate has a nice geometric interpretation.
Revised: Yet, as we motivated, this estimate has a nice geometric interpretation.

Original: However, as we outline in a remark after the proof, this estimate is tight in two important limit cases.
Revised: However, as we outline in a remark after the proof, this estimate is tight in two important limit cases.

Original: Proof.
Revised: Proof.

Original: In this setting
Revised: In this setting

Original: [Equation]
Revised: [Equation]

Original: Equation 1 gives
Revised: Equation 1 gives

Original: [Equation]
Revised: [Equation]

Original: Recall that
Revised: Recall that

Original: [Equation]
Revised: [Equation]

Original: Hence
Revised: Hence

Original: [Equation]
Revised: [Equation]

Original: this is a simple quadratic program with quadratic constraints
Revised: this is a simple quadratic program with quadratic constraints

Original: [Equation]
Revised: [Equation]

Original: and 
Revised: and max

Original: [Equation]
Revised: [Equation]

Original: Further
Revised: Further

Original: [Equation]
Revised: [Equation]

Original: we can change variables and introduce ˜
Revised: we can change variables and introduce ˜

Original: θ i
Revised: θ i

Original: “ θ i a λ Ai { 2 .
Revised: “ θ i a λ Ai { 2 .

Original: The problem gets even simpler: 
Revised: The problem gets even simpler: max

Original: [Equation]
Revised: [Equation]

Original: All in all [Equation] we get
Revised: All in all, we get

Original: [Equation]
Revised: [Equation]

Original: This means
Revised: This means

Original: [Equation]
Revised: [Equation]

Original: where the ﬁrst inequality comes from the monotonicity of the square root function, and the second inequality comes from the fact that (i) the geometric mean is always smaller or equal than the arithmetic mean and (ii) for any sequence of numbers α i ą 1 , max i α i ď
Revised: where the ﬁrst inequality comes from the monotonicity of the square root function, and the second inequality comes from the fact that (i) the geometric mean is always smaller or equal than the arithmetic mean and (ii) for any sequence of numbers α i ą 1 , max i α i ď

Original: ś i
Revised: ś i

Original: α
Revised: α

Original: i .
Revised: i .

Original: Remark 2 (Sanity check) .
Revised: Remark 2 (Sanity check) .

Original: There are two important cases where we can test the bound above.
Revised: There are two important cases where we can test the bound above.

Original: First, if
Revised: First, if

Original: H A “ H B , then I (cid:15) p θ ˚ q “ (cid:15) , and the bound returns I (cid:15) p θ ˚ q ď 2 (cid:15) , since the geometric and arithmetic mean are the same.
Revised: H A “ H B , then I (cid:15) p θ ˚ q “ (cid:15) , and the bound returns I (cid:15) p θ ˚ q ď 2 (cid:15) , since the geometric and arithmetic mean are the same.

Original: Next, say λ
Revised: Next, say λ

Original: Ai “ 0 but λ Bi ą 0 ; then, both the bound and the inconsistency score are 8 (highest possible inconsistency).
Revised: Ai “ 0 but λ Bi ą 0 ; then, both the bound and the inconsistency score are 8 (highest possible inconsistency).

Original: A. P ROOF OF P ROPOSITION
Revised: A. P ROOF OF P ROPOSITION

Original: In this appendix section we consider the AND-masked GD algorithm, introduced at the end of
Revised: In this appendix section we consider the AND-masked GD algorithm, introduced at the end of

Original: Section 2.
Revised: Section 2.

Original: We recall that the masked gradients at iteration k are m t p θ k q d ∇L p θ k q , where m t p θ k q vanishes for any component where there are less than t P t d { 2 ` 1 , . . .
Revised: We recall that the masked gradients at iteration k are m t p θ k q d ∇L p θ k q , where m t p θ k q vanishes for any component where there are less than t P t d { 2 ` 1 , . . .

Original: , d u agreeing gradient signs across environments, and is equal to one otherwise.
Revised: , d u agreeing gradient signs across environments, and is equal to one otherwise.

Original: In a full-batch setting, the algorithm is
Revised: In a full-batch setting, the algorithm is

Original: [Equation]
Revised: [Equation]

Original: where η ą 0 is the learning rate.
Revised: where η ą 0 is the learning rate.

Original: Proposition 1 .
Revised: Proposition 1 .

Original: Let L have L -Lipschitz gradients and consider a learning rate η ď 1 { L .
Revised: Let L have L -Lipschitz gradients and consider a learning rate η ď 1 { L .

Original: After k iterations, AND-masked GD visits at least once a point θ where } m t p θ q d ∇L p θ q} 

Revised: After k iterations, AND-masked GD visits at least once a point θ where } m t p θ q d ∇L p θ q}

Original: ď O p 1 { k q .
Revised: ď O p 1 { k q .

Original: Proof.
Revised: Proof.

Original: Thanks to the component-wise L -smoothness and using a Taylor expansion around θ
Revised: Thanks to the component-wise L -smoothness and using a Taylor expansion around θ

Original: i we have
Revised: i we have

Original: [Equation]
Revised: [Equation]

Original: If we seek η ´ Lη 2 { 2 ě η { 2 , then η ď 1 L , as we assumed in the proposition statement.
Revised: If we seek η ´ Lη 2 { 2 ě η { 2 , then η ď 1 L , as we assumed in the proposition statement.

Original: Therefore,
Revised: Therefore,

Original: L p θ i ` 1 q ď L p θ i q ´ p η { 2 q} m t p θ i q d ∇L p θ i q} 2 , for all i ě 0 .
Revised: L p θ i ` 1 q ď L p θ i q ´ p η { 2 q} m t p θ i q d ∇L p θ i q} 2 , for all i ě 0 .

Original: Summing over i from 0 to a desired iteration k , we getk ´ 1 ÿ
Revised: Summing over i from 0 to a desired iteration k , we get

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: Therefore, min 
Revised: Therefore, min i “

Revised: [Equation]

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: Hence, there exist an iteration i ˚ P t 0 , . . .
Revised: Hence, there exist an iteration i ˚ P t 0 , . . .

Original: , k u such that } m t p θ i ˚ q d ∇L p
Revised: , k u such that } m t p θ i ˚ q d ∇L p

Original: θ i ˚ q} 

Revised: θ i ˚ q}

Original: ď O p 1 { k q .
Revised: ď O p 1 { k q .

Original: A. P ROOF OF P ROPOSITION
Revised: A. P ROOF OF P ROPOSITION

Original: Here we ﬁx parameters θ P R n and assume gradients ∇L e p θ q P R n coming from environments e P E are drawn independently from a multivariate Gaussian with zero mean and σ 2 I covariance.
Revised: Here we ﬁx parameters θ P R n and assume gradients ∇L e p θ q P R n coming from environments e P E are drawn independently from a multivariate Gaussian with zero mean and σ 2 I covariance.

Original: We want to show that, in this random setting, the AND-mask introduced in Section 2.3 decreases the magnitude of the gradient step.
Revised: We want to show that, in this random setting, the AND-mask introduced in Section 2.3 decreases the magnitude of the gradient step.

Original: Proposition 2 .
Revised: Proposition 2 .

Original: Consider the setting we just outlined, with L “ p 1 { d q ř de “ 1 L e .
Revised: Consider the setting we just outlined, with L “ p 1 { d q ř de “ 1 L e .

Original: While E } ∇L p θ q} 2 “
Revised: While E } ∇L p θ q} 2 “

Original: O p n { d q , we have that @
Revised: O p n { d q , we have that @

Original: t P t d { 2 ` 1 , . . .
Revised: t P t d { 2 ` 1 , . . .

Original: , d u , D c
Revised: , d u , D c

Original: P p 1 , 2 s such that E } m t p θ q d ∇L p θ q} 2 ď
Revised: P p 1 , 2 s such that E } m t p θ q d ∇L p θ q} 2 ď

Original: O p n { c d q .
Revised: O p n { c d q .

Original: Proof.
Revised: Proof.

Original: Let us drop the argument θ for ease of notation.
Revised: Let us drop the argument θ for ease of notation.

Original: First, let us consider ∇L (no gradient
Revised: First, let us consider ∇L (no gradient

Original: AND-mask):
Revised: AND-mask):

Original: [Equation]
Revised: [Equation]

Original: where in the ﬁrst equality we used the fact that the ∇L e i are uncorrelated and in the second the fact that E r} ∇L e i } 2 s is the trace of the covariance of ∇L e i .
Revised: where in the ﬁrst equality we used the fact that the ∇L e i are uncorrelated and in the second the fact that E r} ∇L e i } 2 s is the trace of the covariance of ∇L e i .

Original: Next, assume we apply the element-wise AND-mask m t to the gradients, which puts to zero the components (dimensions) where there are less than t P t d { 2 , . . .
Revised: Next, assume we apply the element-wise AND-mask m t to the gradients, which puts to zero the components (dimensions) where there are less than t P t d { 2 , . . .

Original: , d u equal signs.
Revised: , d u equal signs.

Original: Since Gaussians are symmetric around zero, the probability of having exactly u positive j -th gradient component among d environments is Pr p p j “ u q “ ` ˘ d ` du ˘ .
Revised: Since Gaussians are symmetric around zero, the probability of having exactly u positive j -th gradient component among d environments is Pr p p j “ u q “ ` ˘ d ` du ˘ .

Original: Hence, the probability to keep the j -th gradient direction (considering also negative consistency) is
Revised: Hence, the probability to keep the j -th gradient direction (considering also negative consistency) is

Original: [Equation]
Revised: [Equation]

Original: We would now like to compute E ››› m t d
Revised: We would now like to compute E ››› m t d

Original: [Equation]
Revised: [Equation]

Original: di
Revised: di

Original: [Equation]
Revised: [Equation]

Original: 1 ∇L e i ¯››› 2 .
Revised: e i ¯››› 2 .

Original: The difﬁculty lies in the fact that the event m t “ 1 makes gradients conditionally dependent .
Revised: The difﬁculty lies in the fact that the event m t “ 1 makes gradients conditionally dependent .

Original: Indeed, conditioning on both m t “ 1 and r ∇L e s j ą 0 changes the distribution of r ∇L e 

Revised: Indeed, conditioning on both m t “ 1 and r ∇L e s j ą 0 changes the distribution of r ∇L e

Original: s j : this gradient entry is going to be more likely to be positive or negative, depending on the value of r ∇L e s j and on the details of the gradient mask.
Revised: s j : this gradient entry is going to be more likely to be positive or negative, depending on the value of r ∇L e s j and on the details of the gradient mask.

Original: To solve the issue, we our strategy is to reduce the discussion (without loss in generality and with no additional assumption) to the case where gradient entries have all the same sign and hence conditional independence is restored.
Revised: To solve the issue, we our strategy is to reduce the discussion (without loss in generality and with no additional assumption) to the case where gradient entries have all the same sign and hence conditional independence is restored.

Original: We consider the following writing for the quantity we are interested in:
Revised: We consider the following writing for the quantity we are interested in:

Original: [Equation]
Revised: [Equation]

Original: r p j “ ˆ p j s“ n ÿ j “p d ´ t q ÿ ˆ p j “d ÿ ˆ p j  “ t
Revised: s j [Equation] ﬁ ﬂ “ [Equation]

Original: E » – ˜ d d ÿ i “r ∇L e i s j ¸ 2 ˇˇˇˇ p j
Revised: [Equation]

Original: “ ˆ p j ﬁ ﬂ Pr r p j “ ˆ p j s “n ÿ j “d ÿ ˆ p j “ t
Revised: p [Equation] E »

Original: E » – ˜ dd ÿ i “r ∇L e i s j ¸
Revised: [Equation] m [Equation] ˜ d

Original: [Equation]
Revised: [Equation]

Revised: ∇L

Revised: [Equation]

Original: ˇˇˇˇ p
Revised: ˇˇˇˇ p

Revised: [Equation]

Revised: [Equation]

Original: j “ ˆ p j ﬁ ﬂ ˆ ˙
Revised: ﬁ ﬂ Pr [Equation] “ n ÿ [Equation]

Revised: [Equation]

Revised: d

Original: [Equation]
Revised: [Equation]

Original:  ˆ 
Revised: ÿ ˆ p [Equation] [Equation] d ÿ ˆ p j “

Original: [Equation]
Revised: [Equation]

Revised: E

Revised: [Equation]

Revised: [Equation]

Original: ˆ p j ˙ , where we used the deﬁnition of 2 -norm, the law of total expectation, and the symmetry of the problem with respect to positive and negative numbers.
Revised: [Equation] ˇˇˇˇ [Equation] where we used the deﬁnition of 2 -norm, the law of total expectation, and the symmetry of the problem with respect to positive and negative numbers.

Original: Finally, since the gradient components within the same environment are conditionally independent, for any j P t 1 , . .
Revised: Finally, since the gradient components within the same environment are conditionally independent, for any j P t 1 , . .

Original: . , n u we can write
Revised: . , n u we can write

Original: [Equation]
Revised: [Equation]

Original: Finally, we note that the following bound holds:
Revised: Finally, we note that the following bound holds:

Original: [Equation]
Revised: [Equation]

Original: Indeed, if all environments lead to positive (or, symmetrically, negative) and non-interacting gradients in the j -th direction, the average will be the biggest in norm.
Revised: Indeed, if all environments lead to positive (or, symmetrically, negative) and non-interacting gradients in the j -th direction, the average will be the biggest in norm.

Original: Moreover — crucially — conditioned on the event p
Revised: Moreover — crucially — conditioned on the event p

Original: j “ d , gradients coming from different environments are distributed as a positive half-normal distributions.
Revised: j “ d , gradients coming from different environments are distributed as a positive half-normal distributions.

Original: Moreover, they are conditionally independent ; this because, since they are all positive, the value of a gradient in one environment cannot inﬂuence the value of the gradient in another one.
Revised: Moreover, they are conditionally independent ; this because, since they are all positive, the value of a gradient in one environment cannot inﬂuence the value of the gradient in another one.

Original: We remark that conditional independence on the right-hand side is therefore not an assumption , but is intrinsic to the upper bound.
Revised: We remark that conditional independence on the right-hand side is therefore not an assumption , but is intrinsic to the upper bound.

Original: Putting it all together, we have
Revised: Putting it all together, we have

Original: [Equation]
Revised: [Equation]

Original: where in the second line we bounded the squared average of a sum of half normal distributions: let t X i u di “ 1 be a family of uncorrelated positive half-normal distributions derived from a Gaussians with mean zero and variance σ 2 , we have 12 that E r X
Revised: where in the second line we bounded the squared average of a sum of half normal distributions: let t X i u di “ 1 be a family of uncorrelated positive half-normal distributions derived from a Gaussians with mean zero and variance σ 2 , we have 12 that E r X

Original: i s “ σ a 2 { π and E r X 

Revised: i s “ σ [Equation] 2 { π and E r X

Original: i s “ σ 2 .
Revised: i s “ σ 2 .

Original: Also, E r X i X j s “ Er X i s E r
Revised: Also, E r X i X j s “ E

Original: X j s ď σ 2 .
Revised: [Equation]

Original: Therefore,
Revised: Therefore,

Original: E

Original: [Equation]
Revised: [Equation]

Original: Finally, if we set r “ t { d P p 0 . 5 , 1 s , we haveˆ d t ˙ „
Revised: Finally, if we set r “ t { d P p 0 . 5 , 1 s , we have

Original: ˆ r r p 1 ´ r q 1 ´ r ˙ d as d Ñ 8 (discarding all polynomial terms).
Revised: [Equation] r ˙ d as d Ñ 8 (discarding all polynomial terms).

Original: Hence ` dt ˘ is of the form q d , with 1 ď q ă 2 .
Revised: Hence ` dt ˘ is of the form q d , with 1 ď q ă 2 .

Original: So, the quantity σ 2 n p d ´ t q ` dt ˘ ` ˘ d ´ 1 will be exponentially decreasing at a rate O p n {p 2 ´ q q d q .
Revised: So, the quantity σ 2 n p d ´ t q ` dt ˘ ` ˘ d ´ 1 will be exponentially decreasing at a rate O p n {p 2 ´ q q d q .

Original: Notably, if t “ d { 2 , then we lose the exponential rate and get back to O p n { d q .
Revised: Notably, if t “ d { 2 , then we lose the exponential rate and get back to O p n { d q .

Original: B A PPENDIX TO S ECTIONB. S
Revised: B A PPENDIX TO S ECTION

Original: [Table] 
Revised: We used Pytorch Paszke et al. (2017) to implement all experiments in this paper. B. S ECTION 3.

Original: [Table]
Revised: [Table]

Original: B. D ATASET
Revised: B. D ATASET

Revised: [Figure]

Original: Here we report more technical details about the synthetic dataset described inSection
Revised: Here we report more technical details about the synthetic dataset described in

Original: [Table] [Figure]
Revised: Section 3.

Original: Each example is constructed as follows: we ﬁrst choose the label randomly to be either ` 1 or ´ 1 , with equal probability.
Revised: Each example is constructed as follows: we ﬁrst choose the label randomly to be either ` 1 or ´ 1 , with equal probability.

Original: The example is a vector with d S ` d M entries, consisting of the shortcut and the mechanism .
Revised: The example is a vector with d S ` d M entries, consisting of the shortcut and the mechanism .

Original: In our experiments, d M “ 2 and d S “ 32 .
Revised: In our experiments, d M “ 2 and d S “ 32 .

Original: The Gaussian shortcuts are obtained by ﬁrst sampling one random vector x s P R d S per environment.
Revised: The Gaussian shortcuts are obtained by ﬁrst sampling one random vector x s P R d S per environment.

Original: Its components x s,i are sampled independently from a Normal distribution: x s,i „ N p 0 , 0 .
Revised: Its components x s,i are sampled independently from a Normal distribution: x s,i „ N p 0 , 0 .

Original: 1 q .
Revised: 1 q .

Original: We use x s for class 1, and ´ x s for class -1.
Revised: We use x s for class 1, and ´ x s for class -1.

Original: In the test set, all shortcut components are sampled i.i.d.
Revised: In the test set, all shortcut components are sampled i.i.d.

Original: from the same Normal distribution.
Revised: from the same Normal distribution.

Original: Effectively, each example of the test set belongs to a different domain.
Revised: Effectively, each example of the test set belongs to a different domain.

Original: The mechanism is implemented as the two interconnected spirals shown in Figure 13 by sampling the radius r „ Unif p 0 .
Revised: The mechanism is implemented as the two interconnected spirals shown in Figure 13 by sampling the radius r „ Unif p 0 .

Original: q
Revised: q

Original: and then computing the angle as α “ 2 πnr where n is the number of revolutions of the spiral.
Revised: and then computing the angle as α “ 2 πnr where n is the number of revolutions of the spiral.

Original: We add uniform noise in the range r´ 0 .
Revised: We add uniform noise in the range r´ 0 .

Original: 02 s to the radii afterwards.
Revised: 02 s to the radii afterwards.

Original: The training dataset consists of 1280 examples per environment and we use D “ 32 environments unless otherwise mentioned.
Revised: The training dataset consists of 1280 examples per environment and we use D “ 32 environments unless otherwise mentioned.

Original: The training datasets consists of 2000 examples.
Revised: The training datasets consists of 2000 examples.

Original: B. E XPERIMENT
Revised: B. E XPERIMENT

Original: We train all networks for t 3000 { D u epochs, dropping the learning rate by a factor 10 halfway through, and again at three-quarters of training.
Revised: We train all networks for t 3000 { D u epochs, dropping the learning rate by a factor 10 halfway through, and again at three-quarters of training.

Original: For computational reason, we stop each trial before completion if the training accuracy exceeds 97% and the test accuracy is below 60%.
Revised: For computational reason, we stop each trial before completion if the training accuracy exceeds 97% and the test accuracy is below 60%.

Original: All networks are MLPs with
Revised: All networks are MLPs with

Original: LeakyReLU activation functions and a cross-entropy loss on the output.
Revised: LeakyReLU activation functions and a cross-entropy loss on the output.

Original: We run a hyperparameter search over the ranges shown in Table 1.
Revised: We run a hyperparameter search over the ranges shown in Table 1.

Original: For IRM and the AND-mask, we select the best-performing run and re-run it 50 times with different random seeds.
Revised: For IRM and the AND-mask, we select the best-performing run and re-run it 50 times with different random seeds.

Original: For DANN and the standard baselines nothing produced results signiﬁcantly better than chance.
Revised: For DANN and the standard baselines nothing produced results signiﬁcantly better than chance.

Original: B.3. S TANDARD REGULARIZERS AND AND- MASK
Revised: B.3. S TANDARD REGULARIZERS AND AND- MASK

Original: The networks with the L1, L2, Dropout and Batch-normalization regularizers, have hyperparameters that were randomly selected from Table 1.
Revised: The networks with the L1, L2, Dropout and Batch-normalization regularizers, have hyperparameters that were randomly selected from Table 1.

Original: For the AND-mask we used the very same ranges.
Revised: For the AND-mask we used the very same ranges.

Original: The regularizers L1 and L2 are never combined; instead, one weight regularization type out of L1, L and none is selected and we sample from the respective range afterwards.
Revised: The regularizers L1 and L2 are never combined; instead, one weight regularization type out of L1, L and none is selected and we sample from the respective range afterwards.

Original: The parameters found to work best from the grid search were:  256 hidden units, 3 hidden layers, batch size 128, Adam with learning rate 1e-2, no batch norm, no dropout, L2-regularization with a coefﬁcient of 1e-4, no 
Revised: The parameters found to work best from the grid search were: agreement threshold of 1, 256 hidden units, 3 hidden layers, batch size 128, Adam with learning rate 1e-2, no batch norm, no dropout, L2-regularization with a coefﬁcient of 1e-4, no L1-regularization.

Original: L1-regularization.
Revised: In practice, we often found it helpful to rescale the gradients after masking to compensate for the decreasing overall magnitude. We add the option for gradient rescaling as an additional hyperparameter, as we found it to help in several experiments. It rescales gradient components layer-wise after masking, by multiplying the remaining gradient components by c , where c is the ratio of the number of components in that layer over the number of non-masked components in that layer (i.e. the sum of the binary elements in the mask). We speculate that for very large layers, a less extreme normalization scheme or the additional use of gradient clipping might be appropriate.

Original: B.3. D OMAIN A DVERSARIAL N EURAL N ETWORKS
Revised: B.3. D OMAIN A DVERSARIAL N EURAL N ETWORKS

Original: The experiments using DANN follow a similar pattern.
Revised: The experiments using DANN follow a similar pattern.

Original: The model consists of an embedding network, a classiﬁcation network, and a “domain discrimination” network.
Revised: The model consists of an embedding network, a classiﬁcation network, and a “domain discrimination” network.

Original: All three modules are two-layer multi-layer perceptrons (MLP).
Revised: All three modules are two-layer multi-layer perceptrons (MLP).

Original: The number of hidden units of all MLPs are sampled from the range speciﬁed in Table 1, and we trained 100 models.
Revised: The number of hidden units of all MLPs are sampled from the range speciﬁed in Table 1, and we trained 100 models.

Original: Both label classiﬁer and domain discriminator are applied to the output of the embedding network.
Revised: Both label classiﬁer and domain discriminator are applied to the output of the embedding network.

Original: The label classiﬁer is trained to minimize the cross-entropy-loss between the predicted and the true label.
Revised: The label classiﬁer is trained to minimize the cross-entropy-loss between the predicted and the true label.

Original: Similarly, the domain discriminator is trained to minimize the loss between predicted and true domain-label.
Revised: Similarly, the domain discriminator is trained to minimize the loss between predicted and true domain-label.

Original: The embedding network is trained to minimize the regular task classiﬁcation loss and at the same time to maximize the the domain-loss achieved by the domain discriminator.
Revised: The embedding network is trained to minimize the regular task classiﬁcation loss and at the same time to maximize the the domain-loss achieved by the domain discriminator.

Original: B.3. I NVARIANT R ISK M INIMIZATION
Revised: B.3. I NVARIANT R ISK M INIMIZATION

Original: For the experiments using IRM we used the authors’ PyTorch implementation from https: //github.com/facebookresearch/InvariantRiskMinimization .
Revised: For the experiments using IRM we used the authors’ PyTorch implementation from https: //github.com/facebookresearch/InvariantRiskMinimization .

Original: We perform a random hyperparameter search over with the ranges shown in Table
Revised: We perform a random hyperparameter search over with the ranges shown in Table

Original: B.
Revised: [Table]

Original: . C URVES FOR ALL EXPERIMENTS
Revised: .3. C URVES FOR ALL EXPERIMENTS

Original: In Figure 14 we show the learning curves of training and test accuracy for the different methods.[Table] [Figure]
Revised: In Figure 14 we show the learning curves of training and test accuracy for the different methods.

Original: B.3. C ORRELATION PLOTS
Revised: B.3. C ORRELATION PLOTS

Original: For the correlation plots in Figure 7 we used a randomly initialized MLP with the following conﬁguration: 3 hidden layers, 256 hidden units.
Revised: For the correlation plots in Figure 7 we used a randomly initialized MLP with the following conﬁguration: 3 hidden layers, 256 hidden units.

Original: The dataset was using 16 environments and batches of size 1024.
Revised: The dataset was using 16 environments and batches of size 1024.

Original: The lines in Figure 7 are linear least-squares regressions to the gradient data shown as scatter plots.
Revised: The lines in Figure 7 are linear least-squares regressions to the gradient data shown as scatter plots.

Original: We repeat the experiment 10 times with different network weight seeds, resulting in theregression lines.
Revised: We repeat the experiment 10 times with different network weight seeds, resulting in theregression lines.

Original: Zero gradients are excluded from the regression computation, as most gradients are masked out by the product mask in both cases. 
Revised: Zero gradients are excluded from the regression computation, as most gradients are masked out by the product mask in both cases. [Figure] [Figure]

Original: B. F URTHER VISUALIZATIONS AND EXPERIMENTS
Revised: B. F URTHER VISUALIZATIONS AND EXPERIMENTS

Original: In Figure 15 we show how many environments need to be present for the baseline without AND-mask to switch the decision boundary from the shortcuts to the mechanism.
Revised: In Figure 15 we show how many environments need to be present for the baseline without AND-mask to switch the decision boundary from the shortcuts to the mechanism.

Original: Under the same experimental condition as in the main paper, the baseline ﬁrst succeeds at 1024 environments.
Revised: Under the same experimental condition as in the main paper, the baseline ﬁrst succeeds at 1024 environments.

Original: [Figure]
Revised: [Figure]

Original: B. S ECTION 3.2: CIFAR-10 MEMORIZATION AND LABEL NOISE EXPERIMENTS
Revised: B. S ECTION 3.2: CIFAR-10 MEMORIZATION AND LABEL NOISE EXPERIMENTS

Original: [Figure]
Revised: [Figure]

Original: Memorization experiment In Figure 16, we report the test performance (dashed lines) corresponding to the curves presented in the main paper for the CIFAR-10 memorization experiment.
Revised: Memorization experiment In Figure 16, we report the test performance (dashed lines) corresponding to the curves presented in the main paper for the CIFAR-10 memorization experiment.

Original: The test performance with standard labels decreases slower than the training performance as the threshold increases, and they eventually reach the same value.
Revised: The test performance with standard labels decreases slower than the training performance as the threshold increases [Figure] and they eventually reach the same value.

Original: This is consistent with the hypothesis that by training on the consistent directions, the AND-mask selects the invariant patterns and prunes out the signals that are not invariant.
Revised: This is consistent with the hypothesis that by training on the consistent directions, the AND-mask selects the invariant patterns and prunes out the signals that are not invariant.

Original: Network architecture and training details
Revised: Network architecture and training details

Original: Each trial trains the ResNet “ FastResNet ” from the PyTorch-Ignite example 14 for 80 epochs on the full CIFAR-10 training set.
Revised: Each trial trains the ResNet “ FastResNet ” from the PyTorch-Ignite example 15 for 80 epochs on the full CIFAR-10 training set.

Original: We use the Adam optimizer with a learning rate of 5e ´ 4 , and a 0 . 1 learning rate decay at epoch 40 and 60.
Revised: We use the Adam optimizer with a learning rate of 5e ´ 4 , and a 0 . 1 learning rate decay at epoch 40 and 60.

Original: We ﬁx the batch size to 80.
Revised: We ﬁx the batch size to 80.

Original: We set up 14 trials by evaluating each of the AND-mask-thresholds t 0 , 0 . 05 , 0 . 1 , 0 . 2 , 0 . 4 , 0 . 6 , 0 .
Revised: We set up 14 trials by evaluating each of the AND-mask-thresholds t 0 , 0 . 05 , 0 . 1 , 0 . 2 , 0 . 4 , 0 . 6 , 0 .

Original: 8 u for two datasets: (a) unchanged CIFAR-10, (b) CIFAR-10 with the training labels replaced by random labels.
Revised: 8 u for two datasets: (a) unchanged CIFAR-10, (b) CIFAR-10 with the training labels replaced by random labels.

Original: Note that a threshold of 0 corresponds to not using the AND-mask.
Revised: Note that a threshold of 0 corresponds to not using the AND-mask.

Original: Each trial is run twice with separate random seeds.
Revised: Each trial is run twice with separate random seeds.

Original: Label noise experiment We trained the same ResNet as for the experiment above, once with and once without the AND-mask.
Revised: Label noise experiment We trained the same ResNet as for the experiment above, once with and once without the AND-mask.

Original: We ran each experiment with three different starting learning rates t 5e ´ 4 , 1e ´ 3 , 5e ´ 3 u and a learning rate decay at epoch 60.
Revised: We ran each experiment with three different starting learning rates t 5e ´ 4 , 1e ´ 3 , 5e ´ 3 u and a learning rate decay at epoch 60.

Original: The baseline worked best with a learning rate of 1e ´ 3 , while the AND-mask with 5e ´ 3 , likely to compensate for the masked out gradients.
Revised: The baseline worked best with a learning rate of 1e ´ 3 , while the AND-mask with 5e ´ 3 , likely to compensate for the masked out gradients.

Original: The AND-mask threshold that worked best was 0 . 2 , which is consistent with the results obtain in the experiment above
Revised: The AND-mask threshold that worked best was 0 . 2 , which is consistent with the results obtain in the experiment above.

Original: . B. S ECTION 3.3: B EHAVIORAL C LONING ON C OIN R UN
Revised: B. S ECTION 3.3: B EHAVIORAL C LONING ON C OIN R UN

Original: The target policy π ˚ is obtained by training PPO (Schulman et al., 2017) for 400M time steps using the code 15 for the paper Cobbe et al. (2019a).
Revised: The target policy π ˚ is obtained by training PPO (Schulman et al., 2017) for 400M time steps using the code 16 for the paper Cobbe et al. (2019a).

Original: This policy is trained on the full distribution of levels in order to maximize its generality.
Revised: This policy is trained on the full distribution of levels in order to maximize its generality.

Original: We use π ˚ to generate a behavioral cloning (BC) dataset, consisting of pairs p s, π ˚ p a | s qq , where s are the input-images ( 64 ˆ 64 RGB) and π ˚ p a | s q is the discrete probability distribution over actions output by π ˚ .
Revised: We use π ˚ to generate a behavioral cloning (BC) dataset, consisting of pairs p s, π ˚ p a | s qq , where s are the input-images ( 64 ˆ 64 RGB) and π ˚ p a | s q is the discrete probability distribution over actions output by π ˚ .

Original: The states are sampled randomly from trajectories generated by π ˚ .
Revised: The states are sampled randomly from trajectories generated by π ˚ .

Original: In order to test for generalization performance, the BC training dataset is restricted to 64 distinct levels.
Revised: In order to test for generalization performance, the BC training dataset is restricted to 64 distinct levels.

Original: We generate 1000 examples per training level.
Revised: We generate 1000 examples per training level.

Original: The test set consists of 2000 examples, each from a different level which does not appear in the training set.
Revised: The test set consists of 2000 examples, each from a different level which does not appear in the training set.

Original: A ResNet-18 ˆ π θ is trained to minimize the loss D KL p π ˚
Revised: A ResNet-18 ˆ π θ is trained to minimize the loss D KL p π ˚

Original: ||
Revised: ||

Original: ˆ π θ q .
Revised: ˆ π θ q .

Original: We ran two automatic hyperparameter optimization studies using Tree-structured Parzen Estimation (TPE) (Bergstra et al., 2013) oftrials each, with and without the AND-mask.
Revised: We ran two automatic hyperparameter optimization studies using Tree-structured Parzen Estimation (TPE) (Bergstra et al., 2013) oftrials each, with and without the AND-mask.

Original: The learning rate was decayed by a factor of 10 half-way at at 3 { 4 of the training epochs.
Revised: The learning rate was decayed by a factor of 10 half-way at at 3 { 4 of the training epochs.

Original: The “temporal” version of the AND-mask used for this experiment is reported in Algorithm 1. 
Revised: The “temporal” version of the AND-mask used for this experiment is reported in Algorithm 1. [Equation]

Revised: q 3 a Ð β 3 ¨ a ` p 1 ´ β 3 q ¨ elemwise_sign p g q 4 b Ð 1 r | a | ě τ s 5 θ Ð θ ´ α p m ˝

Revised: b q m ?

Revised: v ` (cid:15)

Original: In blue we highlight the additional lines compared to traditional Adam.
Revised: In blue we highlight the additional lines compared to traditional Adam.

Original: The threshold τ and βare hyperparameters that we included in the 1’024 trials of the search using Tree-structured Parsen
Revised: The threshold τ and βare hyperparameters that we included in the 1’024 trials of the search using Tree-structured Parsen

Original: Estimators.
Revised: Estimators.

Original: For the top 10 runs, hyperparameter values that were selected via the TPE search for the
Revised: For the top 10 runs, hyperparameter values that were selected via the TPE search for the

Original: AND-mask are the following.
Revised: AND-mask are the following.

Original: 1 m Ð β 1 ¨ m ` p 1 ´ β 1 q ¨ g 2 v Ð β 2 ¨ v ` p 1 ´ β 2 q ¨ p g ˝ g q 3 a Ð β 3 ¨ a ` p 1 ´ β 3 q ¨ elemwise_sign p g q 4 b Ð 1 r | a | ě τ s 5 θ Ð θ ´ α p m ˝

Original: b q m ?

Original: v ` (cid:15)

Original: [Table]
Revised: [Table]

Original: We found that applying weight decay as a second independent update after the AND-mask routine improved performance.
Revised: We found that applying weight decay as a second independent update after the AND-mask routine improved performance.

Original: To keep the comparison fair, we added this as a switch in the hyperparameter search for the Adam baseline as well, and it improved performance there as well.
Revised: To keep the comparison fair, we added this as a switch in the hyperparameter search for the Adam baseline as well [Table] and it improved performance there as well.

Original: [Figure]
Revised: [Figure]

Original: C A PPENDIX TO S
Revised: Adam AND-mask

Revised: [Figure]

Original: ECTION
Revised: Epoch [Figure] 3× [Figure]

Revised: s

Revised: [Figure]

Original: C. R ELATED WORK IN CAUSAL INFERENCE
Revised: D K L Adam AND-mask C A PPENDIX TO S ECTIONC. R ELATED WORK IN CAUSAL INFERENCE

Original: Causal graphs and causal factorizations The formalization of causality by directed acyclic graphs (Pearl, 2009) is a key element informing our exposition.
Revised: Causal graphs and causal factorizations The formalization of causality by directed acyclic graphs (Pearl, 2009) is a key element informing our exposition.

Original: According to such formalization, a causal model gives rise to each observed distribution.
Revised: According to such formalization, a causal model gives rise to each observed distribution.

Original: It is thereby possible to exploit properties of the causal factorization of the joint probability distribution over the observed variables.
Revised: It is thereby possible to exploit properties of the causal factorization of the joint probability distribution over the observed variables.

Original: Clearly, there are many ways to factorize a joint distribution into conditionals:
Revised: Clearly, there are many ways to factorize a joint distribution into conditionals:

Original: DAG’s help identifying the one that directly respects the underlying causal graph.
Revised: DAG’s help identifying the one that directly respects the underlying causal graph.

Original: An important feature of the causal factorization is that many of the conditionals, which we can think of as physical mechanisms underlying the statistical dependencies represented, are expected to remain invariant under interventions or changing external conditions.
Revised: An important feature of the causal factorization is that many of the conditionals, which we can think of as physical mechanisms underlying the statistical dependencies represented, are expected to remain invariant under interventions or changing external conditions.

Original: This postulate has appeared in various forms in the literature (Haavelmo, 1943; Simon, 1953; Hurwicz, 1962; Pearl, 2009; Schölkopf et al., 2012).
Revised: This postulate has appeared in various forms in the literature (Haavelmo, 1943; Simon, 1953; Hurwicz, 1962; Pearl, 2009; Schölkopf et al., 2012).

Original: Causal models and robust regression Based on this insight, it was proposed that regression based on causal features should presents desirable invariance and robustness properties (Mooij et al., 2009;
Revised: Causal models and robust regression Based on this insight, it was proposed that regression based on causal features should presents desirable invariance and robustness properties (Mooij et al., 2009;

Original: Schölkopf et al., 2012; Peters et al., 2016; Rojas-Carulla et al., 2018; Heinze-Deml et al., 2018;
Revised: Schölkopf et al., 2012; Peters et al., 2016; Rojas-Carulla et al., 2018; Heinze-Deml et al., 2018;

Original: Kügelgen et al., 2019).
Revised: Kügelgen et al., 2019; Parascandolo et al., 2017).

Original: In this view, the mechanisms can be considered as features of the patterns such that they support stable conditional probabilities.
Revised: In this view, the mechanisms can be considered as features of the patterns such that they support stable conditional probabilities.

Original: Thus learning the mechanisms may help achieve a stable performance across a number of conditions.
Revised: Thus learning the mechanisms may help achieve a stable performance across a number of conditions.

Original: Other works connecting causality and learning through invariances are (Subbaswamy et al., 2018; Heinze-Deml and Meinshausen, 2017), and perhaps – most related to our work – (Arjovsky et al., 2019): we presented a comparison with this method in the following section.
Revised: Other works connecting causality and learning through invariances are (Subbaswamy et al., 2018; Heinze-Deml and Meinshausen, 2017), and perhaps – most related to our work – (Arjovsky et al., 2019): we presented a comparison with this method in the following section.

Original: Causal regularization Recently (Janzing, 2019) showed that biasing learning towards models of lower complexity might in some cases be beneﬁcial for a notion of generalization from observational to interventional regimes.
Revised: Causal regularization Recently (Janzing, 2019) showed that biasing learning towards models of lower complexity might in some cases be beneﬁcial for a notion of generalization from observational to interventional regimes.

Original: Our proposed solution is however different, in that we only indirectly deal with penalizing model complexity, and rather focus on our proposed notion of consistency.
Revised: Our proposed solution is however different, in that we only indirectly deal with penalizing model complexity, and rather focus on our proposed notion of consistency.

Original: C. L EARNING INVARIANCES IN THE DATA
Revised: C. L EARNING INVARIANCES IN THE DATA

Original: Here we are going to compare ILC to other approaches for learning invariances in the data with neural networks, and in particular to Invariant Risk Minimization (IRM)
Revised: Here we are going to compare ILC to other approaches for learning invariances in the data with neural networks, and in particular to Invariant Risk Minimization (IRM)

Original: Arjovsky et al.
Revised: Arjovsky et al.

Original: The authors of IRM analyze a set up where minimizing training error might lead to models which absorb all the correlations found within the training data, thus failing to recover the relevant causal explanation.
Revised: The authors of IRM analyze a set up where minimizing training error might lead to models which absorb all the correlations found within the training data, thus failing to recover the relevant causal explanation.

Original: They consider a multi-environment setting and focus on the objective of extracting data representations that lead to invariant prediction across environments.
Revised: They consider a multi-environment setting and focus on the objective of extracting data representations that lead to invariant prediction across environments.

Original: While the high level objective is close to the one we focused on, the differences become clear when considering the deﬁnition of invariant predictors presented in Arjovsky et al.
Revised: While the high level objective is close to the one we focused on, the differences become clear when considering the deﬁnition of invariant predictors presented in Arjovsky et al.

Original: Deﬁnition 1 .
Revised: Deﬁnition 1 .

Original: A data representation Φ : X Ñ H elicits an invariant predictor w ˝
Revised: A data representation Φ : X Ñ H elicits an invariant predictor w ˝

Original: Φ across environments E if there is a classiﬁer w : H Ñ Y simultaneously optimal for all environments, i.e., wP arg min ¯ w : H Ñ y R
Revised: Φ across environments E if there is a classiﬁer w : H Ñ Y simultaneously optimal for all environments, i.e., w

Original: e p ¯ w ˝ Φ q @ e P E .
Revised: [Equation]

Original: In particular, the objective minimized by IRM is:
Revised: In particular, the objective minimized by IRM is:

Original: [Equation]
Revised: [Equation]

Original: where Φ are the logits predicted by the neural network and w is a dummy scaling variable (see
Revised: where Φ are the logits predicted by the neural network and w is a dummy scaling variable (see

Original: Arjovsky et al.
Revised: Arjovsky et al.

Original: The relevant part is the penalty term λ
Revised: The relevant part is the penalty term λ

Original: ¨ ›› ∇ w | w “ 1 . 

Revised: ¨ ›› ∇ w | w “ 1 .

Original: R e p w ¨ Φ q ›› 2 : One way to interpret it, is that the penalty is large on every environment where the distribution outputted by Φ could be made ‘closer’ to the distribution of the labels by either sharpening ( w ą 1 ) or softening it (i.e., closer to uniform w ă 1 ).
Revised: R e p w ¨ Φ q ›› 2 : One way to interpret it, is that the penalty is large on every environment where the distribution outputted by Φ could be made ‘closer’ to the distribution of the labels by either sharpening ( w ą 1 ) or softening it (i.e., closer to uniform w ă 1 ).

Original: Let us consider the example from IRM, where the authors describe two datasets of images that each contain either a cow or a camel:
Revised: Let us consider the example from IRM, where the authors describe two datasets of images that each contain either a cow or a camel:

Original: In one of the datasets, there is grass on 80% of the images with cows, while in the other dataset there is grass on 90% of them.
Revised: In one of the datasets, there is grass on 80% of the images with cows, while in the other dataset there is grass on 90% of them.

Original: IRM then makes the point that we can learn to ignore grass as a feature, because its correlation with the label cow is inconsistent (80% vs 90%).
Revised: IRM then makes the point that we can learn to ignore grass as a feature, because its correlation with the label cow is inconsistent (80% vs 90%).

Original: The setting we consider in this paper is slightly different: take our example from the CIFARexperiments.
Revised: The setting we consider in this paper is slightly different: take our example from the CIFARexperiments.

Original: Under our concept of invariance, we expect that (depending on the data generating process) even a single dataset where we treat every image as coming from its own ‘environment’ should be sufﬁcient to discover invariances.
Revised: Under our concept of invariance, we expect that (depending on the data generating process) even a single dataset where we treat every image as coming from its own ‘environment’ should be sufﬁcient to discover invariances.

Original: Drawing a connection to the setting from IRM, we would argue that the second dataset should not be necessary to learn that ‘grass’ is not ‘cow’.
Revised: Drawing a connection to the setting from IRM, we would argue that the second dataset should not be necessary to learn that ‘grass’ is not ‘cow’.

Original: If one treats every example as coming from its own environment, there is already sufﬁcient information in the ﬁrst dataset to realize that cows are not grass: Grass is predictive of cows only in 80% of the data, so grass cannot be ‘cow’.
Revised: If one treats every example as coming from its own environment, there is already sufﬁcient information in the ﬁrst dataset to realize that cows are not grass: Grass is predictive of cows only in 80% of the data, so grass cannot be ‘cow’.

Original: The actual cow on the other hand, should be present in 100% of the images, and as such it is the invariance we are looking for.
Revised: The actual cow on the other hand, should be present in 100% of the images, and as such it is the invariance we are looking for.

Original: Note that this is of course a much more strict deﬁnition of invariance
Revised: Note that this is of course a much more strict deﬁnition of invariance: If our dataset contains images labeled as ’cows’ but that have no cows within them, we might start to discard the features of cows as well.