[Paper Review] SAND-mask: An Enhanced Gradient Masking Strategy for the Discovery of Invariances in Domain Generalization

Soroosh Shahtalebi, Jean-Christophe Gagnon-Audet+4 more|arXiv (Cornell University)|Jun 4, 2021

Domain Adaptation and Few-Shot LearningComputer Science37 references22 citations

TL;DR

This paper proposes SAND-mask, a continuous gradient masking strategy that enhances invariance discovery in domain generalization by jointly enforcing agreement in gradient direction and magnitude across domains. It outperforms prior methods, achieving a 6% accuracy gain on Colored MNIST and competitive results across DomainBed benchmarks by dynamically adapting masking based on gradient consistency.

ABSTRACT

A major bottleneck in the real-world applications of machine learning models is their failure in generalizing to unseen domains whose data distribution is not i.i.d to the training domains. This failure often stems from learning non-generalizable features in the training domains that are spuriously correlated with the label of data. To address this shortcoming, there has been a growing surge of interest in learning good explanations that are hard to vary, which is studied under the notion of Out-of-Distribution (OOD) Generalization. The search for good explanations that are extit{invariant} across different domains can be seen as finding local (global) minimas in the loss landscape that hold true across all of the training domains. In this paper, we propose a masking strategy, which determines a continuous weight based on the agreement of gradients that flow in each edge of network, in order to control the amount of update received by the edge in each step of optimization. Particularly, our proposed technique referred to as "Smoothed-AND (SAND)-masking", not only validates the agreement in the direction of gradients but also promotes the agreement among their magnitudes to further ensure the discovery of invariances across training domains. SAND-mask is validated over the Domainbed benchmark for domain generalization and significantly improves the state-of-the-art accuracy on the Colored MNIST dataset while providing competitive results on other domain generalization datasets.

Motivation & Objective

• To address the failure modes of existing gradient masking methods in domain generalization, particularly their sensitivity to optimizer momentum, initialization, and data noise.
• To develop a continuous, adaptive masking strategy that promotes invariance by ensuring agreement not only in gradient direction but also in magnitude across training domains.
• To eliminate the need for manual hyperparameter tuning of annealing schedules in OOD generalization by enabling automatic convergence from no masking to full SAND-masking.
• To improve generalization performance on out-of-distribution data without requiring access to a test set for hyperparameter adjustment.
• To provide a robust, scalable solution for learning invariant features that generalize across diverse data distributions.

Proposed method

• SAND-mask computes a continuous mask per network edge based on the agreement of gradients across different domains, using both direction and magnitude consistency.
• It applies a smoothed AND operation on gradients, replacing the discrete AND-mask used in prior work to reduce sensitivity to noise and optimization artifacts.
• The mask value is computed as a function of cosine similarity and magnitude agreement between gradients from different domains, ensuring stable and differentiable updates.
• The method dynamically evolves during training, starting with minimal masking and gradually increasing based on observed gradient consistency, emulating an implicit annealing schedule.
• The approach is integrated into standard backpropagation without modifying the loss function, making it compatible with existing training pipelines.
• It is evaluated using the DomainBed benchmark and compared against SOTA methods including IRM, REx, CDANN, and AND-mask.

Experimental results

Research questions

• RQ1 Can a continuous gradient masking strategy that enforces agreement in both gradient direction and magnitude improve invariance learning in domain generalization?
• RQ2 Does SAND-mask mitigate the failure modes of discrete AND-mask, such as sensitivity to optimizer momentum and data noise?
• RQ3 Can SAND-mask achieve superior generalization performance without requiring explicit hyperparameter tuning of annealing schedules?
• RQ4 How does SAND-mask perform on benchmark datasets like Colored MNIST and Spiral, especially in comparison to SOTA methods?
• RQ5 Does the dynamic, self-adapting nature of SAND-mask lead to better convergence and robustness than fixed or manually scheduled masking strategies?

Key findings

• SAND-mask achieves a 6% absolute accuracy improvement on the Colored MNIST dataset compared to the previous SOTA, reaching 62.3% ± 1.0.
• On the full DomainBed benchmark, SAND-mask achieves competitive performance, outperforming or matching SOTA methods across multiple datasets including Cifar-10, Cifar-100, and Sketch.
• SAND-mask significantly improves generalization on Colored MNIST, where the performance gap between AND-mask and SAND-mask is approximately 15%, indicating strong invariance learning.
• On the Spiral dataset, SAND-mask underperforms AND-mask (49.2% vs. 88.0%), suggesting limitations in capturing certain types of distribution shifts.
• The method demonstrates robustness to initialization and noise, avoiding the failure modes observed in the original AND-mask due to its continuous and differentiable design.
• SAND-mask effectively learns invariant features without requiring explicit annealing schedules, automatically adapting masking strength based on gradient consistency during training.