Learning a Cross-Modal Schrödinger Bridge for Visual Domain Generalization

We appreciate the reviewer for the positive comments on the clear motivation, technical novelties, extensive experiments and strong performance. The below is a point-by-point response to provide more ablation studies, more detailed analysis on stochastic evolution paths and uncertainty modeling.

Q1: More ablation studies on static alignment baselines, such as direct cosine matching or prompt learning without time-indexed dynamics.

R: Per your suggestion, we've added the experiments when replacing the stochastic evolution with two static alignment methods, namely, direct cosine matching (DCM) and without time-indexed dynamics (w.o. TID). The experiments are conducted when using GTA as the source domain. The results in the attached table~show the stochastic evolution clearly outperforms both two methods, indicating its contribution to the overall performance.

Q2: More analysis of the stochastic evolution paths, such as visualizations of intermediate feature distributions or quantitative metrics like Wasserstein distance over time?

R: Many thanks for your valuable suggestion. Unfortunately, this year NeurIPS does not allow the submission of visualization in 1-page PDF or external links. After consulting the program chair, we are allowed to describe the visualization method and the outcome in text. We extract the intermediate feature embeddings from both the text embedding and image embedding over different time steps during the stochastic evolution. We then display them in the feature space by t-SNE visualization. It is observed that the intermediate features from the text embedding and the image embedding are gradually more uniformly mixed, indicating it is stable and domain meaningful. We also compute the Wasserstein distance over time. The results attached below~show that the Wasserstein distance metric is becoming smaller over time, indicating the feature trajectories are converging, smooth and stable. Both results are promised to be added in the revised manuscript.

Q3: Is it possible to extend the inference stage to quantify uncertainty, for example through sampling multiple evolution paths? Such an addition could enhance the robustness of predictions under domain shift.

R: To extend the robustness of inference under domain shift, we leverage the inherent stochasticity of the Schrödinger Bridge evolution to model uncertainty in the prediction stage. Specifically, we simulate multiple independent realizations of the stochastic process, forming a Monte Carlo ensemble of evolved features.

Monte Carlo Evolution Paths. Given the initial visual tokens $\{\boldsymbol{z} _{0}\} \sim P _{0}$, we simulate $M$ independent evolution paths by sampling distinct noise trajectories $\{\xi _{n}^{(m)}\} _{n=0}^{T-1}$, where for each path we have $m = 1, \ldots, M$. The terminal state of each path is computed via:

where each $\boldsymbol{z} _{t_n}^{(m)}$ is recursively computed from its own noise realization.

Prediction Mean and Uncertainty. We perform prediction using the decoder $\mathcal{D}$ over each sample path:

$\hat{y}^{(m)} = \mathcal{D}(\mathcal{F}'^{(m)}, \{\boldsymbol{q}_{c}\}),$

and estimate the predictive mean and variance:

$\bar{y} = \frac{1}{M} \sum_{m=1}^{M} \hat{y}^{(m)},$

$\mathbb{V}(y) = \frac{1}{M} \sum_{m=1}^M (\hat{y}^{(m)} - \bar{y})^2.$

The variance $\mathbb{V}(y)$ quantifies the epistemic uncertainty, providing a confidence-aware prediction strategy across domains.

As a result, the variance $\mathbb{V}(y)$ is used to report the uncertainty score. The lower variance refers to more confident predictions and less uncertainty across domains. We report the uncertainty score on the domain generalized semantic segmentation, when we conduct 5, 10, 15, 20, 25 and 30 Monte Carlo runs. The attached table~shows that the proposed method does not show clearly higher variance over the baseline, indicating it is robust to domain shift.

Should you have further comments or questions, we are happy to address during the author-reviewer discussion stage.

We appreciate the reviewer for the positive comments on the theoretical analysis and extensive experiments. The below is a point-by-point response to clarify the notations, experimental settings and other presentation issues.

Q1: Confusion between robustness and generalization.

R: We are sorry for the confusion. We will replace the use of robustness to generalization throughout the paper for consistency.

Q2: Clarity of Fig.2. Does it follow the domain generalization assumption?

R： We are terribly sorry for the presentation of Fig.2, which is less clear than intended to properly demonstrate the domain generalization setting. To clarify, as detailed in the experimental section, the data from the target domain(s) is only involved in the inference stage for testing, not fed into the model in the training stage.It follows the basic assumption of the domain generalization. We will use:

• put the data from the source/target domain into a box of solid/dashed line, respectively;

• use the solid/dashed arrows of the data propogation to represent the input in the training/inference stage, respectively, where only source/target domain is involved.

Q3： Notation system. $T$ has been defined three times.

R: Many thanks for your careful checking, and we are sorry for the confusion in the notation system. For clarity and for correctness:

• We will only use $T$ to denote the time indices.

• Instead, we will use $\mathcal{M}$ to denote the deterministic transport map to distinguish.

• We will replace $P^T$ by $P^U$ to denote the distribution of the unseen target domain.

• We will remove the continuous interval ($t \in [0,1]$) and only keep the discrete interval for consistency.

Q4: Missing explicit limitation discussion.

R: Thanks for your careful checking and we are sorry for the ignorance. We will add the following text to explicitly discuss the limitation of the proposed method, written as:

However, the proposed SBGen requires a simulation of multiple-step stochastic differential equation (SDE) for each batch, which additionally adds multiple forward passes and increases the complexity over the baseline. Still, it exhibits a good trade-off between the complexity and the clear performance improvement on unseen target domains.

Q5: Computational cost analysis.

R: Per your suggestion, we've accordingly compared the proposed method with the baseline in terms of the training time, parameter number and model size under the DGSS experimental setting. The attached table~shows that although the proposed method achieves an acceptable trade-off between computational cost and performance improvement over the baseline. Specifically, the increase of GPU hour is 0.2 hours, the parameter number increase is 1.58 million, and the model size increase is 0.01GB. The GPU hour refers to the A100 GPU hardware.

Q6: Stability, ideally justified by the variance of multiple runs.

R: Many thanks for your valuable suggestion. We would like to kindly remind the reviewer that, the result is an average of three independent runs for domain generalized segmentation. We therefore report the standard deviation of these three runs from both the proposed method and the baseline. As shown below, the proposed method shows stability across independent runs.

Q7: t-SNE visualization in classification is more meaningful than segmentation.

R: Thanks for your valuable suggestion. Unfortunately this year NeurIPS does not allow any PDF attachment or external link. After consulting the program chair, we describe the visualization method and the result. The classification experiments are conducted when PACS is used as the unseen target domain. The features from the last layer of the encoder are extracted and projected into the latent space for t-SNE visualization. It is observed that the per-class samples from our method are more uniformly distributed than the baseline. The updated figure will be incorporated into the revised version.

Should you have further comments or questions, we are happy to address during the author-reviewer discussion stage.

Thanks for your swift response. Should you find any other concerns that remain unaddressed, we are happy to address during the following week.

We appreciate the reviewer for acknowledging its novelties over existing methods to align the text and image features. The below is a point-by-point response to clarify the computational cost, hyper-parameter selection and to provide more experimental analysis.

Q1: Computational Cost.

R: Thanks for your insightful comments. We definitely agree with the reviewer that, as the proposed method simulates a multi-step SDE. To inspect its impact, we compute the training time, GFLOPs and model size between the proposed method and the baseline. The outcomes in the attached table show that, the proposed method only moderately increases the computation complexity but leads to a clear performance improvement, achieving a good trade-off between them.
Specifically, the increase of GPU hour is 0.2 hours, the parameter number increase is 1.58 million, and the model size increase is 0.01GB. The GPU hour refers to the A100 GPU hardware.

Q2: Sensitivity of hyper-parameter $k$.

R: We are sorry for the formatting of the supplementary material, which is less clear than intended. Specifically, we've conducted an experiment to test how the selection of $K$ impacts the performance of domain generalized semantic segmentation. The attached table shows that, a too small $K$ (e.g., 0.1) may not be able to select sufficient visual features that are domain-specific for the alignment between the domain-agnostic class text and the domain-specific visual features, which may under-fit the generalization representation. A too large $K$ (e.g., 0.4 and 0.5) may introduce more visual features that are not domain-specific in the representation learning, which may lead to over-fit and result in a slight performance drop.

Q3: Performance on multi-source DG for segmentation.

R: Thanks for your valuable suggestion. We've accordingly added an experiment setting of domain generalized semantic segmentation, where two synthetic domains are jointly used for training. The attached results~show that, compared with the cases when simply using either one of the source domains, SBGen leads to further gains by pooling them together.

Q4: Ablation Study on Classification.

R: Per your suggestion, we've accordingly added an ablation study on each component, tested on the classification benchmarks. As attached, each of the three components shows a similar contribution on classification accuracy.

Should you have further comments or questions, we are happy to address during the author-reviewer discussion stage.

We are glad to see that the weakness & questions in the review report have been addressed. We sincerely thank the reviewer for the continued evaluation and the time spent on our paper.

We respectfully respond to the two new concerns raised in the final comment, which were not listed in the original weaknesses or clarification questions.

Q1: The contribution of this line of work appears limited, as many generalization issues can be addressed by scaling model size & training data.

R: We appreciate the reviewer’s new concern, but respectfully disagree:

• Model scaling and data scaling is an independent research line and could not degrade the contribution and significance of domain generalization. Indeed, larger models and data help, but they are not always available. One of the example is in the medical imaging or remote sensing scenario, where each center (domain) can only collect a small amount of images. Besides, due to the privacy issue and data management policy, the images from some centers (domains) may not be available. As a result, the generalization to unseen domains is still in great demand [1, 2]. Another fundemental issue regarding data scaling is the distribution shift. Different datasets can have different distribution, and always encounter with challenges such as long-tailed distribution [3] and out-of-distribution [4]. As a result, generalization still matters and plays an important role, for both machine learning and data scaling [5].

As emphasized in recent machine learning and computer vision papers, domain generalization remains a crucial challenge under fixed resource and distribution shift settings.

[1] Nguyen, A. Tuan, Philip Torr, and Ser Nam Lim. "Fedsr: A simple and effective domain generalization method for federated learning." Advances in Neural Information Processing Systems 35 (2022): 38831-38843.

[2] Zhang, Ruipeng, et al. "Federated domain generalization with generalization adjustment." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023.

[3] Shao, Jie, et al. "DiffuLT: Diffusion for Long-tail Recognition Without External Knowledge." Advances in Neural Information Processing Systems 37 (2024): 123007-123031.

[4] Dong, Hao, et al. "Multiood: Scaling out-of-distribution detection for multiple modalities." Advances in Neural Information Processing Systems 37 (2024): 129250-129278.

[5] Zhang, Boxuan, et al. "What if the input is expanded in OOD detection?." Advances in Neural Information Processing Systems 37 (2024): 21289-21329.

• Our work provides a principled and lightweight solution, which has potential in the context of data scaling. Specifically, it formulates feature evolution via a Schrödinger Bridge, requiring no extra data and minimal tuning. This benefits the broader community by offering robust generalization under fixed foundation models, especially under the cases when re-training is infeasible.

Q2: Whether the method is applicable at larger scales.

R: We appreciate the reviewer's continued evaluation on scalability. We are sorry that the scaling of training data in the rebuttal stage is presented in a way less clear than intended, where the performance on larger scale has already been validated. Specfically, the original experiment in the submission is conducted on GTA5 (G) dataset as the source domain, with 9.4K images. However, we further test the situation on:

• SYNTHIA (S) dataset as the source domain, with $\sim$25K images.

• GTA5(G) and SYNTHIA (S) datasets jointly as the source domains, with a total of $\sim$34.4K images.

We respectively raise the reviewer's attention that the scales of $\times$2.5 and $\times$3.6 have both been tested, where the proposed method still shows a clear improvement over the foundation model baseline. It indicates that its generalization ability maintains on two much larger data scales.

Besides, technically, we would like to reminder the reviewer that the scaling ablity of the proposed method holds because:

• It builds directly upon pre-trained large-scale VLMs (CLIP, EVA-02) and operates on top of them without changing the backbone, making it naturally compatible with scaling.

• The proposed drift module is lightweight (an MLP over feature tokens) and SDE simulation is parallelizable. In practice, we apply it to multi-class semantic segmentation on full-resolution datasets with no performance bottlenecks.

We appreciate the reviewer for the positive comments on the technical novelties, extensive experiments and module design. The below is a point-by-point response to clarify the optimal transport method and the computational cost.

W1 & Q1: Details on the optimal transport solving method.

R: Thanks for your insightful question. Concretely, our method is an entropy-regularized optimal transport method by solving the stochastic differential equation, which is different from existing methods. The steps are described as follows.

First, we initialize the image feature $\boldsymbol{z}_0$ by the image encoder and the proposed Domain-aware Visual Feature Selection component, which is used as the starting point of the optimal transport.

Then, the image feature $\boldsymbol{z}\_{t}$ of each step of the stochastic evolution is updated by the time embedding $\gamma(t)$ and the drift term, computed as $\text{drift} = \text{mlp}( \boldsymbol{z}\_{t} + \boldsymbol{q}\_{c} + \gamma(t))$ where $\boldsymbol{q}\_c$ is the text embedding, and $\text{mlp}$ denotes the drift network.

Then, in every time step $t$, $\boldsymbol{z}\_{t}$ is updated by the stochastic differential equation. Specifically, based on the noise term (generated by a normal distribution $\mathcal{N}(0, I)$ to mimic the random permutation) and the drift term $\text{drift}$, it is computed as $z_{t+1} = z_{t} + \text{drift} \times dt + \sqrt{2 \epsilon dt} \times \mathcal{N}(0, I)$ where $dt$ is the time interval, and $\epsilon$ is a regularized parameter to control the intensity of the noise.

Then, we compute the $K-L$ divergence of the drift noise's sum square, which will be later used as a regularization to smooth the transport path. It is computed as $\text{KL term} = \frac{ (\text{drift}^2) \cdot dt }{4 \epsilon}$

After $T$ time steps, the final image feature $z_T$ is aligned with the text embedding $\boldsymbol{q}_c$ by the mean square error (MSE) loss, given by $\text{MSE loss} = \text{MSE}( z_T, \boldsymbol{q}_c )$

Theoretically, the final error regularized optimal transport (EROT) loss $\text{EROT loss}$ is a weighted sum of the MSE loss and the $K-L$ loss, given by $\text{EROT loss} = \text{MSE loss} + \epsilon \times \text{KL loss},$ where the $K-L$ loss is an average of all the $T$ time steps, given by $\text{KL loss} = \frac{1}{T} \sum_{t=1}^{T} \text{KL term}_t$

Notice that, the computation of the $K-L$ loss is approximated in practical, where the approximation details are mentioned in L208-209 in the main text.

We will enrich these details in the revised main text and supplementary material.

To further check if different methods have a different result, we further compare the proposed method with the method in [1] and another commonly-used method, namely, Sinkhorn. The attached results show that these methods achieve a very similar result. The proposed method shows a slight improvement, which may be explained that it is more tailored for the alignment between image and text embeddings.

We will add these results and discussions in the revised version.

Since [2] does not have publicly available source code, we will cite and discuss it in the related work.

Q2: Clarify the computational cost.

R: Per your suggestion, we've accordingly compared the proposed method with the baseline in terms of the training time, GFLOPs and model size under the DGSS experimental setting. The attached table shows that although the proposed method achieves an acceptable trade-off between computational cost and performance improvement over the baseline. Specifically, the increase of GPU hour is 0.2 hours, the parameter number increase is 1.58 million, and the model size increase is 0.01GB. The GPU hour refers to the A100 GPU hardware.

Should you have further comments or questions, we are happy to address during the author-reviewer discussion stage.