Skip to content

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

🕒 Published (v1): 2025-01-01 · Source: NeurIPS · Venue: NeurIPS 2025 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

SBGen replaces the static image-text alignment used in VLM-based domain generalization with a stochastic Schrödinger Bridge (SB) process that progressively evolves domain-specific visual features toward domain-agnostic textual semantics via a time-conditioned SDE. The framework consists of text-guided feature selection, stochastic cross-domain evolution, and stochastic interpolation, yielding state-of-the-art results on both DG classification and domain-generalized semantic segmentation (DGSS).

Problem

Existing VLM-based domain generalization methods (e.g., prompt learning, cosine matching, adversarial regularization) perform static, one-shot alignment between domain-specific visual features and class-level text embeddings. This deterministic projection fails when the distribution discrepancy across domains is large, because degraded or spurious visual features cannot be corrected in a single mapping step and the semantic drift caused by domain shift is not explicitly modeled.

Method

SBGen frames the alignment from source-domain visual features to domain-agnostic text embeddings as a controlled stochastic process (Schrödinger Bridge).

  1. Domain-Aware Feature Selection (DFS): Given image features F ∈ ℝ^{H×W×C} from a CLIP/EVA-02 encoder and per-class text embeddings q^c, cosine similarity scores are computed for each spatial token. The top-k (k=30% by default) tokens per class are selected as semantically relevant starting points z₀ ∼ P^S.

  2. Stochastic Cross-Domain Evolution (SCE): A time-indexed SDE governs feature evolution: dz_t = f_θ(z_t, t) dt + √(2ε) dW_t The drift f_θ is an MLP conditioned on sinusoidal time embedding γ(t) and target class text embedding q^c, pushing features progressively toward the text distribution P^U.

  3. Stochastic Domain-Agnostic Interpolation (SDI): The SDE is discretized via Euler–Maruyama with T=5 steps. Terminal features z_T are written back into the original spatial feature map F, producing an enhanced map F′. A task head (linear classifier or Mask2Former-style pixel decoder) operates on F′ and {q^c}.

Training minimizes task loss + λ·KL(Q_θ ‖ P), where the KL term is approximated as Σ‖f_θ‖²Δt/4ε + ‖z_T − q^c‖².

Key Contributions

  • Cross-modal Schrödinger Bridge framework (SBGen) that models visual-to-text alignment as entropy-regularized stochastic optimal transport rather than static projection.
  • Three-stage differentiable pipeline (DFS → SCE → SDI) jointly trainable end-to-end with a task + SB consistency loss.
  • Theoretical generalization bound showing SBGen is provably tighter than static alignment baselines.
  • State-of-the-art on five DG classification benchmarks and two DGSS settings, with minimal added compute (+0.2 GPU-hours, +1.58M parameters on A100).

Results

Classification (leave-one-domain-out, CLIP ViT-B/16 backbone): - Average accuracy across PACS, VLCS, OfficeHome, DomainNet, TerraInc: 80.4%, outperforming second-best CLIPCEIL++ (NeurIPS 2024) by 1.3%. - VLCS: 86.7%, DomainNet: 64.4%, TerraInc: 63.5% (all SotA among VLM-based methods).

DGSS (EVA-02-L backbone, G→C/B/M setting): - mIoU: 71.24 / 62.26 / 71.91, avg 68.74 — outperforms second-best MGRNet (AAAI 2025) by +1.71 / +1.12 / +1.94 mIoU. - C→B/M: 66.03 / 77.90, avg 71.97, outperforms MGRNet by +1.33 / +1.47 mIoU.

Ablation (DGSS, G→C/B/M avg): - Baseline (tqdm): 66.05 → +DFS: 66.52 → +SCE: 67.28 → +SDI (full SBGen): 68.74. - Stochastic evolution outperforms direct cosine matching (+1.56 mIoU) and no time-indexed dynamics (+1.31 mIoU). - Optimal T=5 steps; K=0.3 feature selection ratio; λ=1 KL weight.

Limitations

  • Multi-step SDE simulation requires multiple forward passes per batch, increasing training complexity (though empirically only +0.2 GPU-hours vs. baseline).
  • Requires class-level text input; applicability to settings without predefined class names (e.g., open-vocabulary or instance segmentation) is not demonstrated.
  • Evaluated only on image classification and semantic segmentation; extension to other modalities (audio, video) is left as future work.
  • High-dimensional SB dynamics may not scale efficiently without approximation as feature dimensionality grows.

Relevance to Vision-Language Models

SBGen directly addresses a known weakness of VLM-based transfer: static CLIP-style alignment is brittle under large domain shifts because a single projection cannot account for the full distribution gap between in-domain visual features and text anchors. By reframing visual-to-text alignment as a stochastic process grounded in Schrödinger Bridge theory, this work offers a principled, theory-backed alternative to prompt tuning and cosine matching that operates on the feature level without modifying the pretrained encoders. The approach is backbone-agnostic (validated on CLIP, EVA-02, DINOv2) and connects the growing literature on score-based/diffusion alignment with practical VLM adaptation, suggesting that probabilistic trajectory modeling is a viable design axis for robust multimodal representation learning.