Bayesian Domain Invariant Learning via Posterior Generalization of Parameter Distributions

Thank you for your helpful question. We will correct the notations. $E[q(\omega|\theta_i)]$ means $E_{q(\omega|\theta_i)}[\omega]$, and $VAR[q(\omega|\theta_i)]$ means $VAR_{q(\omega|\theta_i)}[\omega]$. If we use Gaussian distribution as the variational distribution on each domain, then $E[q(\omega|\theta_i)]=\mu_i$,$VAR[q(\omega|\theta_i)]=\sigma^2_i$. However, if the variational distribution is other distributions, Equation 4,5 still hold. $q(\omega|\theta_i)$ means the variational distribution, where $\theta_i$ is the parameter of this distribution. For Gaussian variational inference, $\theta_i=(\mu_i,\sigma^2_i)$ and $q(\omega|\theta_i) = \mathcal{N}(\mu_i,\sigma^2_i)$ (or $\mathcal{N}(\omega;\mu_i,\sigma^2_i)$ for better understanding). $\theta_i$ is deterministic, $\omega$ is random.

Thank you for your helpful question. The computational costs are close to ERM, given that PTG only need a few more iterations (50 iterations, about 1.4 epochs) to aggregate the learned posteriors. Despite training three models, each model is trained on one domain representing 1/3 of the entire training data, so the overall computationally expensive is still same as ERM. The aggregation computations also have close forms. The training duration for ERM-Bayesian is approximately 44 minutes, followed by PTG training which spans approximately 20 minutes.

However, the GPU memory cost is expensive because PTG has to load the parameters of all models for training and computing the aggregated posteriors. We have addressed this phenomenon in Section 5. Simplified structure, parallel computation or simply making aggregations by CPU may be solutions, but they are future works.

Thank you for your helpful question. $\omega$ is the notation of the random-variable-sized parameter. If we want to refer to posterior parameter distribution, we use $p(\omega|\mathcal{D})$. $q(\omega|\theta)$ is the variational distribution of parameter, and $\theta$ is the true learnable object. For example, if we use Gaussian distribution as the variational distribution, then $\theta=(\mu,\sigma^2)$, $\omega$ follows $\mathcal{N}(\mu,\sigma^2)$, and $\mu,\sigma^2$ are the true learnable variable. We use subscript to denote all the objects that belong to a specific domain, expect for subscript 0. For example, $p(\omega|\mathcal{D}_i)$ means the posterior parameter distribution given domain $\mathcal{D}_i$. $q(\omega|\theta_i)$ means the variational parameter distribution on $\mathcal{D}_i$. $f_i()$ is the BNN featurizer on $\mathcal{D}_i$, and its variational parameter distribution is $q(\omega|\theta_i)$. During training, only $\theta_i$ is updated; during testing, we sample from $q(\omega|\theta_i)$ to form several networks, and take the average of their predictions as the final prediction. Please refer to [1] for more details about BNN. There's no domain $\mathcal{D}_0$, but we keep using $q(\omega|\theta_0)$ and $f_0$ to represent the target parameter distribution and target featurizer.

[1]Blundell C, Cornebise J, Kavukcuoglu K, et al. Weight uncertainty in neural network[C]//International conference on machine learning. PMLR, 2015: 1613-1622.

Thank you for your helpful questions. The discussions will be added to the paper.

Our contributions are: (1) We identified a new research direction to Domain Generalization, the analysis of parameter posterior distributions. (2) We proposed a relaxed assumption and a theory to estimate $p(\omega|\mathcal{D}^c)$ without access to $\mathcal{D}^c$. (3) We proposed two simple methods to approximate $p(\omega|\mathcal{D}^c)$, and the experiments showed the effectiveness. We hope our work can inspire further research.

We hope that these answers can show our ideas more clearly. We are looking forward to your feedback.

Thank you for your careful reading and question. We use Monte-Carlo sampling to sample from variational distributions. This is a common operation in variational inference BNN.[1]

[1]Blundell C, Cornebise J, Kavukcuoglu K, et al. Weight uncertainty in neural network[C]//International conference on machine learning. PMLR, 2015: 1613-1622.

Thank you for your helpful question. We initialize a BNN by a ResNet pretrained on ImageNet, which is the common choice in Domain Generalization. The method to initialize a BNN by DNN is adopted from [1]

[1]Krishnan R, Subedar M, Tickoo O. Specifying weight priors in bayesian deep neural networks with empirical bayes[C]//Proceedings of the AAAI Conference on Artificial Intelligence. 2020, 34(04): 4477-4484.

Thank you for your helpful question. We didn't improve BNN memory efficiency because our focus was on Domain Generalization rather than Bayesian deep learning. For easy understanding, we only provide PTG based on primary variational inference BNN [1]. There are many further works that focus on reducing the memory cost, such as [2,3]. As a result, there may be more solutions that can reduce the memory cost of PTG, but it's beyond this paper and could be future works. Our main concern is how to solve the problem of DG from the perspective of parameter posterior, and whether our solution is effective. We theoretically show how to estimate the posterior given invariant information by posteriors on each training domain. We further explain why the aggregated posterior can learn domain invariant information from the view of representation learning. A practical implementation is presented method based on VI, and the effectiveness is validated by experiments.

[1]Blundell C, Cornebise J, Kavukcuoglu K, et al. Weight uncertainty in neural network[C]//International conference on machine learning. PMLR, 2015: 1613-1622.
[2]Kristiadi A, Hein M, Hennig P. Being bayesian, even just a bit, fixes overconfidence in relu networks[C]//International conference on machine learning. PMLR, 2020: 5436-5446.
[3]Deng Z, Zhou F, Zhu J. Accelerated Linearized Laplace Approximation for Bayesian Deep Learning[J]. Advances in Neural Information Processing Systems, 2022, 35: 2695-2708.

Thank you for your insightful question. From Bayesian view, our aim is to learn $p(\omega|\mathcal{D}^c)$, which means the posterior parameter distributions that use all domain invariant information for prediction. There may be a misunderstanding that by "domain invariant posterior", we refer to $p(\omega|\mathcal{D}^c)$ rather than parameters that vary little among training domains. We made this description in Section 1 paragraph 4, and we will improve our expression for better clarity. We explain in Section 3.5 about the relationship between domain invariant parameters $p(\omega|\mathcal{D}^c)$ and its variation rate.

We also explain why domain specific features may still contain domain invariant information in Section 3.5. In Fig 2, there are stripes on feature maps. These strips are domain specific features because one is straight and one is wavy. If we use MSE to evaluate the distance between these features, they are far away (only few points overlap). We show by this toy model that this information can be extracted by parameters, but if we only want to extract domain invariant features, this information will be ignored. Fig 2 is a toy model that, for easy understanding, we assume the domain information exists in the feature map. However, we don't know exactly how the information is represented or how to get it, and we don't care. We only assume that $\mathcal{D}^c$ and $\mathcal{D}^v$ exist and are independent, and the domain distribution $p(\mathcal{D})$ can somehow be represented by them. One of our contributions is that we can estimate $p(\omega|\mathcal{D}^c)$ without access to $\mathcal{D}^c$.

Thank you for your helpful questions. The discussions will be added to the paper.

Our contributions are: (1) We identified a new research direction to Domain Generalization, the analysis of parameter posterior distributions. (2) We proposed a relaxed assumption and a theory to estimate $p(\omega|\mathcal{D}^c)$ without access to $\mathcal{D}^c$. (3) We proposed two simple methods to approximate $p(\omega|\mathcal{D}^c)$, and the experiments showed the effectiveness.

We hope our work can inspire further research. We hope that these answers can show our ideas more clearly. We are looking forward to your feedback.

Thank you for your insightful question. First, we didn't update classifier on each domain, and after the change of $f_0$, we have to update classifier to adapt it to $f_0$. Second, in theory PTG don't need to update $f_0$ after the aggregation. However, we are afraid that after some iterations, $f_i$ may get diversified too much that most posteriors in $f_0$ may have too large variance. We are not sure that whether a good generalization ability can be achieved in this stage, but too large variance may make a BNN divergent. As a result, we further update $f_0$ by just one iteration, so $f_0$ will not change a lot, and it can still extract meaningful features.

Thank you for your careful reading and question. Equation 4 is a common, simplified notation in variational inference BNN[1]. $\omega$ represents a vector of the whole parameters. We flatten all the parameters in a BNN and concatenate them to form $\omega$(just in concept, not in practice). Usually, we use a diagonal Gaussian variational distribution to approximate $\omega$. As a result, we can represent the mean and covariance of the variational distribution by vector $\mu$ and $\sigma^2$. In addition, BNN can still be trained by backward propagation, which means different parts of $\omega$ can be trained individually. For example, if we want to train a BNN convolutional layer, there are in fact two trainable convolutional layers: one represents $\mu$ and one represents $\sigma$. They have the same structure as the intended BNN convolutional layer. Please refer to [1] for more BNN training details.

[1]Blundell C, Cornebise J, Kavukcuoglu K, et al. Weight uncertainty in neural network[C]//International conference on machine learning. PMLR, 2015: 1613-1622.

There seems to by a misunderstanding of BNN. $\omega$ usually means the learnable parameters in DNN. However, in variational BNN, $\omega$ is the notation of the random-variable-sized parameter, $q(\omega|\theta)$ is the variational distribution of parameter, and $\theta$ is the true learnable object. For example, if we use Gaussian distribution as the variational distribution, then $\omega$ follows $\mathcal{N}(\mu,\sigma^2)$, and $\mu,\sigma^2$ are the true learnable variable. During training, we only learn $\mu,\sigma^2$ by Equation 4 and 5. During testing, we sample $\omega$ from $\mathcal{N}(\mu,\sigma^2)$ for 5 times to form 5 different models, and use the average of their predictions as the final prediction.

Thank you for your helpful question. As stated in Weakness 4, we don't have to specify these aspects. One of our contributions is to estimate $p(\omega|\mathcal{D}^c)$ without access to $\mathcal{D}^c$.

Thank you for your helpful question. There are two initialization procedure. First, we initialize a BNN by a ResNet pretrained on ImageNet, which is the common choice in Domain Generalization. The method to initialize a BNN by DNN is adopted from [1]. Second, after the ERM training, we initialize the BNN on each domain by the ERM trained BNN. We are sorry that we didn't draw a clear distinction between these prior models. We will modify the expression.-

[1]Specifying weight priors\mu in bayesian deep neural networks with empirical bayes.

We are sorry that we caused this confusion. Feature and parameters are just common defined. The domain invariant features are also same as common works. Domain invariant parameters (domain invariant posteriors) denote the posterior distribution of parameters given domain invariant information $p(\omega|\mathcal{D}^c)$. We give the definitions in Section 1 paragraph 4 and section 3.1, and we will add the definition by definition formula.

Thank you for your insightful question. We train the initial network by ERM just on all the given training domains. There's no additional training dataset. The ERM initialization process is part of our method, and since ERM introduces no generalization techniques, we think the improvements comes from our work.

Thank you for your careful reading. We will modify these problems. The result of PTG on ResNet 18 without prior network is:

Thank you for your insightful question. In Fig 2, there are stripes on feature maps. These strips are domain specific features because one is straight and one is wavy. If we use MSE to evaluate the distance between these features, they are far away (only few points overlap). We show by this toy model that this information can be extracted by parameters, but if we only want to extract domain invariant features, this information will be ignored. Fig 2 is a toy model that, for easy understanding, we assume the domain information exists in the feature map. However, we don't know exactly how the information is represented or how to get it, and we don't care. We only assume that $\mathcal{D}^c$ and $\mathcal{D}^v$ exist and are independent, and the domain distribution $p(\mathcal{D})$ can somehow be represented by them. One of our contributions is that we can estimate $p(\omega|\mathcal{D}^c)$ without access to $\mathcal{D}^c$.

Thank you for your valuable suggestions. Our comparative analysis remains contemporary considering that more than half of the models(8/14) are introduced after 2021, and the rest are classical counterparts. We selected our competitors for a fair comparison. Model selection plays a vital role in domain generalization, so we only compare with models available on DomainBed, which is a fair test platform that automatically select models. We did not compare PTG with many other SOTA methods, because they either use CLIP pretrained model, SWAD weight averaging strategy, large-scale ensemble learning or other strategies.

We are sorry that we can't combine PTG with EoA, which is the best model in the benchmark NICO++. Because EoA is an ensemble model, while PTG requires a single model as its prior. However, we want to mention that although CORAL is old, it still gets the third performance in NICO++ and there is only a small distance.(EoA 78.88, MIRO 78.65, CORAL 78.38). As you suggested, we provide the result of the best two models on NICO++, but please note we still give results on DomainBed, and all models use ResNet50. The results of MIRO and EoA are collected form [1]. Please refer to Section 4.1 paragraph 3 and Appendix G for more details.

We admit that EoA is better than PTG, but we want to remind again that EoA is an ensemble method while PTG just has one model in the testing stage. We already provide more combinations in Appendix H, and we are sorry that we can't provide the combination with MIRO due to the time limit.

[1] Arpit D, Wang H, Zhou Y, et al. Ensemble of averages: Improving model selection and boosting performance in domain generalization[J]. Advances in Neural Information Processing Systems, 2022, 35: 8265-8277.

Thank you for your insightful questions. We are sorry that we caused some misunderstanding.

The key misunderstanding is, $\mathcal{D}^c$ and $\mathcal{D}^v$ are not dataset! The second misunderstanding is, it is $p(\mathcal{D}^c)$ that remains constant, rather than $\mathcal{D}^c$.

$\mathcal{D}^c$ and $\mathcal{D}^v$ are two abstract concepts: $\mathcal{D}^c$ means the collection of all domain invariant information from a domain $\mathcal{D}$, and $\mathcal{D}^v$ is the opposite. $\mathcal{D}$ refers to a domain rather than a dataset, $p(\mathcal{D}^v)$ is its distribution, and the corresponding dataset are samples from $p(\mathcal{D}^v)$. We use $\mathcal{D}_i$ to distinguish between different domains. We assume that $\mathcal{D}^c$ and $\mathcal{D}^v$ are two independent sufficient statistics of $\mathcal{D}$. To be specific, there exist a bijection between $p(\mathcal{D})$ and $p(\mathcal{D}^c)p(\mathcal{D}^v)$, and $p(\mathcal{D}^c)$ and $p(\mathcal{D}^v)$ can be inferred by the dataset of $\mathcal{D}$. However, we don't care about (1) the specific algebra expression of $\mathcal{D}^c$, (2) the specific bijection formula, (3) the specific formula to estimate $p(\mathcal{D}^c)$ from dataset. One of our contributions is to estimate $p(\omega|\mathcal{D}^c)$ without access to $\mathcal{D}^c$.

In your example, $\mathcal{D}$ means a domain of dogs, and $\mathcal{D}_i$ may refer to a specific kind of dog, like golden retriever. $\mathcal{D}^c$ means the common information that all dogs share, such as the overall body shape, the facial patterns, the teeth and so on. $\mathcal{D}^v_i$ means the specific information that golden retriever have, such as the color, the hair length. As stated before, we don't care how the information is represented or how to get it, we just assume that it exists. We can directly estimate the posterior parameter distributions that use the common information for recognition. This is why we claim that our assumption is more relaxed than domain invariant features.

We explained our idea in Section 1 paragraph 5, Section 3.1 and Appendix A. Thanks again for your helpful question, we will add these discussions and your example in the paper.

Thank you for your insightful question. One of our contributions is that we are the first one to align parameter posteriors for domain generalization. As introduced in Section 1 paragraph 3 and Section 2.3, all former works use BNN to show the distribution of latent features, and they align the feature distributions similar as other Non-Bayesian ways. We find that the alignment of parameter distributions can also bring improvements. Besides, we propose a method to estimate the posterior given domain invariant information $p(\omega|\mathcal{D}^c)$, which can achieve the best domain generalization ability in theory. There's a problem that $\mathcal{D}^c$ is unknown, but we can still estimate $p(\omega|\mathcal{D}^c)$ without access to $\mathcal{D}^c$. In addition, although we propose the way to estimate $p(\omega|\mathcal{D}^c)$, we solve some difficulties in implementation, such as disordered dimensions of parameters. We explained these difficulties in section 3.3 and 3.4. We propose two simple yet effective implementations to solve these problems, and experiments show their superiority. In conclusion, we identify a new research direction to Domain Generalization, the analysis of parameter posterior distributions, and we propose the basic framework for it.

Thank you for your helpful suggestion. **Maybe there is a misunderstanding between $\int P(X|Y=y)P(Y=y)dy$ and $\int P(X|Y=y,Z)P(Y=y)dy$. **Marginalization usually means the process that requires summing over the possible values of one variable to determine the marginal contribution of another, to be specific, $P(X)=\int P(X,Y=y)dy=\int P(X|Y=y)P(Y=y)dy$. However, our theorem, $p(\omega|\mathcal{D}^c)=\int p(\mathcal{D}^v)p(\omega|\mathcal{D}^c,\mathcal{D}^v)d\mathcal{D}^v$, follows another Bayesian principle. The proof of Theorem 3.1 is shown in appendix B. Although $p(\omega|\mathcal{D}^c)=\mathbb{E}_{p(\mathcal{D}^v)}[p(\omega|\mathcal{D}^c,\mathcal{D}^v)]$ appears to be true, it only holds under the condition that $\mathcal{D}^c$ and $\mathcal{D}^v$ are independent. Besides, Theorem 3.1 shows that we can estimate $p(\omega|\mathcal{D}^c)$ by the posterior given full information, which plays a vital role in the following context.

We are sorry that we made the confusion. In Bayesian Neural Networks, posterior usually refers to the posterior distribution of parameters given data $p(\omega|D)$. We didn't make new definitions for parameter and feature. Parameter means model parameters and feature means latent features. We sample for 5 times from the variational distribution $q(\omega|\theta)$, and $q(\omega|\theta)$ is the approximation of the true parameter posterior $p(\omega|D)$. The dimensionality of posterior is the dimensionality of parameter. For example, the posterior of a 3X3X3 conv layer in a BNN has the dimensionality 3X3X3. The variational distribution $q(\omega|\theta)$ also share the same dimensionality as $p(\omega|D)$. The principle of BNN and the procedure of variational inference is shown in Section 2.3, 2.4 and 3.1, and we will add more descriptions.

Thank you for your helpful questions. The discussions will be added to the paper. We hope that these answers can express our ideas more clearly. We are looking forward to your feedback.

Thank you for providing suggestions for the paper. The responses may be too long, so we summarize each respnese into one sentence.

weakness 1: Parameter means model parameters, feature means latent features and during testing we sample from parameter distribution.

weakness 2: Theorem 3.1 is not marginalization, and it plays a vital role in the following context.

weakness 3: We are the first one to align parameter posteriors for domain generalization, exploring a new research direction.

weakness 4: $\mathcal{D}^c$ and $\mathcal{D}^v$ are not datasets.

weakness 5: Our experiment remains contemporary, and we add some suggested competitors.

mino : we add the suggested ablation study, and we will fix typos.

question 1: ERM is done on all the given training domains. There's no additional training dataset.

question 2: Feature and parameters are just common defined. Domain invariant parameters means $p(\omega|\mathcal{D}^c)$.

question 3: Both two initialization procedures exist.

question 4: We don't have to specify these aspects to build PTG.

question 5/6: Sorry the responses can't be simplified

We are looking forward to your additional feed back. Please feel free to ask any question, and we will reply immediately.

Thanks so much to the authors for the further detailed elaborations and significant effort during the author response period. The authors have now addressed most of my concerns. For now, I am happy to raise my score to 6. Meanwhile, I would still like to see a fairer comparison when the baseline methods are empowered with ERM initialization (to have the same initialization as the proposed method). However, as indicated by the authors, ERM initialization is also a part of their contribution. I would like to suggest the authors include it in the ablation studies.

I might need more time to think about it and re-evaluate the paper again. If the authors agree to provide the ablation studies (including comparison with the baselines empowered with ERM initialization) in the camera-ready version and my re-evaluation ends up telling me that it is not a big concern, I will further raise my score to 8. I will finalize my score during the discussion period with ACs. Thanks to the authors for their patience.

I would like to also mention, that proposing the alignment on the model parameter posteriors across domains is very novel, to other reviewers and AC. It offers a new baseline to enhance domain generalization using the core concept from BNN, which is significant to the machine learning community in general.

Weakness 1

Our major contribution is that we propose a method to estimate $p(\omega|\mathcal{D}^c)$ without $\mathcal{D}^c$. The analysis of parameter distributions has been widely studied so far, and variational inference is a simple way to align parameters. Besides Bayesian theorem, we also solve problems that may occur during implementation. Furthermore, our experiment can prove that our effort works.

Weakness 2

Although PTG-Lite is not BNN, its principle is still Bayesian. It is Theorem 3.1, which shows how to get the domain invariant posterior $p(\omega|\mathcal{D}^c)$ with only access to $\mathcal{D}$, that matters, rather than way to approximate $p(\omega|\mathcal{D}^c)$. PTG-Lite uses deterministic values to approximate $p(\omega|\mathcal{D}^c)$, but the key idea is still to estimate the invariant posterior.

Weakness 3

We explained why $\mathcal{D}^c$ and $\mathcal{D}^v$ should be independent in Appendix A. The proof of Theorem 3.1 is shown in Appendix B. Although Theorem 3.1 appears to be always true, in fact it holds under the condition that $\mathcal{D}^c$ and $\mathcal{D}^v$ are independent. Besides, Theorem 3.1 shows that we can estimate $p(\omega|\mathcal{D}^c)$ by the posterior given full information, which plays a vital role in the following context.

Weakness 4

We didn't report the mentioned approaches because they are not appliable on DomainBed. For fair comparison, we only compare with models on DomainBed.

Weakness 5

We will correct typos and improve the writing.

Question 1

We don't need to specify which information from the data is invariant. In BNN, we care how to estimate $p(\omega|\mathcal{D})$ rather than what's the data distribution $p(\mathcal{D})$. Similarly, we care how to estimate $p(\omega|\mathcal{D}^c)$, rather than $p(\mathcal{D}^c)$. $\mathcal{D}^c$ means any information that keeps invariant between domains, and $p(\omega|\mathcal{D}^c)$ is parameter posterior distribution that use $\mathcal{D}^c$ for prediction.

Question 2

These are typos. $E[q(\omega|\theta_i)]$ means $E_{q(\omega|\theta_i)}[\omega]$, and $VAR[q(\omega|\theta_i)]$ means $VAR_{q(\omega|\theta_i)}[\omega]$.