Understanding Hessian Alignment for Domain Generalization
Sobhan Hemati Guojun Zhang Amir Estiri Xi Chen
Out-of-distribution (OOD) generalization is a critical ability for deep
learning models in many real-world scenarios including healthcare and
autonomous vehicles. Recently, different techniques have been proposed to
improve OOD generalization. Among these methods, gradient-based regularizers
have shown promising performance compared with other competitors. Despite this
success, our understanding of the role of Hessian and gradient alignment in
domain generalization is still limited. To address this shortcoming, we analyze
the role of the classifier's head Hessian matrix and gradient in domain
generalization using recent OOD theory of transferability. Theoretically, we
show that spectral norm between the classifier's head Hessian matrices across
domains is an upper bound of the transfer measure, a notion of distance between
target and source domains. Furthermore, we analyze all the attributes that get
aligned when we encourage similarity between Hessians and gradients. Our
analysis explains the success of many regularizers like CORAL, IRM, V-REx,
Fish, IGA, and Fishr as they regularize part of the classifier's head Hessian
and/or gradient. Finally, we propose two simple yet effective methods to match
the classifier's head Hessians and gradients in an efficient way, based on the
Hessian Gradient Product (HGP) and Hutchinson's method (Hutchinson), and
without directly calculating Hessians. We validate the OOD generalization
ability of proposed methods in different scenarios, including transferability,
severe correlation shift, label shift and diversity shift. Our results show
that Hessian alignment methods achieve promising performance on various OOD
benchmarks. The code is available at
\url{https://github.com/huawei-noah/Federated-Learning/tree/main/HessianAlignment}.