📊 statistics

Constrained Density Estimation via Optimal Transport

This paper proposes a novel framework for constrained density estimation that minimizes the Wasserstein distance to a prior under expectation and regularization constraints, utilizing an annealing-like algorithm to effectively address non-smooth constraints in both synthetic and real-world financial applications.

Original authors: Yinan Hu, Esteban G. Tabak

Published 2026-02-24

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yinan Hu, Esteban G. Tabak

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 a master chef trying to recreate a famous, complex dish (the Target) based on a recipe you have in your head (the Prior). However, you don't have the exact recipe for the final dish. Instead, you only have a few strict rules from a food critic:

"The total amount of protein must be exactly 50 grams."
"The dish must not contain any ingredients heavier than 200 grams."

Your goal is to tweak your original recipe just enough to satisfy these rules while changing the dish as little as possible from your original idea.

This paper proposes a new, smarter way to solve this problem, specifically for data and probability distributions. Here is the breakdown using everyday analogies.

1. The Old Way vs. The New Way

Traditionally, statisticians used a method called KL Divergence (think of this as "Vertical Stretching").

The Analogy: Imagine your original dish is a flat pancake. To add more protein, the old method would simply pour a thick layer of syrup on top of the pancake. The pancake gets taller (higher probability density) in some spots and stays flat in others.
The Problem: This creates "artifacts." You end up with weird, sharp edges or spots where the pancake is impossibly thick, and other spots where it's non-existent. It looks unnatural.

The authors propose a new method based on Optimal Transport (think of this as "Horizontal Shifting").

The Analogy: Instead of pouring syrup on top, imagine you have a pile of dough (your data). To satisfy the protein rule, you physically move chunks of dough from one side of the plate to another. You stretch the dough horizontally to fill the gaps or compress it to fit the rules.
The Benefit: This feels more natural. It respects the "shape" of the data. If you need more mass in a specific area, you move it there, rather than just inventing new mass out of thin air.

2. The "Rough Edge" Problem

The authors realized that while "moving dough" is better, it can still create ugly results.

The Issue: If you have a strict rule like "No ingredients allowed in this specific zone," the optimal way to move the dough is to push everything right up to the edge of that zone. This creates a massive pile of dough right on the boundary and a sudden empty hole (a gap) right next to it.
Real-world consequence: In finance, this looks like a price distribution that says "There is a 100% chance the stock price is exactly $50, and 0% chance it is $49.99." That's unrealistic. Prices usually wiggle around, they don't snap to a single line.

3. The Solution: "Smoothing" and "Annealing"

To fix these jagged edges and gaps, the authors introduced two clever tricks:

A. Smoothing Constraints (The "Spread the Butter" Rule)
They added a rule that says: "Don't let the dough pile up too high in one spot, and don't let it get too thin in another."

This forces the algorithm to spread the mass out more evenly, filling in the gaps and softening the sharp piles. It turns a jagged mountain range into a rolling hill.

B. The Annealing Algorithm (The "Warm-up" Strategy)
Solving this math problem is like trying to walk through a dark, foggy maze to find a specific exit. If you try to solve the hard rules immediately, you might get stuck.

The Trick: The authors start with "soft" rules. Imagine the rules are made of jelly. You can wiggle through them easily. You find a solution that satisfies the "jelly" rules.
The Process: Then, they slowly make the jelly harder and harder (like cooling it down or "annealing" metal). With each step, the rules become stricter, but because you started with a good solution, you don't get stuck. You gradually morph your solution until it perfectly satisfies the hard, real-world rules without getting lost.

4. Why This Matters: The Finance Example

The paper tests this on Option Pricing (predicting the future price of a stock).

The Scenario: You know the prices of simple "Vanilla" options (like a bet that the stock goes above $100). You want to figure out the probability of the stock hitting weird prices to price a complex "Exotic" option.
The Result: When they used their new "Horizontal Shifting" method with smoothing, they predicted the prices of these complex options much more accurately than the old "Vertical Stretching" method.
The Takeaway: The new method didn't just guess; it found a distribution that looked like real market data (smooth, continuous) rather than a mathematical artifact with sharp spikes.

Summary

Think of this paper as a new recipe for Data Reconstruction.

Old Method: Stretch the data vertically (like stretching taffy). It's easy but looks fake and breaks easily.
New Method: Move the data horizontally (like rearranging furniture). It's more natural.
The Innovation: They added "smoothing" to prevent furniture from piling up in corners and used a "warm-up" strategy to solve the puzzle without getting stuck.

This allows scientists and financial experts to build better models that respect the rules of the real world without creating impossible, jagged mathematical monsters.

1. Problem Statement

The paper addresses the problem of density estimation under expectation constraints. In many practical scenarios (e.g., finance, economics, nutritional planning), the target probability distribution (measure) is not explicitly known but is constrained by specific requirements, such as the expected values of certain functions (e.g., option prices, total protein intake).

Goal: Given a known prior distribution $\rho$ (source) and a set of expectation constraints $\mathbb{E}_\mu[f_k] = \bar{f}_k$ , find the target distribution $\mu$ that is "closest" to the prior while satisfying these constraints.
Metric: Unlike traditional methods that minimize the Kullback-Leibler (KL) divergence (a "vertical" measure of discrepancy), this framework minimizes the Wasserstein distance (Optimal Transport cost), which provides a "horizontal" measure of the displacement required to transform the prior into the target.
Challenge: When constraints involve non-smooth functions (e.g., indicator functions or ReLU functions used in option pricing), the resulting optimal target distribution often exhibits singularities, such as mass accumulation on boundaries or gaps (regions of zero probability), which may be undesirable in real-world modeling.

2. Methodology

A. Theoretical Framework

The authors formulate the problem using both Monge's (deterministic map) and Kantorovich's (transport plan) formulations of Optimal Transport.

Objective: Minimize $\int c(x, T(x)) d\rho(x)$ (or the transport cost $\pi$ ) subject to $\int f_k(y) d\mu(y) = \bar{f}_k$ .
Key Theoretical Insight: For one-dimensional problems with quadratic cost, the optimal transport map remains monotone, even under constraints.
Comparison with KL Divergence:
- KL Approach: Tends to "stretch" the prior density vertically. This often results in sharp corners or artifacts when constraints are non-smooth.
- Wasserstein Approach: Shifts mass horizontally. This leads to different artifacts: mass accumulation at boundaries and gaps in the support.

B. Analytical Solutions

The paper derives exact or semi-exact solutions for specific cases using the squared Euclidean distance ( $c(x,y) = \|x-y\|^2$ ):

Indicator Function Constraints: When the target support is restricted to an interval $[a, b]$ , the KL solution renormalizes the prior within the interval, while the OT solution restricts the prior to the interval and accumulates the displaced mass as Dirac deltas at the boundaries ( $a$ and $b$ ).
RELU Constraints: Used for option pricing ( $f(x) = \max(x-\omega, 0)$ ). The OT solution involves piecewise constant horizontal shifts of the prior mass, whereas the KL solution applies an exponential tilt.
Multiple Constraints: The paper provides a backward recursive algorithm to solve for multiple RELU constraints, determining the shift parameters and boundary points iteratively.

C. Numerical Algorithm & Regularization

Since analytical solutions are intractable for general cases and the prior is often given as a set of samples, the authors propose a numerical algorithm with the following features:

Sample-Based Optimization: The prior $\rho$ is replaced by an empirical measure $\hat{\rho}$ (a set of samples). The goal is to find a set of target samples $\{y_i\}$ that minimize the transport cost.
Annealing-like Schedule: To handle non-smooth constraints (like indicators or ReLUs) which cause vanishing gradients, the authors introduce a mollification parameter $\epsilon$ $ϵ$ .
- The constraints are smoothed (e.g., using $\tanh$ approximations).
- The algorithm starts with a large $\epsilon$ (smooth constraints) and gradually reduces it (annealing) to converge to the true non-smooth constraints.
Smoothing Inequalities: To mitigate the artifacts of mass accumulation and gaps, two additional inequality constraints are introduced on the target density $p_\mu$ $p_{μ}$ :
1. $\int p_\mu^2 \leq M_\delta$ : Prevents measure accumulation in small sets (controls $L^2$ norm).
2. $\int_D \frac{1}{p_\mu} \leq M_0$ : Prevents regions of vanishing density (controls $L^{-1}$ norm).
  These are enforced using barrier functions within the objective function.
Optimization: The problem is solved via gradient descent with adaptive learning rates (backtracking line search satisfying the Armijo-Goldstein condition).

3. Key Contributions

Novel Framework: Proposes a constrained density estimation framework based on Wasserstein distance rather than KL divergence, offering a geometrically distinct approach to handling expectation constraints.
Analytical Characterization: Provides explicit solutions for 1D problems with indicator and RELU constraints, highlighting the fundamental differences between OT and KL-based estimators (horizontal shifting vs. vertical stretching).
Regularization Strategy: Introduces a novel method to mitigate singularities (gaps and accumulations) in the OT solution using smoothing inequality constraints on the density.
Robust Algorithm: Develops an annealing-based numerical algorithm capable of handling non-smooth constraints and sample-based priors, ensuring convergence without getting stuck in local minima caused by vanishing gradients.

4. Results

The framework was validated through synthetic examples and a proof-of-concept application in quantitative finance.

Synthetic Examples:
- Visual comparisons (Figures 1 & 2) show that while KL divergence creates sharp corners, standard OT creates gaps and boundary spikes.
- The regularized OT (with smoothing inequalities) successfully redistributes mass to fill gaps and smooth boundaries, producing a more realistic target distribution.
- In 2D examples (circles and half-planes), the regularized method effectively mollifies the Dirac masses found in the unregularized exact solution.
Financial Case Study (Option Pricing):
- Task: Recover the risk-neutral pricing measure of an asset given vanilla option prices (constraints) and use it to price exotic options (Down-and-out, Cash-or-nothing, Asset-or-nothing).
- Setup: Prior = Lognormal(1,1); True Surrogate = Lognormal(2,1).
- Performance:
  - Both OT and KL methods significantly outperformed the unadjusted prior.
  - Wasserstein-based methods (both standard and smooth variants) achieved lower relative errors in 5 out of 6 test cases compared to the KL divergence method.
  - The "Smooth Wasserstein" approach provided the most robust pricing, minimizing errors in exotic option valuation.

5. Significance

Geometric Interpretation: The paper demonstrates that Wasserstein distance is a superior metric for problems where the "location" of probability mass matters (e.g., asset prices), as it respects the underlying geometry of the space better than KL divergence.
Handling Non-Smoothness: The proposed annealing algorithm offers a practical solution to the common issue of non-differentiable constraints in optimal transport, making the method applicable to real-world problems like option pricing with discrete barriers.
Artifact Mitigation: The introduction of smoothing inequalities addresses a critical limitation of pure OT (singularities), making the estimated densities more suitable for downstream applications where zero-probability regions or infinite densities are unrealistic.
Financial Application: The results suggest that OT-based calibration is a powerful tool for pricing exotic derivatives, potentially outperforming traditional maximum likelihood or KL-based approaches in markets with limited data.

In conclusion, this paper establishes a rigorous theoretical and numerical foundation for Constrained Optimal Transport, proving its efficacy in density estimation and offering a robust alternative to KL divergence, particularly in fields requiring geometric fidelity and handling of non-smooth constraints.