Gradient-Guided Annealing for Domain Generalization

Figure 1. (left) Decision boundaries of a 4 th -degree polynomial logistic regression model with 2D input. In this example, feature x1 is class-specific and x2 is domain-specific, while color represents classes and shapes represent domains. The samples with solid red and green colors are included in the training data, whereas the fainted samples are part of the hidden held-out test set. As a result, domain shift is represented by a change in x2. Although the classifier should only infer based on x1, traditional gradient descent leads to overfitting (top-left). The proposed method, GGA (bottom-left), introduces an annealing process that depends on gradient agreement, leading to models that generalize well to new, unobserved target domains. (right) Schematics of the parameter updates of ERM (top-right) and GGA (bottom-right). Parameters updated via ERM are driven by gradient conflict, whereas GGA searches for a point where gradients align before continuing descending towards a minima.