Weight Diffusion for Future: Learn to Generalize in Non-Stationary Environments

Sincerely thanks for your efforts in reviewing the paper. Below, we respond to your questions in detail.

Q1: Discussion with the two related papers [1, 2].

Thanks. Our method differs from [1, 2] as follows:

Focused problem: G.pt [1] focuses on supervised learning and reinforcement learning, while our W-Diff addresses the evolving domain generalization (EDG) in the domain-incremental setting. In EDG, distribution shifts often hinder models that are trained with full supervision on in-distribution (ID) data from generalizing to out-of-distribution data. Hence, our goal is to address the distribution shift, instead of generating diverse high-performance parameters for ID data.

D2NWG [2] focuses on transfer learning to provide better model parameter initialization for faster fine-tuning convergence on new datasets. However, target domains are unlabeled in EDG, preventing supervised fine-tuning. This makes D2NWG unsuitable for domain generalization. In contrast, our method generates model parameters applicable directly to the unlabeled target domain without fine-tuning.

Condition design for diffusion model: G.pt [1] collects the loss/error/return of task model checkpoints during training as the condition for the diffusion model. It is designed for a single dataset to which the training data belongs, thus struggling with distribution shifts.

D2NWG [2] uses CLIP to extract features for each sample and Set Transformer to generate dataset encoding from these features. The dataset encoding is used as the condition for diffusion model, while training set samples of a new dataset are required to obtain the dataset encoding. This is infeasible in unlabeled target domains.

Different from [1, 2], we use the classifier weights of historical domain (referred as reference point) and the prototypes of current domain as the condition of diffusion model to generate the difference of classifier weights between reference point and anchor point (i.e., the classifier weights of current domain). Considerations of such design are 1) the difference between reference point and anchor point denotes the evolution of parameters from historical to current domain, which conduces to modeling the crucial evolving pattern across domains in EDG; 2) the reference point provides initialization-like information, while current prototypes offer some information about the desired decision boundary, helping to explore the relationship between generated parameters and the given domain.

Moreover, we compare W-Diff with G.pt [1] in the following datasets. G.pt shows worse generalization on target domain, due to the limitations discussed above.

----Error rate comparison (%)---

Q2: The idea is ok, but still needs a careful discussion and ablation study.

Thanks. We have provided the results of ablation study in Table 4 of the PDF file. Concretely, we explore the following variants:

• Variant A ablates the consistency loss $\mathcal{L}^t_{con}$ for learning domain-shared feature encoder.

The performance drop of Variant A suggests that learning domain-invariant feature representations is necessary for EDG in the domain-incremental setting. Otherwise, the feature encoder could easily overfit to current domain, prohibiting task model from generalization.

• Variant B ablates the conditional diffusion model and directly uses the incrementally trained classifier for inference.

• Variant C directly uses the historical classifier weights in the reference point queue $Q_r$ to construct the average weight ensemble $\bar{\mathbf{W}}^{test}$ for inference: $\bar{\mathbf{W}}^{test}=\frac{1}{|Q_r|}\sum_{\ddot{\mathbf{W}}^{t^{\prime}}\in Q_r}\ddot{\mathbf{W}}^{t^{\prime}}$.

The inferior results of variant B and C indicate that W-Diff benefits from generating meaningful and customized classifier weights via controlling the condition of diffusion model.

Q3: The performance gain over different datasets looks marginal.

Thanks. To comprehensively evaluate the effectiveness of W-Diff, we have conducted a significance test (t-test) on different datasets. Concretely, a significance level of 0.05 is applied, and if p-value is less than 0.05, the accuracy difference between EvoS [8] and W-Diff is statistically significant. For clearer explanation, the -log(p) of each p-value has shown in red line. In Fig. 1(b) of the PDF file, majority of the -log(p) of the performance comparison between EvoS and W-Diff are larger than -log(0.05), which means W-Diff is statistically superior to EvoS at most datasets.

Besides, we also extend W-Diff to regression tasks, where the regression datasets and results are from DRAIN [4]. In Table 3 of the PDF file, W-Diff still works well. Overall, given the generalization performance on regression and text/image/multi-variate classification datasets, W-Diff is of more versatility.

Q4: More ablation studies, e.g., just leveraging the past classifiers for the ensemble prediction, or test-time adaptive classifier ensemble as in [3]?

Thanks. As you suggested, we have added more ablation study results in Table 1 of the PDF file.

• Variant C in the reply to Q2 denotes leveraging the past classifiers for the ensemble prediction.

• Variant D denotes using the APT-batch manner in [3], which first conducts unsupervised adaptation of classifier by back-propagating the gradient of entropy loss to source-trained classifier and then makes independent predictions on each batch.

• Variant E denotes using the APT-online manner in [3], which first updates source-trained classifier with the cumulative moving average of update directions from the first target data batch to current batch, and then makes predictions on current batch.

In Table 1 of the PDF file, W-Diff outperforms these variants, because they ignore the evolving pattern which is crucial for EDG.

Thanks for your efforts in reviewing the paper and the constructive comments. Below, we have tried our best to address your concerns in detail.

Q1: Can this method scale to larger networks and generate the entire network?

Thanks. Firstly, we have tried larger networks on the fMoW dataset by replacing the DenseNet-121 with DenseNet-161, DenseNet-169, DenseNet-201, respectively. The results are provided in Table 2 of the PDF file, where further performance improvements are obtained.

Secondly, our method can be extended to generate more parameters, but some modifications are needed to solve the different shapes of parameters from multiple layers and the training efficiency of diffusion models when the number of parameters to be generated is high. One possible solution to the shape difference is to flatten parameters from different layers and then concatenate all the flattened parameters into a vector. The generated parameters in a vector format are finally converted into their original shapes. As for the training efficiency, we can additionally train a VAE to encode the high-dimensional parameter vector into a low-dimensional embedding and then decode the embedding to reconstruct the parameter vector. Later, the diffusion model is trained in the low-dimensional embedding space and the decoder of VAE is used to recover parameters from the generated embedding via diffusion model in the inference stage. As stated in the Limitations section, we will leave this for a future work due to time constraints.

Thirdly, as mentioned above, generating the entire network is possible, but it is not a good choice for evolving domain generalization (EDG).

1. The massive parameters of the entire network require a large number of sequential source domains to accurately model the evolving pattern of the entire network.

2. Not all parameters are unshared. Shallow layers of the network usually extract general knowledge that can be shared, while deep layers are specific to tasks or datasets [5].

3. Previous EDG work [7] has provided theoretical evidence, showing that solely learning dynamic features is insufficient. [6] and [7] both use a static variational encoding network as the feature encoder to extract invariant features and train a domain-adaptive classifier. Hence, based on the experience and conclusion of previous works, we choose to only generate the parameters of task head, e.g., the classifier for classification tasks.

4. Moreover, in the paper, we compare our W-Diff with DRAIN [4] which leverages LSTM to learn the evolving pattern of the entire network. Experimentally, on 2-Moons and ONP datasets, the error rate of target domain is 3.2% (DRAIN) vs. 1.5% (W-Diff), 38.3% (DRAIN) vs. 32.9% (W-Diff). These results partially show that generating entire network is less effective for EDG.

Q2: Comparison with VAEs.

Thanks for your advice. We have added the results when replacing the diffusion model with VAE to generate classifier weights, conditioned on the reference point and prototypes. Please refer to variant F in Table 4 of the PDF file. The better generalization performance when using diffusion model to learn the evolving pattern of classifiers shows its superiority in modeling complex distributions and generating high-quality data, which has also been proven by many diffusion-based generation works. Besides, to reduce the difficulty of training diffusion models, we only generate the parameters of the classifier. The number of parameters (MB) for the conditional diffusion model is given in Table 6 of the PDF file, from which we see that the diffusion model is small and easy to train.

Q3: Comparison of training time with SOTA method.

Thanks. The training time complexity of our method mainly comes from the diffusion model. For simplicity, we take the U-Net with all convolutional layers as an example. Assuming that there are $L$ convolutional layers, the size of feature map in the $i$-th layer is $H_i \times W_i×C_i$ and the kernel size is $k_i\times k_i$. Then the time complexity of one forward pass is $\mathcal{O}(\sum_{i=1}^L H_i \times W_i \times C_i \times C_{i-1} \times k_i^2)$. Let $S$ denote the total time step in diffusion model. Then the time complexity of training diffusion models for $I$ iterations can be approximated as $\mathcal{O}(I \times S \times \sum_{i=1}^L H_i\times W_i \times C_i\times C_{i-1}\times k_i^2)$.

Moreover, in Table 5 of the PDF file, we compare the training time and GPU memory of our W-Diff, DRAIN [4], EvoS [8] and GI [11] on the RMNIST and Huffpost datasets. We acknowledge that our method has no significant advantage in terms of training time, due to the diffusion model. But this is not a limitation unique to our approach and most methods based on diffusion models have this limitation. As part of future work, we will investigate and try to address the limitation to further enhance the training efficiency.

In addition, it is worth mentioning that GI and DRAIN require huge computational resources during the training process, when they are applied on relatively large networks. Specifically, DRAIN needs to generate the entire network parameters and the fine-tuning stage of GI requires second-order gradients. On the Huffpost dataset with the backbone of DistilBERT-base and a batch size of 64, GI and DRAIN encounter the issue of GPU memory explosion. By contrast, our method utilizes diffusion model to generate only classifier weights, and as shown in Table 6 of the PDF file, the diffusion model is small to train, without the explosion of GPU memory when using the same batch size of 64. Overall, the main contribution of our work is providing a new perspective to address EDG in the domain-incremental setting via delicately tailoring the conditional diffusion model.

Q4: Include a notation table in the appendix.

Thanks for your advice. We will add a notation table in the revision to ease the burden of readers.

We are grateful for your efforts in reviewing the paper as well as your constructive comments. Below, we do our utmost to address your concerns.

Q1: The specific advantages that diffusion models offer for DG.

Thanks for your comment. Firstly, instead of learning a deterministic classifier like previous EDG methods [4, 8], we consider the classifier as a distribution, which conduces to improving the robustness and reducing miscalibration. Yet, the prior knowledge of distribution type is unknown and we only have the observed data (i.e., saved classifier checkpoints) from the unknown distribution. Fortunately, diffusion models are quite powerful in modeling complicated and unknown distribution based on observed data. Moreover, we have also tried other generative models, e.g., VAE, to generate the classifier weights. The results on RMNIST dataset are shown in the following table. We see that the generalization performance is worse when using VAE, which demonstrates the superiority of diffusion model in modeling complicated classifier distributions.

-----Accuracy (%) on RMNIST (K = 3)-----

Secondly, in the paper, we focus on the problem of evolving domain generalization (EDG), where the domain gradually evolves over time in an underlying pattern. In addition to the invariant feature learning, it is also critical to excavate the underlying evolving pattern across domains for better predicting the model status on future domains. Inspired by the fact that conditional diffusion models excel at generating specific images using additional information, we delicately design the condition including both the classifier weights of historical domain and the prototypes of current domain, and train the conditional diffusion model to generate the discrepancy of classifier weights between historical domain and current domain. The discrepancy represents the evolving of classifier across domains. In such way, we convey the information on classifier evolving from past to present to the conditional diffusion model.

Thirdly, during inference, we can inject the target data information through the prototypes in the condition to generate target-customized classifiers. Besides, since the diffusion model is stochastic in nature, multiple noise samplings can generate different classifiers, forming an ensemble, which leads to better robustness for generalizing on new domains.

Q2: The computational complexity of the diffusion model training.

Thanks for your advice. For simplicity, we take the U-Net with all convolutional layers as an example. Assuming that there are $L$ convolutional layers, the size of feature map in the $i$-th layer is $H_i \times W_i×C_i$ and the kernel size is $k_i\times k_i$. Then the time complexity of one forward pass is $\mathcal{O}(\sum_{i=1}^L H_i \times W_i \times C_i \times C_{i-1} \times k_i^2)$. Let $S$ denote the total time step in diffusion model. Then the time complexity of training diffusion models for $I$ iterations can be approximated as $\mathcal{O}(I \times S \times \sum_{i=1}^L H_i \times W_i \times C_i \times C_{i-1} \times k_i^2)$.

Moreover, in Table 5 of the PDF file, we compare the training time and GPU memory of our W-Diff, DRAIN [4], EvoS [8] and GI [11] on the RMNIST and Huffpost datasets. We acknowledge that our method has no significant advantage in terms of training time, due to the diffusion model. But this is not a limitation unique to our approach and most methods based on diffusion models have this limitation. As part of future work, we will investigate and try to address the limitation to further enhance the training efficiency.

In addition, it is worth mentioning that GI and DRAIN require huge computational resources during the training process, when they are applied on relatively large networks. Specifically, DRAIN needs to generate the entire network parameters and the fine-tuning stage of GI requires second-order gradients. On the Huffpost dataset with the backbone of DistilBERT-base and a batch size of 64, GI and DRAIN encounter the issue of GPU memory explosion. By contrast, our method utilizes diffusion model to generate only classifier weights, and as shown in the following table, the diffusion model is small to train, without the explosion of GPU memory when using the same batch size of 64. Overall, the main contribution of our work is providing a new perspective to address EDG in the domain-incremental setting via delicately tailoring the conditional diffusion model.

-----------------Number of parameters (MB) of conditional diffusion model $\mathcal{E}_{\boldsymbol{\theta}}$----------------

Q3: The evaluation is limited to classification tasks; applicability to other types of tasks, such as regression, remains unexplored.

Thanks for your advice. We have extended our method to the two regression datasets (i.e., House and Appliance) in DRAIN [4]. For regression tasks, the prototype is calculated at the domain level, i.e., the average of features in a single domain. And the cross-entropy loss is replaced with the mean squared error (MSE) loss and the Kullback-Leibler divergence in the consistency loss is replaced with MSE. Specifically, the results are provided in the following table, where our method still achieves better generalization performance.

-----------------Mean absolute error (MAE) for regression tasks----------------

Thanks a lot for your efforts in reviewing the paper. Below, we respond to your questions in detail.

Q1: Concerns about the assumption of dataset accessibility in the inference phase.

Thanks. Firstly, we would like to make it clear that for discriminative tasks, test data (i.e., target data in our paper) must have been given, when the model is used for testing. Thus, our method doesn't need to infer future data distributions. We can estimate the prototype matrix directly by feeding test data into the model to get their features and class probability predictions.

Secondly, previous evolving domain generalization (EDG) works [4, 6, 7, 11] are only evaluated on the next one target domain. Following [8, 10], we evaluate on the next $K$ target domains, hoping models can also generalize well on the target domain in the farther future. When $K=1$, it aligns with previous standard evaluations. Besides, the evaluation on each domain is conducted independently.

Finally, given that the whole data of a target domain may not come at once, we also offer a way for estimating the prototype matrix in a batch-data stream scenario, where the prototype matrix for the $j$-th target data batch is estimated using the cumulative moving average. In Fig. 1(a) of the PDF file, we provide the results on RMNIST and Huffpost when evaluating W-Diff in the batch-data stream scenario.

Q2: Why not include timestamp difference into the condition?

Thanks. In our method, timestamp difference information is implicitly considered, as domains evolve over time and arrive in chronological order during training. Besides, the difference between classifier weights of current and historical domains also contains implicit timestamp difference information. Moreover, due to different optimization starting points and stochastic gradient descent, it is unlikely that classifiers respectively trained for two domains are exactly the same. Hence, we do not explicitly include timestamp difference in the condition.

Nevertheless, as you suggested, we have also tried to explicitly include the timestamp difference $\Delta t$. Results on RMNIST are shown in Table 1 of the PDF file, where no obvious improvement is obtained, possibly due to reasons above.

Q3: The framework works very well experimentally but lacks formal theory.

Thanks. The gradient descent algorithm (GDA) is known to be theoretically supported, and we find the denoising process of diffusion models can be seen as a learnable GDA with momentum update and small noise perturbation. Let us consider the following gradient descent algorithm for optimizing the model with parameters $\Phi$ on dataset $D$ via minimizing loss $L$:

$\Phi_{i+1}=\Phi_i-\lambda \cdot \nabla_{\Phi}L(\Phi_i, D),\quad i=0,\ldots,T-1,$

where $\lambda$ controls the update step size.

Alternatively, we generate classifier weights via conditional diffusion model $\mathcal{E}_{\theta}$ with the following denoising process:

$W^t\_{s-1}=\frac{1}{\sqrt{\alpha_s}}\left(W^t\_s - \frac{\beta\_s}{\sqrt{1-{\bar{\alpha}}\_s}} \mathcal{E}\_\theta(W^t\_s,s,\ddot W^{t^{\prime}}, \mu^t)\right) + \sigma\_s \epsilon, \quad \epsilon \sim N(\mathbf 0,\mathbf I),$

where $s$ is diffusion step, $W^t_0$ is the overfitted classifier weights on the $t$-th domain $D^t$, $\mu^t$ is the prototype matrix of $D^t$, and $\ddot W^{t^{\prime}}$ is the saved overfitted classifier weights on historical domain $D^{t^{\prime}}, t^{\prime} < t$.

Then, we can reformulate the denoising process as

$W^t\_{s-1}=W^t\_s-\underbrace{\frac{\beta\_s}{\sqrt{\alpha\_s}\sqrt{1-{\bar{\alpha}}\_s}}}\_{\lambda} \mathcal{E}\_{\theta}(W^t\_s,s, \ddot W^{t^{\prime}}, \mu^t)+\underbrace{(\frac{1}{\sqrt{\alpha\_s}}-1)W^t\_s}\_{\text{momentum update}}+\underbrace{\sigma\_s \epsilon}\_{\text{noise perturbation}}.$

Compared with GDA, the conditional diffusion model can be viewed as learning a descent path of $\nabla_W L(W^t_s, D^t|\ddot W^{t^{\prime}}, \mu^t)$, so that conditioned on historical classifier weights and current prototype matrix, the diffusion model can directly generate classifier weights suitable for current domain.

Q4: Typos and grammar errors.

Thanks. We have proofread the paper carefully and corrected the errors.

Q5: Settings of $\beta_s$ in U-Net.

Thanks. Eq.(3) is obtained by using $\alpha_s=1-\beta_s$. And $\beta_s$ is set via the following code, and betas is fixed throughout training.

betas=(torch.linspace(1e-4 ** 0.5, 2e-2 ** 0.5, 1000) ** 2)

Q6: Explanation of Figure 3a.

Thanks. What we want to convey is that our generated classifier weights are diverse and generally high-performance. $W^7$ is obtained by fine-tuning the classifier on domain $\mathcal{D}^7$, which may be trapped in local optima. The parameter space is too big to fully explore during fine-tuning. It does not mean that only the area around $W^7$ is good. $\hat W^{7|5}$ generally performs well, meaning that our generated classifier weights are diverse and potentially cover better but unexplored area in the fine-tuning process.

Q7: What is Variant B in ablation study?

Thanks. Variant B ablates the conditional diffusion model and directly uses the incrementally trained classifier along with the learned domain-shared feature encoder for inference.

Q8: Could it be modified to predict future data distribution?

Thanks. Predicting future data distribution is possible by saving partial historical data to learn the evolving pattern at the data level, but it is not wise for discriminative tasks (e.g., classification), which requires further fine-tuning on generated instances to generalize discriminative models to target domain. Moreover, due to varied data types (e.g., image, text, multivariate in our experiments), data diffusion is cumbersome, which requires quite different architectures and lacks of universality. By contrast, our weight diffusion is more general for discriminative tasks on different datatypes.