🔢 mathematics

Scalable Bi-causal Optimal Transport via KL Relaxation and Policy Gradients

This paper introduces a scalable stochastic-optimization framework for computing bi-causal optimal transport couplings by employing a KL-penalized relaxation and policy-gradient algorithms, thereby overcoming computational barriers in continuous path spaces and enabling applications in robust finance and sequential uncertainty quantification.

Original authors: Haoyang Cao, Jesse Hoekstra, Renyuan Xu, Yumin Xu, Ruixun Zhang

Published 2026-05-19

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Haoyang Cao, Jesse Hoekstra, Renyuan Xu, Yumin Xu, Ruixun Zhang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot to walk exactly like a human. You have a video of a real human walking (the "target"), and you want the robot to mimic that movement perfectly.

However, there is a catch: The robot cannot see the future.

If the robot tries to move its foot before the human does, just because it "guessed" the human would step there, it's cheating. In the real world, you can only react to what has already happened, not what is about to happen. This is what the paper calls a "non-anticipative" constraint.

This paper solves a very hard math problem: How do you make two different things (like two stock markets, or a low-quality weather forecast and a high-quality one) move together perfectly over time, without either one peeking into the future of the other?

Here is the breakdown of their solution using simple analogies:

1. The Problem: The "Impossible Jigsaw Puzzle"

In the past, trying to match two complex, moving patterns (like stock prices over 100 days) was like trying to solve a jigsaw puzzle where the pieces keep changing shape every time you touch them.

The Old Way: Researchers tried to force the robot to match the human's path exactly at every single step. This worked for small, simple puzzles but crashed the computer when the puzzle got big or complex.
The Result: It was too slow and too hard to use for real-world problems like predicting financial risks or improving weather models.

2. The Solution: The "Soft-Constraint" Relaxation

The authors came up with a clever trick. Instead of forcing the robot to match the human perfectly at every single step (which is like a rigid, unbreakable rule), they introduced a "penalty system."

The Analogy: Imagine a coach telling the robot: "You don't have to match the human's step exactly right now, but if you drift too far off, you get a 'fine' (a penalty)."
The Math: They used a concept called KL Divergence (think of it as a "distance meter" between two probability clouds). If the robot's path starts to look different from the human's path, the "fine" gets bigger.
The Magic: By making the "fine" very large, the robot is forced to match the human almost perfectly, but because the rule is now a "soft penalty" instead of a "hard wall," the computer can solve the puzzle much faster using a technique called Policy Gradients (which is like the robot learning through trial and error, getting better with every attempt).

3. The "Dynamic" Learning Process

The paper proves that this "soft" method actually leads to the exact same result as the "hard" method if you crank the penalty up high enough.

The Recursive Structure: The authors showed that you don't need to plan the whole 100-day walk at once. You can just decide the next step based on where you are right now. This turns a massive, impossible calculation into a series of small, manageable steps (like a video game where you only need to plan the next jump, not the whole level).

4. Real-World Applications Tested

The authors didn't just do math on paper; they tested this on two specific real-world scenarios:

A. Robust Hedging (Financial Safety)

The Scenario: Imagine you are an investor trying to protect your money against a market crash. You need to know the "worst-case scenario" price for a financial product.
The Test: They used their method to find the safest possible price for a financial contract.
The Result: Their method found a price that was almost identical to the theoretical "perfect" price (within 1% error), but it did it much faster than previous methods. It successfully learned how to simulate market crashes that respected the rule: "You can't know the crash before it happens."

B. Time Series Statistical Downscaling (Weather & Data)

The Scenario: Imagine you have a blurry, low-resolution weather map (like a pixelated photo) and you want to turn it into a sharp, high-resolution map.
The Problem: If you just try to "sharpen" the blurry photo, you might invent fake weather patterns that don't make sense (e.g., rain appearing out of nowhere).
The Test: They used their method to "de-bias" the blurry data first, ensuring the low-res data matched the statistical rules of the real world, then generated the high-res version.
The Result: Their method created high-resolution weather patterns that were much more accurate and realistic than just guessing or using standard sharpening tools. It preserved the "flow" of time correctly.

Summary

This paper provides a scalable, fast, and accurate way to make two complex, moving systems mimic each other over time without cheating (looking into the future).

Old Way: Rigid, slow, and breaks on big problems.
New Way: Uses a "penalty system" to guide the learning, making it fast enough to run on modern computers while still being mathematically perfect.

It's like upgrading from trying to force a square peg into a round hole by hammering it (slow and damaging) to using a flexible mold that naturally shapes the peg to fit perfectly (fast and efficient).

Technical Summary: Scalable Bi-causal Optimal Transport via KL Relaxation and Policy Gradients

Problem Statement
Bi-causal optimal transport (OT) provides a framework for comparing and coupling stochastic processes under non-anticipative information constraints. Unlike classical OT, bi-causal OT restricts admissible couplings such that the conditional law of each process's next state depends only on the joint history up to the current time, ensuring neither process anticipates the other's future. This structure is critical for applications in robust finance (e.g., hedging under model uncertainty) and sequential uncertainty quantification (e.g., time series downscaling).

However, the practical application of bi-causal OT is hindered by computational intractability. Existing methods, such as entropic regularization with Sinkhorn iterations or fitted value iteration, struggle with continuous distributions, high state dimensions, and long time horizons. The primary difficulty lies in enforcing hard marginal constraints over the path space, which renders direct optimization combinatorially difficult and computationally expensive.

Methodology
The authors propose a scalable stochastic-optimization framework that replaces hard marginal constraints with a differentiable Kullback–Leibler (KL) penalty, enabling the use of policy-gradient methods.

KL-Penalized Relaxation:
The authors relax the original bi-causal OT problem by enlarging the feasible set of couplings to include arbitrary joint transition kernels. They introduce a KL-divergence penalty to regulate deviations from the prescribed marginal transitions ( $T_n, T'_n$ ) and initial marginals ( $\mu_0, \mu'_0$ ). The relaxed objective function $W_{rel}$ minimizes the expected cumulative cost plus a penalty term weighted by a parameter $\beta$ :
$W_{rel}(\mu, \mu') = \inf_{\pi \in \Pi_{rel}} \mathbb{E}_\pi \left[ \sum_{n=0}^N c_n(Y_n, Y'_n) + \beta \cdot \text{KL penalties} \right]$
This formulation preserves the recursive structure of the problem, allowing for dynamic programming.
Theoretical Foundations:
- Dynamic Programming: The paper establishes dynamic programming principles for both the original constrained problem and the relaxed formulation. This reduces the global path-space optimization to a sequence of one-step optimizations over initial couplings and transition kernels.
- $\Gamma$ -Convergence: The authors prove that as the penalty parameter $\beta \to \infty$ , the relaxed objective and its optimal couplings converge to the original bi-causal OT problem and its optimal solutions. This ensures the relaxation is a principled surrogate rather than just a heuristic.
Policy Gradient Algorithm:
- Parameterization: The coupling is parameterized using expressive conditional models, specifically normalizing flows (neural spline layers), which allow for exact log-density evaluation.
- Gradient Representation: The authors derive an explicit policy-gradient representation for the relaxed objective. The gradient decomposes into two components:
  - A value component ( $\nabla J_{val}$ ): An expectation under the current parameterized coupling, weighted by the remaining cumulative cost.
  - A KL component ( $\nabla J_{KL}$ ): An expectation under the prescribed marginal dynamics (reference measure), capturing the penalty term.
- Estimation: The algorithm employs unbiased mini-batch gradient estimators. To reduce variance, it utilizes state-dependent control variates fitted on auxiliary trajectories.
- Convergence: The paper provides non-asymptotic convergence guarantees, showing that the average regret is of order $O(1/K)$ under Polyak–Łojasiewicz conditions, where $K$ is the number of iterations.

Key Contributions

Relaxation Theory: The establishment of dynamic programming principles for the KL-relaxed problem and the proof of $\Gamma$ -convergence to the original constrained problem.
Scalable Algorithm: The development of a policy-gradient algorithm with unbiased estimators and variance reduction, capable of handling general marginals (including continuous and high-dimensional cases) without requiring discretization of state or action spaces.
Convergence Guarantees: Derivation of non-asymptotic regret bounds ( $O(1/K)$ ) for the proposed stochastic optimization method.
Empirical Validation: Demonstration of the method's ability to capture both marginal laws and temporal dependence in synthetic settings and real-world applications.

Results
The paper validates the framework through synthetic experiments and two specific applications:

Synthetic Experiments:
- The method accurately recovers marginal distributions (both Gaussian and non-Gaussian/Beta) and temporal dependencies (adjacent-step correlations) in both unimodal and bimodal settings.
- Performance metrics (Wasserstein distance, KS distance, Sliced Wasserstein distance) show low discrepancy between generated and ground-truth distributions.
- The method scales effectively to higher dimensions ( $d=3$ ) and longer horizons ( $N=10$ ).
Robust Hedging (Financial Application):
- The framework is applied to compute robust subhedging prices for financial derivatives under non-anticipative constraints.
- The learned coupling recovers near-optimal pricing bounds with relative errors below 1% compared to closed-form solutions.
- The method preserves martingale dynamics and outperforms existing approaches in computational speed (e.g., faster than fitted value iteration methods).
Time Series Statistical Downscaling:
- The method is used to generate high-fidelity time series from low-fidelity observations by first debiasing the low-resolution data via bi-causal OT.
- The "OT-adjusted" conditional diffusion model significantly outperforms unconditional generation and direct conditioning on coarse data, achieving lower distributional discrepancies and more accurate temporal correlations.

Significance and Claims
The paper claims to provide a scalable computational approach to bi-causal OT, bridging the gap between theoretical frameworks and practical implementation in settings where non-anticipative information constraints are essential.

Computational Feasibility: By converting the constrained path-space problem into a differentiable stochastic optimization problem, the method makes bi-causal OT tractable for continuous distributions and long horizons, overcoming the "curse of dimensionality" and memory costs associated with previous methods.
Generality: The framework is applicable to general marginals, not limited to discrete or empirical samples.
Practical Utility: The results suggest that bi-causal OT can serve as a practical computational primitive for larger optimization and uncertainty-quantification pipelines in operations research and quantitative finance, specifically for tasks like robust hedging and statistical downscaling where respecting information flow is critical for decision quality.

The authors maintain a modest tone regarding future work, noting that while the current analysis focuses on the Markov setting, the framework could be extended to non-Markovian settings and path-dependent costs, though this would require further theoretical development.