Skip to content

Demystifying Multimodal Biomolecular Co-design with Intrinsic Geodesic Coupling

Conference: ICML 2026
arXiv: 2606.01628
Code: TBD
Area: Scientific Computing / Biomolecular Co-design / Multimodal Generation / Optimal Transport
Keywords: Biomolecular Co-design, Temporal Coupling, Optimal Transport, Bayesian Optimization, Flow Matching

TL;DR

The authors remodel the co-generation problem of heterogeneous modalities—"sequence + 3D structure"—as a Temporal Optimal Transport (TOT) problem. Using bi-level optimization and a Gaussian Process surrogate (GeoCoupling), the model automatically learns non-diagonal temporal coupling curves during training (i.e., allowing structure and sequence to be denoised at their respective optimal paces). This approach outperforms both "synchronous coupling" and "random coupling" baselines in SBDD and unconditional protein co-design tasks, revealing a universal "structure-leading" law where geometry precedes semantics in generation.

Background & Motivation

Background: The function of biomolecules (proteins, ligands) is determined by the coupling of sequence and 3D structure. Consequently, co-design (joint structural and sequential generation) has become the mainstream paradigm for de novo drug and protein design. Representative methods include MultiFlow, DPLM-2, La-Proteina (proteins), as well as TargetDiff, MolCRAFT, MolPilot, and DrugFlow (SBDD). These methods essentially perform diffusion or flow matching on a heterogeneous product manifold \(\mathbb{R}^{N\times 3} \times \mathbb{R}^{N\times K}\).

Limitations of Prior Work: Almost all co-design models implicitly adopt synchronous coupling, where all modalities share the same timestep \(t\) and evolve from noise to data at equal speeds. This is a strong implicit inductive bias, assuming identical denoising difficulty and convergence rates for all modalities. Recent works like Campbell et al. 2024 attempted to alleviate this with random coupling—sampling \((t_r, t_h) \sim [0,1]^2\) independently during training—but this introduces training-inference inconsistency (as inference usually follows a specific curve) and high-variance supervision.

Key Challenge: By observing SBDD training dynamics (Paper Fig. 1C), the authors found that under synchronous coupling, structural MSE remains high for most of the trajectory, dropping only very late. Switching to an asynchronous coupling allows structural error to drop earlier and improves validity. This indicates that the optimal generation trajectory is not the diagonal of the product manifold, but a geometrically curved geodesic where different modalities should be allocated time budgets based on their "learning complexity."

Goal: Elevate "how inter-modal time is coupled" from a hard-coded design choice to a learnable first-order design variable with controllable computational overhead.

Key Insight: Treat the training loss \(\mathcal{L}_\text{MSE}(\theta, \gamma)\) of multimodal generation as the transport cost in the temporal domain. The entire scheduling curve \(\gamma:[0,1] \to [0,1]^2\) corresponds to a coupling measure \(\pi_\gamma \in \mathcal{P}([0,1]^2)\). This translates "finding the optimal coupling" into "finding the lowest energy geodesic on the product manifold."

Core Idea: Use bi-level optimization + GP surrogate + Bayesian Optimization to learn this geodesic \(\gamma^*\) online within the training loop. The inner loop trains \(\theta\) with a fixed \(\gamma\), while the outer loop searches for a better \(\gamma\) on the loss surface provided by \(\theta^*\). The GP surrogate amortizes the cost of retraining required for each \(\gamma\) update.

Method

Overall Architecture

GeoCoupling abstracts multimodal generation as finding a monotonic curve \(\gamma\) on a 2D time square \([0,1]^2\) (Structure time \(t_r\) × Sequence time \(t_h\)) such that the transfer energy of the flow model trained along this curve is minimized. The framework follows a nested loop:

  • Inner Loop (MSE Training): Train the vector field \(v_\theta\) using standard flow matching/BFN/diffusion objectives under the current schedule \(\gamma\), where \(\theta^* = \arg\min_\theta \mathcal{L}_\text{MSE}(\theta, \gamma)\).
  • Outer Loop (Coupling Search): Store \((t_r, t_h, \mathcal{L})\) triplets observed during training into a rolling buffer \(\mathcal{B}\) with capacity \(N_\max = 1000\). Fit a cost surface \(c(t_r, t_h)\) using a Gaussian Process (GP), then find a new low-energy geodesic \(\gamma^*\) using Bayesian Optimization on the GP.
  • EMA Smoothing: Apply Exponential Moving Average (EMA) to the learned schedules to prevent sudden outer-loop changes from destabilizing the inner-loop training.

Input: Heterogeneous modal priors \(\pi_0 = p(\boldsymbol r) \otimes p(\boldsymbol h)\); Output: A coupled flow from \(\pi_0\) to the joint data distribution \(\pi_1 = p_\text{data}(\boldsymbol r, \boldsymbol h)\), along with the learned temporal coupling curve \(\gamma^*\).

Key Designs

  1. Temporal Optimal Transport Formulation (TOT):

    • Function: Transitions the traditional OT perspective of "pairing \(x_0, x_1\) in sample space" to the temporal domain. It treats the scheduling curve \(\gamma\) as a push-forward measure \(\pi_\gamma := \gamma_\# \lambda \in \mathcal{P}([0,1]^2)\), with transport cost \(\mathcal{E}(\gamma) = \int c(t_r, t_h)\, d\pi_\gamma\), where \(c(t_r, t_h) := \mathbb{E}_x[\mathcal{L}_\text{MSE}(x, (t_r, t_h))]\).
    • Mechanism: The authors prove (Prop. 3.2) that training loss integrated along \(\gamma\) decomposes into \(\mathcal{E}(\gamma) = \int [\,\underbrace{\|v_\theta - u^\gamma\|^2}_\text{Bias} + \underbrace{\mathrm{Var}(\mathbf{u}_t^\gamma \mid \mathbf{x}_t)}_\text{Variance}\,]\, dt\). Synchronous coupling is "high Bias, low Variance," while random coupling is "low Bias, high Variance." The geometric optimal \(\gamma^*\) seeks the minimum point between them.
    • Design Motivation: Provides a clean geometric and statistical explanation for why coupling needs to be learned—it is not an engineering trick, but a real optimal geodesic on the product manifold.

  2. Bi-level Optimization Target:

    • Function: Decouples finding \(\gamma\) from training \(\theta\), avoiding the computationally infeasible requirement of calculating hypergradients over the entire training trajectory.
    • Mechanism: Outer loop \(\min_{\gamma\in\Gamma} \mathcal{J}(\gamma) = \mathbb{E}_x[\int_0^1 \mathcal{L}_{\theta^*}(x, \gamma(t))\, dt]\), inner loop \(\theta^* = \arg\min_\theta \mathcal{L}_\text{MSE}(\theta, \gamma)\). Prop. 3.3 states that once bias is reduced by the inner loop, the optimal coupling is \(\gamma^* = \arg\min_\gamma \mathbb{E}_{t,x}[\mathrm{Var}(u_t^\gamma \mid \mathbf{x}_t)]\), giving the outer loop a clear, estimable target.
    • Design Motivation: Directly backpropagating through long inner-loop training is neither differentiable nor affordable; the bi-level variance perspective allows the outer loop to provide gradient signals just by "observing training loss."

  3. GP-BO Outer Loop:

    • Function: Solves the outer loop \(\gamma^*\) online and cheaply, allowing inner and outer loops to advance alternately.
    • Mechanism: Models the cost surface as a GP: \(c(\mathbf{t}) \sim \mathcal{GP}(\mu(\mathbf{t}), k(\mathbf{t},\mathbf{t}') + \sigma_n^2 \delta)\). A rolling buffer \(\mathcal{B}\) keeps only the most recent training observations to ensure the GP reflects the current capacity of the model. The outer loop uses acquisition functions from Bayesian Optimization to select candidate time pairs and applies a shortest-path algorithm on the GP surface to find a monotonic geodesic for the new \(\gamma\).
    • Design Motivation: A brute-force discrete grid search would require \(O(N^K)\) cost evaluations (measured at 1213.6s per update). GP-BO reduces this to 21.5s (56x speedup), enabling high-frequency outer-loop integration.

Loss & Training

The inner loop uses the native training objectives of the underlying models (Flow Matching / Diffusion MSE / BFN ELBO, etc.). The only change is sampling \((t_r, t_h)\) along the current \(\gamma\) instead of independent or synchronous sampling. Rolling buffer updates and EMA smoothing of the learned \(\gamma\) stabilize the training. The total training steps are roughly the same as the original models.

Key Experimental Results

Main Results

Structure-Based Drug Design (CrossDock, 100 test pockets × 100 molecules):

Category Method PB-Valid↑ Vina Score↓ (avg) Vina Dock↓ (avg) scRMSD<2Å↑
Reference - 95.0% -6.36 -7.45 34.0%
Sync MolCRAFT 84.6% -6.55 -7.67 46.8%
Sync DrugFlow 79.6% -5.12 -6.99 23.1%
Random MolPilot 95.9% -6.88 -7.92 41.1%
Learned GeoCoupling 94.3% -7.16 -8.32 43.1%

GeoCoupling leads comprehensively in binding affinity (Vina Score / Min / Dock), with PB-Valid comparable to MolPilot.

Unconditional Protein Co-design (Length 100-500, N=100):

Method Co-design↑ pLDDT↑ 1 - Pairwise TM↑ FS Clusters↑ Max TM↓
MultiFlow 0.72 79.39 0.63 0.56 0.83
La-Proteina (tri) 0.77 85.32 0.59 0.36 0.85
DPLM2 0.31 83.69 0.63 0.49 0.96
Ours 0.79 80.15 0.63 0.48 0.83
Ours (post-hoc → MultiFlow) 0.74 79.23 0.64 0.73 0.83

GeoCoupling achieves the highest co-designability. Its learned coupling also works as a plug-and-play component on MultiFlow checkpoints, increasing FS Clusters from 0.56 to 0.73.

Ablation Study

Configuration Connected↑ Vina Score↓ (mean) Vina Min↓ (mean) Description
Full (Ours) 93.5% -7.12 -7.57 Bi-level + EMA
Fixed \(\gamma^*\) 91.1% -6.97 -7.45 Fixed schedule before training
w/o EMA 91.9% -6.50 -7.24 No smoothing for outer schedule

Key Findings

  • Structure-leading is a universal law: For both SBDD (small molecules) and proteins, the learned \(\gamma^*\) shows a shape where "structure \(t_r\) advances fast early, and sequence \(t_h\) denoises rapidly only after structure stabilizes." This suggests geometric context is a necessary prior for sequence decoding.
  • Advantage in OOD lengths: When protein length \(\geq 400\), MultiFlow co-designability drops below 0.3, while GeoCoupling maintains \(> 0.6\), proving the coupling is a robust transport plan rather than an overfitted trick.
  • BO is indispensable: Dense-grid search takes 1213.6s per update vs. 21.5s for GP-BO (56x speedup), allowing the outer loop to run in real-time.
  • MolPilot is a special case: It is equivalent to running the outer loop once after training convergence. GeoCoupling achieves better results with 1x training steps.

Highlights & Insights

  • Elevating inter-modal temporal coupling to a learnable variable is the cleanest contribution. Previous work used either diagonal (sync) or uniform random sampling. This paper systematically shows both are extreme ends of a Bias-Variance trade-off, with the optimal solution lying on a geometric curve.
  • Unified Transport Perspective: Placing "sample space OT" and "temporal schedule OT" in a single framework. The former optimizes spatial coupling \(\pi(x_0, x_1)\), the latter optimizes temporal coupling \(\pi_\gamma(t_r, t_h)\). This duality bridges two major research lines in diffusion/flow matching.
  • Physical Interpretability of Structure-Leading: The automatically learned coupling confirms biological priors like induced fit—"build the scaffold before deciding the sequence."
  • Post-hoc Plug-and-Play: The learned \(\gamma^*\) can be transferred to existing checkpoints (like MultiFlow) without retraining, illustrating excellent engineering utility.

Limitations & Future Work

  • The authors acknowledge that GP-BO is a noisy approximate outer search without global optimality guarantees. The curse of dimensionality remains for \(K > 2\) modalities.
  • The learned coupling is optimal in the aggregate sense—using the same \(\gamma\) for all samples. Future work could introduce amortized conditional coupling \(\gamma(x)\).
  • Experiments did not cover all-atom proteins or protein-protein docking, and SBDD evaluation relies on Vina without wet-lab or rigorous physical simulation verification.
  • vs. MolPilot (Qiu et al., 2025): MolPilot performs a one-time search (VOS) after training, which is a degenerate version of the bi-level framework. GeoCoupling's co-evolution of coupling and model capacity allows it to outperform MolPilot with fewer training steps.
  • vs. MultiFlow / DPLM-2: These represent random coupling. This paper explains their training-inference inconsistency as "high-variance supervision."
  • vs. Classic OT Flow Matching: Previous works focused on sample space OT (straightening \(x_0 \to x_1\)). This work focuses on time domain OT (straightening the \(t_r \to t_h\) coupling). The two are orthogonal and additive.

Rating

  • Novelty: ⭐⭐⭐⭐⭐ Elevating temporal coupling to a learnable variable with a TOT framework is a significant conceptual innovation.
  • Experimental Thoroughness: ⭐⭐⭐⭐ Covering SBDD and proteins for both ID and OOD scenarios, though lacking wet-lab validation.
  • Writing Quality: ⭐⭐⭐⭐⭐ Clear propositions; Fig. 1 explains motivation, method, and phenomena lucidly.
  • Value: ⭐⭐⭐⭐⭐ The learned \(\gamma^*\) is plug-and-play for existing models, and the "structure-leading" discovery provides universal design guidance for AI for Science.