[Paper Review] On the benefits of representation regularization in invariance based domain generalization

Changjian Shui, Boyu Wang+1 more|arXiv (Cornell University)|May 30, 2021

Domain Adaptation and Few-Shot Learning4 citations

TL;DR

This paper identifies that learning invariant representations alone in domain generalization is insufficient for robust generalization to unseen environments due to non-smooth representations. It proposes Jacobian matrix regularization to enforce Lipschitz continuity of the representation function, which theoretically bounds test error and empirically improves performance across multiple datasets and invariance criteria without modifying existing algorithms.

ABSTRACT

A crucial aspect in reliable machine learning is to design a deployable system in generalizing new related but unobserved environments. Domain generalization aims to alleviate such a prediction gap between the observed and unseen environments. Previous approaches commonly incorporated learning invariant representation for achieving good empirical performance. In this paper, we reveal that merely learning invariant representation is vulnerable to the unseen environment. To this end, we derive novel theoretical analysis to control the unseen test environment error in the representation learning, which highlights the importance of controlling the smoothness of representation. In practice, our analysis further inspires an efficient regularization method to improve the robustness in domain generalization. Our regularization is orthogonal to and can be straightforwardly adopted in existing domain generalization algorithms for invariant representation learning. Empirical results show that our algorithm outperforms the base versions in various dataset and invariance criteria.

Motivation & Objective

• To investigate why invariant representation learning alone fails in domain generalization under distribution shifts.
• To identify the root cause of poor generalization in unseen environments, particularly non-smooth or complex representations.
• To derive a theoretical upper bound on test error that depends on representation smoothness.
• To propose a practical regularization method—Jacobian matrix regularization—that enforces Lipschitz continuity of the representation function.
• To demonstrate that the proposed method improves performance across diverse datasets and invariance criteria without altering existing domain generalization frameworks.

Proposed method

• Theoretical analysis derives an upper bound on the test environment error in terms of the representation function's Lipschitz constant.
• It establishes that Lipschitz continuity is a sufficient condition to control the generalization error in unseen environments.
• The paper proposes Jacobian matrix regularization as a practical method to enforce the Lipschitz property of the representation function.
• The regularization is applied as a penalty on the Frobenius norm of the Jacobian matrix of the representation network.
• The method is orthogonal to existing invariance-based domain generalization algorithms and can be directly integrated into them.
• Empirical validation is conducted on multiple benchmarks (e.g., PACS, Office-Home, digits) using various invariance criteria and network architectures.

Experimental results

Research questions

• RQ1 Why does learning invariant representations fail to generalize to unseen environments despite minimizing source risks?
• RQ2 What theoretical conditions are necessary to ensure small test error in domain generalization?
• RQ3 How does the smoothness of the representation function affect generalization performance?
• RQ4 Can Jacobian-based regularization improve robustness in domain generalization without modifying the core invariance objective?
• RQ5 How does the proposed regularization compare to data augmentation in terms of implicit regularization of representation smoothness?

Key findings

• Theoretical analysis shows that the Lipschitz constant of the representation function directly bounds the generalization error in unseen environments.
• Non-smooth representations—especially complex or highly varying ones—are vulnerable to distribution shifts, even when invariance criteria are satisfied.
• Jacobian matrix regularization effectively controls the Lipschitz constant and improves robustness across multiple datasets and invariance criteria.
• Empirical results show consistent performance gains over base models across PACS, Office-Home, and digits datasets under different invariance objectives.
• The regularization is effective even when combined with strong data augmentation, suggesting it captures a distinct and complementary form of inductive bias.
• The method is universally applicable and can be seamlessly integrated into existing domain generalization frameworks without architectural changes.