Balancing Invariant and Specific Knowledge for Domain Generalization with Online Knowledge Distillation

1 Introduction

The success of deep neural networks largely depends on the assumption that training (source domain) and testing (target domain) data are independently and identically distributed (i.i.d.). However, this assumption is often violated in realworld scenarios due to discrepancies between training and testing data, known as the domain shift problem, leading to significant performance degradation [Wang et al., 2022]. To address this problem, domain adaptation has been explored to transfer knowledge from source to target domains [Pan and Yang, 2009]. Unsupervised domain adaptation, in particular, leverages unlabelled data from target domains, thereby eliminating the need for target domain annotations [Xu et al., 2019]. Despite their effectiveness, unsupervised domain adaptation methods necessitate data collection and model tuning for each target domain, making them impractical in many situations [Yue et al., 2019]. Consequently, domain generalization (DG) has emerged as a prominent alternative. DG aims to learn a universal representation frommultiple labelled source domains, enabling robust generalization to unseen domains [Wang et al., 2022]. Existing approaches typically fall into three categories: data augmentation [Zhou et al., 2020], domain-invariant representation [Wang et al., 2022], and specialized training strategies [Zhao et al., 2024].

Knowledge distillation is a training strategy that has demonstrated effectiveness in DG [Wang et al., 2021; Huang et al., 2023]. However, most DG methods based on knowledge distillation focus on extracting domain-invariant knowledge, assuming that domain-specific knowledge impedes generalization [Lee et al., 2022]. This assumption does not always hold, as domain-invariant knowledge derived from a limited number of source domains may not generalize well to unseen domains [Zhang et al., 2023b]. Conversely, domainspecific knowledge from source domains is able to enhance DG performance when target domains share characteristics with particular sources (Figure 1). Nonetheless, simultaneously distilling invariant and specific knowledge presents two key challenges. First, the teacher must ensure that learning specific knowledge does not compromise the invariant knowledge it has already obtained. Existing approaches that rely on ensembles [Zhou et al., 2021] or batch normalization layers [Seo et al., 2020; Zhang et al., 2023b] without an independent extractor for invariant knowledge risk distorting the invariant knowledge while incorporating specific knowledge. Second, it is essential to balance distilled invariant and specific knowledge to prevent the distilled specific knowledge from undermining the distilled invariant knowledge, thereby ensuring the student's generalization performance when the source and target domains do not share characteristics. Additionally, current methods typically employ an offline distillation strategy, where the teacher remains fixed during the distillation process. This approach restricts the teacher from adapting to the student's evolving requirements during training, potentially resulting in suboptimal knowledge transfer.

To address these challenges, we propose Balanced Online knowLedge Distillation (BOLD), which distills both invariant and specific knowledge from a multi-domain expert teacher to a student, dynamically balancing their contributions using uncertainty weighting within an online distillation strategy. BOLD leverages adapter [Gao et al., 2024] techniques to construct a multi-domain expert teacher. Specifically, it integrates multiple adapters into a pretrained backbone, with each adapter specializing in capturing knowledge for a specific domain. This design separates invariant knowledge within the backbone from specific knowledge in the adapters, allowing the student to distill invariant knowledge from the backbone and specific knowledge from the corresponding expert adapter. To tackle the second challenge, BOLD incorporates Pareto optimization principles [Lin et al., 2019] and uncertainty weighting [Kendall et al., 2018] to ensure that both types of knowledge are optimized simultaneously without compromising either. Furthermore, BOLD employs an online distillation strategy, enabling domain experts to train concurrently with the student. During this process, the domain experts minimize the discrepancy between their and student outputs. This online approach enables the domain experts to dynamically adapt based on the student feedback throughout training, facilitating more effective knowledge transfer.

Our contributions are as follows: (1) We demonstrate that leveraging an appropriate optimization strategy effectively enhances the model's generalizability by distilling both invariant and specific knowledge. (2) We illustrate that adapting the teacher in response to student feedback using an online distillation strategy improves knowledge transfer and strengthens the student's generalizability. (3) We provide the theoretical insights that underline the importance of domainspecific knowledge and establish the rationale for utilizing uncertainty weighting. Experiments against state-of-the-art baselines validate the effectiveness of the BOLD framework.

2 RelatedWork

Domain Shift refers to the degradation in performance caused by discrepancies between the source (training) and target (testing) domains [Pan and Yang, 2009]. Domain adaptation has been proposed to address this issue by aligning the marginal [Baktashmotlagh et al., 2013] or conditional [Luo et al., 2020] distributions of the source and target domains or by fine-tuning models trained on source domains to adapt to the target domain [Long et al., 2015]. To reduce the cost associated with annotating target domain data, domain adaptation has been explored in semi-supervised [Saito et al., 2019] and unsupervised [Long et al., 2017] scenarios, utilizing partially labelled or unlabelled target domain data during training. However, these methods still rely on pre-collected target domain data, which presents a practical limitation, as obtaining such data is not always feasible [Yue et al., 2019]. This limitation highlights the need for approaches that are able to generalize to unseen domains without requiring target domain data collection in advance [Wang et al., 2022].

Domain Generalization (DG) was first introduced by [Blanchard et al., 2011] and later formalized by [Muandet et al., 2013]. Existing DG approaches primarily fall into three categories: data augmentation [Zhou et al., 2020], domaininvariant representation learning [Wang et al., 2022], and specialized learning strategies [Zhao et al., 2024; Pang et al., 2025]. Recently, knowledge distillation has attracted attention in the context of DG. [Wang et al., 2021] first proposed a gradient regularization method to regularize the domain-invariant knowledge distilled from the teacher. [Lee et al., 2022] introduced a self-distillation framework where a group of students collectively form a teacher, with each student distilling domain-invariant knowledge from the ensemble teacher. [Huang et al., 2023] leverages the text encoder of a Vision-Language model to distil domain-invariant knowledge. [Zhang et al., 2023b] suggested distilling domain-aware knowledge from a large pre-trained teacher model. Most existing methods focus exclusively on distilling domain-invariant knowledge, overlooking the significance of domain-specific knowledge [Seo et al., 2020; Bui et al., 2021]. Additionally, these methods typically employ an offline distillation strategy, where the teacher remains fixed after initial training. In contrast, our framework distills both invariant and specific knowledge using an online distillation strategy, allowing the teacher to adapt based on feedback from the student.

Knowledge Distillation was initially developed for model compression, with the goal of making the output of a smaller student model similar to that of a larger, existing teacher model [Hinton et al., 2014]. [Luo et al., 2016] demonstrated that training a student model using knowledge from a teacher via knowledge distillation can lead to better performance than direct training with one-hot ground truth labels. Knowledge distillation methods can be categorized into offline and online approaches, depending on whether the teacher is updated concurrently with the student [Gou et al., 2021]. In offline distillation, knowledge is transferred from a pre-trained teacher to a student, typically following a twostage training process [Zagoruyko and Komodakis, 2017;

Mirzadeh et al., 2020; Li et al., 2020]. Conversely, online distillation allows for the simultaneous updating of both teacher and student and supports an end-to-end trainable knowledge distillation framework [Anil et al., 2018; Zhang et al., 2018; Chen et al., 2020; Wu and Gong, 2021]. While offline distillation has proven effective in DG [Wang et al., 2021; Lee et al., 2022; Huang et al., 2023], the potential of online distillation remains unexplored. To our knowledge, this work is the first to explore how online distillation enhances DG.

3 Balanced Online Knowledge Distillation

This section begins with an overview of domain generalization and knowledge distillation. We then present Balanced Online Knowledge Distillation (BOLD) in three parts: 1) We describe how the teacher acquires specific knowledge and how the student distills invariant and specific knowledge from the teacher; 2) We explain how BOLD incorporates the Pareto optimization principle and uncertainty weighting to balance the distilled invariant and specific knowledge; 3) We discuss how the teacher dynamically adapts to student feedback, enhancing the knowledge transfer. Additionally, we provide theoretical insights that highlight the importance of domainspecific knowledge and justify the use of uncertainty weighting. Figure 2 illustrates the BOLD framework.

3.1 Preliminary

Notation. Let X denote an input feature space, with dimension d, and Y a target class label space. A domain, D, is composed of data sampled from a joint distribution P(X, Y ) on X × Y, where D = (xi, yi)ni=1 ∼ P(X, Y ), x ∈ X ⊂ Rd , y ∈ Y ⊂ R and n is the number of data in the domain. Here, X and Y denote the corresponding random variables [Zhou et al., 2022; Wang et al., 2022].

Domain Generalization. For the task of domain generalization, the input is N source domains (training set), S = {Dj | j = 1, · · · ,N}, where Dj = {(xji , yj i )}nj i=1 denotes the j-th domain and nj denotes the number of examples in j-th domain. The joint distributions between each pair of domains are different: P(X, Y )(j) ̸= P(X, Y )(k) , j ̸= k. The goal of domain generalization is to learn a robust and generalizable predictive function f : X → Y from the N source domains to achieve a minimum prediction error on an unseen target domain T , where T cannot be accessed during training and P(X, Y )(T ) ̸= P(X, Y )(j) for j ∈ {1, · · · ,N}.

Knowledge Distillation. Let T(x) and S(x) denote the outputs of the teacher and student models for a given input x. The knowledge distillation loss LKD is typically defined as the Kullback-Leibler (KL) divergence between the outputs of the teacher and student models: LKD = KL (T (x) ∥ S(x)).

3.2 Distilling Invariant and Specific Knowledge

Teacher Model. We adopt Contrastive Language-Image Pretraining (CLIP) [Radford et al., 2021] as the backbone for the teacher model, which includes both an image encoder and a text encoder. CLIP was selected for its strong generalization ability in associating images with their corresponding textual descriptions. For extracting invariant knowledge, the teacher leverages the pretrained image encoder without additional fine-tuning. To capture specific knowledge, we integrate adapters [Gao et al., 2024], a parameter-efficient tuning method, where each adapter specializes in a specific domain. Figure 2 illustrates that multiple domain-specific adapters are appended to the image encoder, with the number of adapters corresponding to the number of source domains. This design segregates invariant knowledge within the backbone and specific knowledge within the adapters. By ensuring the teacher preserves invariant knowledge while acquiring specific knowledge, the student is able to effectively distill invariant knowledge from the backbone and specific knowledge from the corresponding expert adapter.

We employ cross-entropy loss, Lce, for each expert adapter E. Unlike approaches that rely solely on maximizing similarity, cross-entropy allows us to maximize the similarity between an image and its ground-truth prompt while minimizing the similarity with unmatched prompts, ensuring comprehensive optimization [Radford et al., 2021]. For each class c, we generate m×N prompts in the format: "a picture of a {Dj }{ck }.", where m is the number of classes, Dj represents the j-th domain and ck represents the k-th class. The text encoder of the teacher model converts these prompts into text embeddings, yielding m text embeddings per domain, corresponding to the m classes. When processing an image from domain Di , the corresponding expert adapter Ei calculates the cross-entropy loss LE for each domain, as defined in Equation 1. Here, Timg denotes the image encoder of the teacher model, Ej and T j represent the expert and text embeddings for the j-th domain, respectively, where j ∈ {1, · · · ,N}. The similarity measurement, sim(·, ·), evaluates the similarity of image-text pairs, and we adopt cosine similarity following prior works [Radford et al., 2021].

After calculating LE for each domain, BOLD computes the domain loss Lidomain for expert adapter i by minimizing the loss for its corresponding domain while maximizing the loss for other domains, as outlined in Equation 2.

Student Model. To distill both invariant and specific knowledge from the teacher, we introduce two distillation losses: Invariant Distillation Loss (Linv) and Student-Specific Distillation Loss (Lsspc), as defined in Equations 3 and 4. The loss Linv minimizes the KL divergence between the outputs of the student and the image encoder of the teacher, while Lisspc minimizes the KL divergence between the student's outputs and the outputs of the relevant domain expert Ei corresponding to the domain of the input data. Since KL divergence is an asymmetric distance measure, the direction of distribution guidance is crucial. In our approach, the distribution of the teacher's output is used to guide the student's output when distilling knowledge from the teacher to the student.

Additionally, the student model learns independently by minimizing the cross-entropy of the given input. Equation 5 outlines the complete loss function.

The combination of Linv and Lsspc enables the student to capture both shared and unique features across domains, which is essential for enhancing the model's ability to generalize to unseen domains when the target domain shares features with some source domains. Furthermore, minimizing the divergence between student and teacher outputs acts as regularization, mitigating the risk of overfitting to the source domain data. Since the teacher's output represents a full probability distribution over all classes, the student learns to not only predict the correct label but also to approximate this distribution, thereby accounting for uncertainty. Additionally, reducing the divergence between student and teacher outputs enables the student to capture implicit information embedded in the teacher's soft outputs regarding inter-class relationships. This includes subtle correlations and patterns that are not evident through hard labels [Wang et al., 2021].

3.3 Balancing Invariant and Specific Knowledge

Simultaneously distilling invariant and specific knowledge into a single model presents a critical challenge: balancing the contributions of potentially conflicting losses. To address this, we leverage the principles of Pareto optimization [Lin et al., 2019], which suggest that when conflicts arise between multiple optimization objectives, a well-designed weighting strategy can optimize one objective without compromising the others. Specifically, we adopt uncertainty weighting [Kendall et al., 2018], introducing two learnable parameters, σinv and σsspc, to dynamically adjust the contributions of the invariant and specific distillation losses.

We begin by recalling the key definitions of Pareto optimality [Liang et al., 2021]. Let θ ∈ Θ represent the model parameters and consider n loss functions. Generally, it is infeasible to find a single θ that minimizes all losses simultaneously due to inherent conflicts among them. However, it is possible to identify a set of solutions known as Pareto optimal solutions, which balance these competing losses effectively.

Definition 1 (Pareto Dominance). Let θA, θB ∈ Θ represent two parameter vectors. We say that θA Pareto dominates θB (denoted as θA ≺ θB ) if li(θA) ≤ li(θB ), ∀i ∈ {1, 2, · · · , n} and li(θA) < li(θB ), ∃i ∈ {1, 2, · · · , n}.

Definition 2 (Pareto Optimality). A parameter vector θ∗ is Pareto optimality if there is no other parameter vector θ that dominates it. Formally, θ∗ is Pareto optimal if there does not exist θ such that θ ≺ θ∗ .

Definition 3 (Pareto Front). The Pareto front is the set of all Pareto optimal parameter vectors in the loss space, where each point corresponds to a unique parameter vector.

We address the simultaneous learning of invariant and specific knowledge within a Pareto optimization framework. Traditional Pareto-based methods, such as weighted sums, the ϵ-constraint technique, Chebyshev distance minimization, and evolutionary algorithms, can approximate the Pareto front but only rely on static or manually tuned weight assignments [Miettinen, 1999], which are difficulty to calibrate. To overcome this limitation, we employ uncertainty weighting to dynamically adjust the contributions of each loss function. This approach eliminates the need for repeated hyperparameter tuning and ensures robustness against gradient fluctuations. The updated loss function for the student is defined in Equation 6. Compared to directly learning the weights using a simple linear combination of losses, uncertainty weighting offers a significant advantage by preventing the weights from converging to zero [Kendall et al., 2018].

To justify our approach in relation to uncertainty weighting, we demonstrate that minimizing the KL divergence between student features and both invariant and specific features, corresponding to the invariant and specific distillation losses, is equivalent to maximizing a multi-task likelihood for these two objectives (Equation 7). This probabilistic interpretation aligns seamlessly with the uncertainty framework, providing a robust theoretical foundation. Here, zinv denotes the invariant features output by the teacher's backbone, zspc denotes the specific features output by the domain expert adapter, and zstudent denote the features output by the student.

3.4 Online Knowledge Distillation

In contrast to existing DG methods based on knowledge distillation that utilize a fixed teacher model, we implement an online distillation strategy, enabling the teacher to adapt to student feedback. To accomplish this, we introduce the Teacher-Specific Distillation Loss, Ltspc, defined in Equation 8 and integrate it into the teacher's learning objective, as illustrated in Equation 9. Unlike the Student-Specific Distillation Loss, the Teacher-Specific Distillation Loss leverages the student's output to guide the teacher's output. During training, only the domain expert corresponding to the input data's domain is updated, while the teacher's image encoder and domain experts for unrelated domains remain unaffected. Here, LiT denotes the loss for the domain expert associated with the i-th domain while Lidomain and Litspc are the domain and teacher-specific distillation losses for the i-th domain.

The online distillation strategy enables the teacher to adapt in real-time based on feedback from the student. Unlike fixed teacher models, which may become outdated as the student evolves, this dynamic adaptation ensures that the transferred knowledge remains relevant and continuously refined, resulting in more effective knowledge transfer. Moreover, the online distillation approach supports an end-to-end training process, eliminating the need for a separate training phase.

4 Experiments

We evaluate our approach using the DomainBed [Gulrajani and Lopez-Paz, 2021] benchmark across five datasets: PACS [Li et al., 2017], OfficeHome [Venkateswara et al., 2017], VLCS [Fang et al., 2013], Terra Incognita [Beery et al., 2018], and DomainNet [Peng et al., 2019]. Additionally, we assess performance on Digits [Zhou et al., 2020] and on NICO++ [Zhang et al., 2023a].

4.1 Experimental Results

Table 1 reports the average accuracy for all baselines across all benchmarks. The best performance is highlighted in bold, while the second-best performance is underlined.

Overall Average Accuracy Across Benchmarks. Table 1 highlights three key findings: (1) BOLD consistently outperforms state-of-the-art approaches, achieving the highest accuracy on the PACS, OfficeHome, VLCS, DomainNet, and NICO++ datasets, demonstrating its effectiveness in improving model generalizability to unseen domains. Notably, BOLD's strong performance on large-scale datasets such as DomainNet and NICO++ underscores its scalability. (2) Knowledge distillation-based methods (NKD, RISE, and BOLD) show weaker performance on the Terra and Digits datasets. This underperformance is attributed to limitations of the teacher model, CLIP, which performs poorly on these datasets. Consequently, students trained to mimic the teacher's outputs inherit these limitations. (3). Despite overall lower performance on Terra and Digits, BOLD surpasses NKD and RISE by a clear margin, achieving an improvement of approximately 5% on the Terra dataset.

Effectiveness Across Different Backbones. Table 2 presents evaluation results of BOLD, NKD, and RISE using different backbones across all benchmarks. The table includes evaluations of knowledge distillation from ResNet-50 and ViT-B/32 to ResNet-18 and from ResNet-50 to ResNet-50. Results for distilling knowledge from ViT-B/32 to ResNet-50 are shown in Table 1. These results demonstrate that BOLD consistently outperforms NKD and RISE across all backbones, highlighting its effectiveness in domain generalization.

4.2 Ablation Study

Effectiveness of Domain-Specific Knowledge and Online Distillation Strategy. Table 3 presents the ablation study results, validating the effectiveness of distilling domain-specific knowledge and employing the online distillation strategy. Here, Invariant represents the setup where the student distills only domain-invariant knowledge. Spcoffline refers to the setup where the student distills both types of knowledge but offline, meaning the teacher does not adapt to feedback from the student. Spconline refers to the setup where the student distills both types of knowledge online, allowing the teacher to adapt dynamically based on student feedback during training.

Based on the results in Table 3, we make three key observations: (1) When invariant knowledge is highly representative, the benefits of distilling specific knowledge are relatively minor. (2) When the target domain shares knowledge with the source domains, distilling specific knowledge results in substantial improvements, as observed in PACS and DomainNet. (3) The online distillation strategy effectively enhances knowledge transfer, particularly when the original teacher demonstrates limited capability, as observed in Terra.

Effectiveness of Uncertainty Weighting. Table 4 presents an ablation study evaluating the uncertainty weighting for balancing invariant and specific distillation losses. We compare it with three alternatives: static weighting (fixed weight of 0.5 for each loss), GradNorm [Chen et al., 2018], and ParetoMTL [Lin et al., 2019]. The results show that while GradNorm and ParetoMTL outperform Uncertainty Weighting in specific domains, Uncertainty Weighting consistently achieves either the best or second-best performance across all domains, leading to the highest overall performance. These findings underscore its robustness and effectiveness.

4.3 Further Analysis

Varying of σInv and σsspc. Figure 3 illustrates changes in σinv, σsspc, and their ratio σsspc/σinv. From these results, we make three key observations: (1) The optimal values of σinv and σsspc vary across datasets, making manual weight tuning impractical and highlighting the necessity for a dynamic weighting strategy. (2) A higher ratio of σsspc/σinv in datasets such as DomainNet and OfficeHome suggests that specific knowledge is more important in these contexts than in datasets like Terra, Digits, and VLCS. This observation aligns with Figure 4, where Terra is the only dataset exhibiting gradient conflict between invariant and specific distillation losses, while Digits and VLCS exhibit low gradient similarity. (3) The values of σinv consistently exceed those of σsspc, indicating that the invariant distillation loss contributes more to learning than the specific distillation loss. This finding supports the intuition that specific knowledge complements rather than dominates

5 Conclusion

Our Balanced Online Knowledge Distillation (BOLD) framework leverages both domain-invariant and domain-specific knowledge through an online distillation strategy to improve domain generalization. BOLD employs uncertainty weighting to dynamically balance loss contributions, eliminating the need for manual tuning and enhancing robustness across diverse datasets. The framework is supported by theoretical insights, providing a theoretical foundation for its design. Experiments and ablation studies validate the effectiveness of BOLD. Future work will explore developing more advanced distillation strategies to address limitations in the teacher model's capabilities and extend BOLD to more complex tasks, such as object detection and re-identification.

Acknowledgements

This research was partially supported by New Zealand MBIE Strategic Science Investment Fund (SSIF) Data Science platform - Time-Evolving Data Science / Artificial Intelligence for Advanced Open Environmental Science (UOWX1910).

invariant knowledge during training.

Knowledge Conflict. Figure 4 illustrates how cosine similarity between the gradients of invariant and specific distillation losses evolves over 50 epochs across all benchmarks. As training progresses, the gradient similarity converges. For the Terra dataset, the similarity converges to approximately -0.2, indicating a gradient conflict between invariant and specific distillation losses. In contrast, for the remaining benchmarks, the similarity converges to a positive value. These findings suggest that invariant and specific knowledge are not inherently conflicting; rather, their conflict is dataset-dependent, reinforcing the need for an appropriate balancing strategy.

T-SNE Visualization. Figure 5 presents the T-SNE visualization for ERM, NKD, RISE, and BOLD on PACS. As shown, distilling knowledge from a large teacher allows NKD, RISE, and BOLD to produce a more separable embedding space than ERM, highlighting the effectiveness of knowledge distillation. Furthermore, by incorporating specific knowledge, BOLD achieves an even more distinct and well-separated embedding space than NKD and RISE, demonstrating the potential of domain-specific knowledge for effective DG.

Imbalanced Dataset & Scalability. Imbalanced dataset distribution poses a practical challenge in providing sufficient training for domain experts. However, our framework employs lightweight adapters as domain experts instead of largescale neural networks. This design enables effective training even in domains with only a few hundred images. We also compare the parameter count across different backbones and expert adapters relative to the number of source domains.

We also gratefully acknowledge the Centre for Machine Learning for Social Good at the University of Auckland.