Wasserstein-Fisher-Rao Geometry: Unified Transport on Hybrid State Spaces#
TLDR#
When state is hybrid (discrete \(K\) + continuous \(z_n\)), belief evolution requires a geometry that handles both flow (transport within charts) and jumps (reaction between charts).
The Wasserstein–Fisher–Rao (WFR / Hellinger–Kantorovich) metric gives a single variational principle for this hybrid belief dynamics.
WFR resolves “duct-tape” product-metric heuristics by pricing transport and reaction consistently, and it yields implementable consistency diagnostics (WFRCheck).
Think operationally: belief is a fluid that can move and transform; WFR measures the cheapest way to do both.
This chapter connects the metric law (capacity) to dynamics (geodesic jump-diffusion) and to the information bound (area-law-like limits).
Roadmap#
Why product metrics fail for hybrid state spaces.
Define WFR and interpret transport vs. reaction components.
Connect WFR to filtering/control objectives and to implementation diagnostics.
Let me tell you about one of the most elegant solutions I’ve ever seen to a problem that seems hopelessly messy at first.
Here’s the situation. Our agent has an internal representation that mixes discrete and continuous parts. The discrete part says “what kind of situation is this?”—are we in the kitchen or the living room, is the object a cup or a bottle? The continuous part says “where exactly within that situation?”—the precise position, orientation, all the fine details.
Now, the standard approach is to handle these separately. You have some kind of graph or finite-state machine for the discrete part, and a Riemannian manifold for the continuous part, and you glue them together with duct tape and hope for the best. When does the agent “jump” from one discrete state to another? How do you compare paths that involve different combinations of jumping and moving? The whole thing becomes a computational and conceptual nightmare.
But there’s a beautiful way out of this mess, and it comes from a surprising place: optimal transport theory. The key insight is deceptively simple: stop thinking about the agent’s state as a point that moves around, and start thinking about it as a distribution that flows and transforms.
The latent bundle \(\mathcal{Z} = \mathcal{K} \times \mathcal{Z}_n \times \mathcal{Z}_{\mathrm{tex}}\) (Section 2.2a) combines a discrete macro-state \(K\) with continuous nuisance coordinates \(z_n\). The product metric \(d_{\mathcal{K}} \oplus G_n\) (Definition 2.2.1) and the Sasaki-like warped metric (Section 7.11.3) were heuristic constructions that treat the discrete and continuous components separately. These constructions are constrained by the agent’s Definition 2.
Researcher Bridge: Handling Distribution Shift
Standard Bayesian filters fail during “surprises” because they can’t handle mass appearing or disappearing (Unbalanced Transport). The Wasserstein-Fisher-Rao (WFR) metric allows the agent’s belief to both flow (smooth tracking) and jump (teleporting probability mass). This provides a unified variational principle for both continuous state-tracking and discrete hypothesis-switching.
This section introduces the Wasserstein-Fisher-Rao (WFR) metric—also known as Hellinger-Kantorovich [Chizat et al., 2018, Liero et al., 2018]—which provides a rigorous, unified variational principle. The key insight is to treat the agent’s internal state not as a point in \(\mathcal{Z}\), but as a measure (belief state) \(\rho_s \in \mathcal{M}^+(\mathcal{Z})\) evolving on the bundle.
Motivation: The Failure of Product Metrics#
Before we dive into the solution, let’s make sure we really understand the problem. Why doesn’t the obvious approach work?
The obvious approach is what I call the “Sasaki-like construction”—you take the metric on the discrete part, you take the metric on the continuous part, and you combine them. It’s like saying: the distance between two states is the discrete hop distance plus the continuous Riemannian distance.
This seems reasonable, but watch what happens when things get interesting.
The Problem with Sasaki-like Constructions.
The metric tensor from Section 7.11.3 (where \(\rho_{\text{depth}}\) denotes resolution depth, not density):
assumes a fixed point moving through the bundle. This creates two problems:
Discontinuous Jumps: When the agent transitions from chart \(K_i\) to chart \(K_j\), the metric provides no principled way to measure the “cost” of the jump versus continuous motion along an overlap.
No Mass Conservation: A point either is or isn’t at a location. But the agent’s belief can be partially in multiple charts simultaneously (soft routing, Section 7.8).
Let me make this very concrete. Suppose the agent is tracking an object that suddenly moves behind an occluder. The belief distribution should smoothly transition from “I’m pretty sure where it is” to “it could be in several places.” But if we’re tracking a point, we have to decide: does the point stay where it was, or does it jump? Neither option is right—the situation calls for a distribution that spreads out.
Or consider this: the agent is 90% confident it’s in scenario A and 10% confident it’s in scenario B. What’s the “position” of that belief? There isn’t one! You need a distribution.
The Core Problem
Think of a particle versus a probability cloud. A particle has to be somewhere—it can move, but it can’t be in two places at once. A probability cloud can spread, concentrate, flow, and even split. The agent’s belief is fundamentally a cloud, not a particle. Treating it as a particle forces artificial discretization: when do you “switch” hypotheses? The WFR framework says: you don’t have to choose. Mass can continuously redistribute.
The WFR Solution.
The Wasserstein-Fisher-Rao metric resolves both issues by lifting dynamics to the space of measures \(\mathcal{M}^+(\mathcal{Z})\). In this space:
Transport (Wasserstein): Probability mass moves along continuous coordinates via the continuity equation.
Reaction (Fisher-Rao): Probability mass is created/annihilated locally, enabling discrete chart transitions.
The metric determines the optimal path by minimizing the total cost: transport cost \(\int\|v\|_G^2\,d\rho\) plus reaction cost \(\int\lambda^2|r|^2\,d\rho\).
Here’s the beautiful idea. Instead of asking “where is the agent’s belief point?”, we ask “what is the agent’s belief distribution?” And instead of asking “how does the point move?”, we ask “how does the distribution evolve?”
This distribution can do two things: it can flow (mass moves from here to there while conserving total probability) or it can react (mass appears or disappears locally). The first is what happens when you track a moving object. The second is what happens when you suddenly realize “wait, I was wrong about which scenario I’m in.”
The WFR metric gives us a principled way to measure the “cost” of any combination of flowing and reacting. For fixed endpoints, metric, and reaction scale, the Benamou–Brenier formulation is convex in the appropriate flux variables such as \((\rho,\rho v,\rho r)\). That removes a combinatorial search from the variational problem, but the choice of model, chart maps, and \(\lambda\) remains a modelling decision.
The WFR Metric (Benamou-Brenier Formulation)#
Now let’s get precise. The Benamou-Brenier formulation is a beautiful way to think about optimal transport: instead of asking “what’s the cheapest way to rearrange mass from configuration A to configuration B?”, you ask “what’s the most efficient process that transforms A into B over time?”
It’s like the difference between asking “what’s the shortest path between two cities?” and asking “what’s the most fuel-efficient way to drive between them, considering traffic and terrain?” The second question embeds the problem in time and lets you think about dynamics.
Let \(\rho(s, z)\) be a time-varying density on the latent bundle \(\mathcal{Z}\), evolving in computation time \(s\). The WFR distance is defined by the minimal action of a generalized continuity equation.
Definition 53 (The Generalized WFR Action)
The squared WFR distance \(d^2_{\mathrm{WFR}}(\rho_0, \rho_1)\) is the infimum of the undriven generalized energy functional:
subject to the Unbalanced Continuity Equation:
where:
\(v_s(z) \in T_z\mathcal{Z}\) is the velocity field (transport/flow)
\(r_s(z) \in \mathbb{R}\) is the reaction rate (growth/decay of mass)
\(\lambda > 0\) is the length-scale parameter balancing transport and reaction
\(G\) is the Riemannian metric on the continuous fibres (Section 2.5) Conservative and driven cases. The displayed distance contains no vector potential. A non-conservative reward field is represented by the separately defined driven action \(\mathcal S_{\mathbf A}:=\mathcal E_{\mathrm{WFR}}-2\beta_{\mathrm{curl}}\int\!\langle\mathbf A,v\rangle\,d\rho\,ds\); it is a control objective, not a squared distance. Only the curvature \(\mathcal F=d\mathbf A\) is gauge invariant. If \(H^1_{\mathrm{dR}}(\mathcal Z)=0\) (or the harmonic component is set to zero) and \(\mathcal F=0\), a gauge with \(\mathbf A=0\) may be chosen; otherwise a harmonic component remains.
The driven action is not claimed to be gauge invariant or to yield the second-order Lorentz–Langevin SDE without an additional reduction and noise model. The SDE is defined independently in Definition Definition 69.
Forward reference (Boundary Conditions). Section 23.5 specifies how boundary conditions on \(\partial\mathcal{Z}\) (sensory and motor boundaries) constrain the WFR dynamics: Waking imposes Dirichlet (sensors) + Neumann (motors) BCs; Dreaming imposes reflective BCs on both, enabling recirculating flow without external input.
Let me unpack this piece by piece, because there’s a lot going on.
The unbalanced continuity equation is the heart of the matter: \(\partial_s \rho + \nabla \cdot (\rho v) = \rho r\). On the left side, we have how the density changes with time (\(\partial_s \rho\)) plus how it flows due to velocity (\(\nabla \cdot (\rho v)\)). On the right side, we have the reaction term (\(\rho r\))—if \(r > 0\), mass is being created; if \(r < 0\), mass is being destroyed.
In ordinary optimal transport, the right side is zero: mass is conserved, it just moves around. That’s the “balanced” case. But we need the “unbalanced” case because when an agent switches hypotheses—goes from “I think it’s scenario A” to “I think it’s scenario B”—mass has to disappear from A and appear in B. That’s not transport; that’s reaction.
The action functional measures the total cost of a path. You pay for velocity (moving mass around) and you pay for reaction (creating or destroying mass). The parameter \(\lambda\) sets the exchange rate: how much is one unit of transport worth compared to one unit of reaction?
And that vector potential term? It belongs to a separately driven action for situations where the reward one-form has curl—where going around a loop need not return the same value. It is a linear bias on candidate paths, not part of the squared WFR distance defined above. Only its curvature \(\mathcal{F}=d\mathbf A\) is gauge invariant; interpreting the bias as reward requires the declared field, units, and control reduction.
Example: Belief Update as WFR Flow
Imagine a robot tracking a ball. Initially, the belief is concentrated near position \(x_0\). Then the ball moves quickly to \(x_1\). What happens to the belief?
Pure transport (\(r = 0\)): The belief distribution flows smoothly from \(x_0\) to \(x_1\). This is what happens during normal tracking when the ball moves predictably.
Pure reaction (\(v = 0\)): The belief at \(x_0\) shrinks while belief at \(x_1\) grows. This is what happens during a “surprise”—the ball teleports (occlusion, fast motion), and rather than flowing smoothly, the belief essentially jumps.
Mixed: Usually both happen. The belief flows toward where you expect the ball to go, but also mass is transferred to alternative hypotheses (“maybe it bounced off something I didn’t see”).
The WFR metric finds the optimal mix. If \(x_1\) is close to \(x_0\), transport dominates (just track it). If \(x_1\) is far away, reaction dominates (teleport the belief).
Physics Isomorphism: Wasserstein-Fisher-Rao Geometry
In Physics: The Wasserstein-Fisher-Rao (WFR) or Hellinger–Kantorovich metric is a distinguished cone-space/Benamou–Brenier construction that combines optimal transport with Hellinger reaction geometry. It is one member of a broader family of unbalanced transport metrics [Chizat et al., 2018, Liero et al., 2018].
In Implementation: The finite non-negative belief measure \(\rho\) evolves under the WFR metric on \(\mathcal{M}^+(\mathcal{Z})\):
Correspondence Table:
Optimal Transport |
Agent (Belief Dynamics) |
|---|---|
Wasserstein distance \(W_2\) |
Transport cost for belief |
Fisher-Rao distance |
Information cost for reweighting |
Transport velocity \(v\) |
Belief flow in \(\mathcal{Z}\) |
Reaction rate \(r\) |
Mass creation/annihilation |
Benamou-Brenier formula |
Dynamic formulation |
Geodesic interpolation |
Optimal belief transition |
Significance: WFR unifies transport and reaction on the cone of non-negative finite measures \(\mathcal M^+(\mathcal Z)\). The pure-reaction restriction is the scaled Hellinger geometry on measures; its tangent tensor restricts to Fisher–Rao after fixing total mass, while a finite simplex requires an additional discrete-state restriction.
Remark 4 (Units)
\([v] = \text{length}/\text{time}\), \([r] = 1/\text{time}\), and \([\lambda] = \text{length}\) after taking the metric coordinates as length units. The ratio \(\|v\|/(\lambda |r|)\) is a local cost ratio; the exact crossover for a pair of Dirac masses depends on the normalization of the Hellinger term.
The units tell you something important. Velocity has units of length per time—that’s obvious. Reaction rate has units of inverse time—it’s a growth rate, like an interest rate. And \(\lambda\), the crossover parameter, has units of length.
So how do we compare transport and reaction? The ratio \(\|v\|/(\lambda |r|)\) is a useful local cost indicator when both fields and their units have been fixed. It is not a universal decision rule: the pairwise crossover for two measures also depends on mass normalization and the convention used for the Hellinger term.
The picture of “walking” versus “teleporting” is therefore a way to inspect a calibrated model. The variational problem, rather than the slogan, decides which admissible combination is cheaper.
Transport vs. Reaction Components#
Now let’s look at the two mechanisms separately before understanding how they combine.
The belief state \(\rho_s\) evolves on the bundle \(\mathcal{Z}\) via two mechanisms.
1. Transport (Wasserstein Component): The density evolves via the continuity equation \(\partial_s\rho + \nabla\cdot(\rho v) = 0\) along the continuous coordinates \(z_n\). The transport cost is \(\int \|v\|_G^2\, d\rho\). In the limit \(r \to 0\), the dynamics reduce to the standard Wasserstein-2 (\(W_2\)) optimal transport on the Riemannian manifold.
2. Reaction (Hellinger Component): The density undergoes local mass creation/annihilation via the source term \(\rho r\). This corresponds to discrete chart transitions: mass decreases on Chart A (\(r < 0\)) and increases on Chart B (\(r > 0\)). The reaction cost is \(\int \lambda^2|r|^2\, d\rho\). In the limit \(v \to 0\), the induced distance is the scaled Hellinger geometry on \(\mathcal M^+(\mathcal Z)\); its Fisher–Rao tensor appears only after restricting to fixed total mass (and, for a simplex, a finite chart register).
Here’s a way to think about the difference.
Transport is like rearranging furniture in a room. You can slide the couch from here to there, but the couch is still the same couch, and the total amount of furniture is conserved. The cost depends on how far you move things and how heavy they are.
Reaction is like a chemical reaction. You put in reactants, you get out products. Mass isn’t conserved locally—it appears and disappears. In our case, the “mass” is belief: probability assigned to different hypotheses.
Both mechanisms are doing something profound. Transport handles the question “how does my estimate of position change?” Reaction handles the question “how does my estimate of what situation I’m in change?” In a rich agent, both happen simultaneously.
The Two Extreme Cases
Limit |
What dominates |
Physical picture |
Agent behavior |
|---|---|---|---|
\(r \to 0\) |
Transport |
Incompressible fluid flow |
Smooth tracking within a hypothesis |
\(v \to 0\) |
Reaction |
Chemical kinetics |
Switching between hypotheses |
General |
Both |
Compressible reactive flow |
Tracking with hypothesis revision |
Most interesting agent behavior lives in the “general” regime. A robot tracking an object while considering alternative interpretations is doing both transport and reaction.
3. The Coupling Constant \(\lambda\) (Reaction-Transport Crossover Scale):
This parameter sets the characteristic length scale in the action. It does not by itself put the pairwise crossover at \(\|z_A-z_B\|_G=\lambda\); the numerical threshold depends on the mass normalization and, for Dirac endpoints, on the cone-angle convention.
Operational interpretation: choose \(\lambda\) as a declared WFR hyperparameter. It may be calibrated against a router-overlap scale, but the atlas overlap is not an intrinsic radius and is not identical to \(\lambda\) by definition.
The parameter \(\lambda\) sets the relative price of transport and reaction in the chosen action. The road and airline picture is useful, but the crossover is determined by the full variational cost, mass normalization, and cone convention, not by a universal distance equal to \(\lambda\).
A router-overlap scale can help calibrate \(\lambda\), but the two quantities are not identical by definition. Reaction is a local source term on the measure cone, so a transition need not be a literal jump between two discrete charts; in a finite chart register, the same mechanism can represent a change of hypothesis.
Definition 54 (Canonical length-scale)
Let \(G_n(\cdot;K)\) be the metric on the continuous fibre of chart \(K\). If the interior fibres have a positive uniform injectivity-radius lower bound, a geometric calibration is
where the infimum is taken over the chosen chart interiors. The discrete factor has no exponential map, and on an untrimmed open chart the infimum may be zero; in either case use a calibrated positive hyperparameter instead of calling this quantity canonical.
Default value. If the injectivity radius is unknown or the metric is learned, a practical default is:
where \(\ell_{\mathrm{step}}\) is a declared coordinate step and \(\bar G_n\) is the metric averaged over the sampled chart interiors. This is a calibration heuristic, not an intrinsic injectivity-radius identity.
Cross-reference: The screening length \(\ell_{\text{screen}} = 1/\kappa_{\mathrm{scr}}\) from Section 24.2 is a spatial discount-induced/value-field scale. It is distinct from the WFR reaction length \(\lambda\) (and from the temporal discount rate used to define \(\kappa_{\mathrm{scr}}\)).
Why the injectivity radius? Because that’s the scale at which the manifold’s topology starts to matter. Within the injectivity radius, the space looks like flat Euclidean space—you can move in any direction without running into weird topological obstacles. Beyond the injectivity radius, paths might wrap around, and the geometry becomes nontrivial.
If you’ve never encountered the injectivity radius before, here’s the intuition. Imagine you’re on the surface of a sphere. From any point, you can draw geodesics (great circles) in all directions. How far can you go before some of those geodesics start meeting again? That distance is the injectivity radius. On a sphere of radius \(R\), it’s \(\pi R\)—the distance to the antipodal point.
In the agent’s latent space, the injectivity radius tells you: how far can you transport mass along smooth geodesics before you have to worry about the discrete structure of the chart atlas?
Reconciling Discrete and Continuous#
Now comes the payoff. The WFR metric doesn’t just let us handle discrete and continuous separately—it actually unifies them in a way that respects the structure of both.
Proposition 12 (Limiting Regimes)
The WFR metric seamlessly unifies discrete and continuous dynamics:
Continuous Movement (Flow): When moving within a chart, \(r \approx 0\). The dynamics are dominated by \(\nabla \cdot (\rho v)\), and the metric reduces to \(W_2\) (Wasserstein-2). This recovers the Riemannian manifold structure of the nuisance fibres.
Discrete Movement (Jump): When the flow reaches a topological obstruction (chart boundary without overlap), transport can become prohibitively expensive. It can then be cheaper to use the source term \(r\):
\(r < 0\) on the old chart (mass destruction)
\(r > 0\) on the new chart (mass creation) On the full non-negative-measure cone this gives the Hellinger/Fisher–Rao-type pure-reaction metric; a Fisher–Rao simplex is obtained only after restricting to a finite normalized chart register.
Mixed Regime (Overlap): In chart overlaps, both \(v\) and \(r\) are active. The optimal path smoothly interpolates between transport and reaction.
Proof sketch. The cone-space representation lifts a pointwise measure to \((z,2\lambda\sqrt{\rho})\) with the cone metric. With \(r=0\) the projection is the \(W_2\) geodesic; with \(v=0\) it is the Hellinger geodesic on \(\mathcal M^+(\mathcal Z)\). A finite Fisher–Rao simplex is the discrete normalized special case. \(\square\)
This is the useful picture: WFR contains transport and reaction in one variational problem. With \(r=0\) it gives the Wasserstein transport restriction; with \(v=0\) it gives the Hellinger/Fisher–Rao-type cone geometry. A finite Fisher–Rao simplex appears only after restricting to a finite normalized chart register. In a mixed regime the minimizer can use both terms, subject to the stated endpoints and regularity assumptions.
The “cone-space representation” is a technical change of variables involving \(\sqrt{\rho}\). It makes the local formula easier to analyze, but it does not mean that every coordinate representation is linear or that supports move smoothly in the ordinary Euclidean sense. The continuity is in the measure geometry.
Analogy: Highway vs. Airplane
Imagine you’re in a landscape of cities connected by highways and airports.
Transport (Wasserstein): Driving on highways. You can go anywhere, but it takes time proportional to distance.
Reaction (Fisher-Rao): Taking a locally priced shortcut in weight, pictured as flying; in the full cone it can occur throughout the domain, while a finite chart register supplies the airport analogy.
WFR: Finding the optimal combination. For short trips, drive. For long trips, drive to the nearest airport, fly, then drive from the destination airport.
The length scale \(\lambda\) is like the crossover set by the relative price of driving and flying. In the actual model it is a declared WFR parameter, possibly calibrated against chart geometry. Reaction is a local source term on the whole measure cone; discrete “airports” are only an analogy for the finite chart-register specialization. The variational problem chooses a combination only after the endpoints, metric, and admissible fields have been fixed.
Connection to RL #26: Distributional RL as a related limit
The General Law (Fragile Agent): Belief states evolve on \(\mathcal{M}^+(\mathcal{Z})\) via Wasserstein-Fisher-Rao dynamics:
subject to the unbalanced continuity equation \(\partial_s \rho + \nabla \cdot (\rho v) = \rho r\).
The Degenerate Limit: Restrict to value distributions at single states (no spatial transport). Use Euclidean metric (\(G \to I\)).
Related Standard-RL equation:
The distributional Bellman equation is a useful comparison, but it is not derived from the reaction-only WFR action. C51, QR-DQN, and IQN add their own projection or quantile objectives [Bellemare et al., 2017, Dabney et al., 2018].
What the generalization offers:
Unified transport-reaction: WFR handles continuous flow (within charts) and discrete jumps (between charts) in one framework
Belief geometry: The metric on \(\mathcal{M}^+(\mathcal{Z})\) respects both \(W_2\) (spatial) and Fisher-Rao (probabilistic)
Teleportation length: \(\lambda\) determines when transport beats reaction (Proposition Proposition 12)
GKSL embedding: Quantum-like master equations embed naturally (Section 20.5)
Connection to GKSL / Master Equation (Section 12.5)#
Now let me show you a connection that has rigorous mathematical foundations, but only after we restrict the objects. A finite reversible classical master equation is a gradient flow in a discrete Wasserstein-type geometry. A GKSL equation reduces to that setting when the density matrix stays diagonal and the Hamiltonian and jump operators satisfy the stated hypotheses. That restricted classical limit is the connection; an arbitrary GKSL evolution is not automatically a WFR gradient flow.
The WFR framework connects rigorously to the GKSL (Lindblad) master equation via the classical limit. We state this precisely.
Theorem 7 (Classical Master Equation as a Discrete-Wasserstein Gradient Flow)
Let \(\mathcal{K} = \{1, \ldots, K\}\) be a finite state space with transition rates \(W_{jk} \geq 0\) (rate of jumping from \(k\) to \(j\)). The classical master equation
If the chain is reversible with respect to a strictly positive stationary distribution \(\pi\) (detailed balance \(W_{jk}\pi_k=W_{kj}\pi_j\)), this is the gradient flow of the relative entropy \(H(p \| \pi) = \sum_j p_j \log(p_j / \pi_j)\) with respect to a discrete Wasserstein-type metric. It is not, by this statement alone, a WFR reaction flow [Chow et al., 2012, Maas, 2011, Mielke, 2011].
Corollary 3 (GKSL Classical Limit)
Assume that \(\varrho = \mathrm{diag}(p_1, \ldots, p_K)\), that \(H\) is diagonal in the same basis, and that the jump operators preserve diagonal matrices (for example, they are linear combinations of matrix units). Then the GKSL equation reduces to a classical master equation with rates
The commutator term vanishes under the diagonal-\(H\) hypothesis. If the resulting rates satisfy detailed balance, Theorem Theorem 7 identifies the evolution as a gradient flow in the discrete-Wasserstein geometry.
Correspondence Table (Classical Limit):
GKSL Component (diagonal \(\varrho\)) |
WFR Interpretation |
|---|---|
\(-i[H, \varrho]\) (Commutator) |
Vanishes when \(H\) is diagonal in the belief basis |
\(\sum_j \gamma_j(L_j \varrho L_j^\dagger - \frac{1}{2}\{L_j^\dagger L_j, \varrho\})\) (Dissipator) |
Graph transport on \(\mathcal K\) (Maas metric), distinct from pointwise WFR reaction |
Probability conservation \(\sum_k \dot{p}_k = 0\) |
Balanced reaction (\(\int \rho r \, d\mu = 0\) globally) |
Jump operators \(L_j\) |
Transition kernels (where mass teleports to) |
Remark 5 (Full Quantum Case)
For non-diagonal density matrices (quantum coherences), the appropriate geometric structure is the quantum Wasserstein distance of Carlen & Maas [Carlen and Maas, 2017, Carlen and Maas, 2014]. The GKSL equation is the gradient flow of quantum relative entropy with respect to this metric. This framework handles coherences but is more complex than the classical WFR theory used here.
Let me be precise about what is rigorous and what is not.
Rigorous in the stated setting: A finite reversible classical master equation is a gradient flow of relative entropy in a discrete Wasserstein-type metric. Under the diagonal-density, diagonal-Hamiltonian, and jump-preservation assumptions, the GKSL equation reduces to that master equation. This is a theorem, not a claim about every Lindblad model.
Rigorous but different: For full quantum states with coherences, Carlen & Maas constructed a quantum Wasserstein framework in which suitable GKSL dynamics are gradient flows of quantum relative entropy. That space contains density matrices, not classical probability measures. The continuous WFR results concern their own measure-valued variational problem; they do not automatically convert a coherent GKSL evolution into the classical WFR problem.
The practical upshot: Use the classical WFR geometry when the belief state, rates, reversibility, and boundary conditions meet its hypotheses. For coherent or nonreversible dynamics, retain the appropriate quantum or nonequilibrium analysis instead of importing the WFR conclusion.
Why This Connection Matters
The GKSL/Lindblad structure is useful when one actually has a quantum state and a GKSL generator. It then supplies structural properties such as complete positivity, trace preservation for a trace-preserving generator, and Markovian composition. A WFR implementation over classical densities does not inherit those guarantees merely by using similar symbols. Its positivity, normalization, and Markov properties must be enforced or proved for the chosen discretization and update.
The Unified World Model#
Now let’s see how this theory translates into something you can implement. A world-model component can expose a velocity and a reaction rate in one interface, while a separate planner or loss trains those fields against a target. Whether the macro predictor and micro dynamics share parameters is an architectural choice; the WFR definition does not require one monolithic network.
The WFR formulation enables a single World Model that predicts both transport and reaction, eliminating the need for separate “macro predictor” and “micro dynamics” modules.
Definition 55 (WFR World Model)
The action-conditioned world model outputs a generalized velocity field \((v,r)\); the policy selects actions, and a separate planner or loss can train it to minimize WFR path length to a target distribution.
import torch
import torch.nn as nn
from typing import Tuple
class WFRWorldModel(nn.Module):
"""
Unified World Model using Unbalanced Optimal Transport dynamics.
Predicts the 'Generalized Velocity' (v, r) for belief particles.
No separate 'discrete' and 'continuous' modules.
"""
def __init__(
self,
macro_embed_dim: int,
nuisance_dim: int,
action_dim: int,
hidden_dim: int = 256,
):
super().__init__()
# Input: particle state + action
# State includes: macro embedding, nuisance coords, mass (weight)
input_dim = macro_embed_dim + nuisance_dim + 1 + action_dim
# Single MLP backbone for unified dynamics
self.dynamics_net = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
nn.SiLU(),
nn.Linear(hidden_dim, hidden_dim),
nn.SiLU(),
)
# Head 1: Transport velocity (Riemannian motion on fibre)
self.head_v = nn.Linear(hidden_dim, nuisance_dim)
# Head 2: Reaction rate (Fisher-Rao mass creation/destruction)
self.head_r = nn.Linear(hidden_dim, 1)
def forward(
self,
z_t: torch.Tensor, # [B, D] latent state (macro_embed + nuisance)
mass_t: torch.Tensor, # [B, 1] particle weight (belief mass)
action_t: torch.Tensor, # [B, A] action
dt: float = 0.1,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
"""
Predict next state via WFR dynamics.
Returns:
z_next: [B, D] next latent state
mass_next: [B, 1] next particle mass
v_t: [B, nuisance_dim] transport velocity
r_t: [B, 1] reaction rate
"""
# Unified prediction
inp = torch.cat([z_t, mass_t, action_t], dim=-1)
feat = self.dynamics_net(inp)
v_t = self.head_v(feat) # Transport velocity
r_t = self.head_r(feat) # Reaction rate (log-growth)
# Integrate dynamics (Euler step)
# Position update (Transport): z' = z + v * dt
z_next = z_t.clone()
z_next[..., -self.head_v.out_features:] += v_t * dt
# Mass update (Reaction): m' = m * exp(r * dt)
# If r > 0: hypothesis gaining probability (jumping in)
# If r < 0: hypothesis losing probability (jumping out)
mass_next = mass_t * torch.exp(r_t * dt)
return z_next, mass_next, v_t, r_t
How this handles the “Jump” seamlessly:
Deep inside a Chart: Model predicts \(r \approx 0\) and \(v \neq 0\). Particle moves normally.
Approaching a Boundary: Model sees invalid description (high prediction error). Predicts \(r < 0\) for current chart, \(r > 0\) for neighboring chart particles.
Result: Probability mass smoothly “tunnels” between charts without hard discrete switching.
Look at what this reference implementation exposes. The network outputs a velocity \(v\) and a reaction rate \(r\), and a simple Euler step gives
Position updates additively: \(z' = z + v \cdot dt\)
Particle masses multiplicatively: \(m' = m \cdot \exp(r \cdot dt)\)
The multiplicative update is the local reaction law: \(r>0\) grows a particle weight, \(r<0\) shrinks it, and \(r=0\) leaves that weight unchanged. It does not by itself conserve the total mass or implement the spatial divergence in the continuity equation; normalization, boundary flux, and a discretized divergence must be supplied by the chosen solver.
A trained model may learn small reaction rates in predictable regions and larger rates near a chart transition, but that behavior is a modelling hypothesis to validate. The code is an illustrative world-model component, not a proof that the learned fields select the WFR minimizer or switch charts correctly.
Particle Filter Interpretation
You can use this implementation as a particle-filter analogy, provided you keep the distinction clear. Particles carry weights and move with \(v\); the local update \(m' = m \cdot \exp(r \cdot dt)\) changes those weights. There is no literal resampling in this code, and the particle approximation still needs normalization and a rule for representing the spatial divergence.
The differentiable update makes end-to-end training possible for this component. It does not establish convergence to the WFR geodesic, preservation of total probability, or consistency with a boundary-value problem; those are separate numerical and analytic checks.
Scale Renormalization (Connection to Section 7.12)#
Now here’s something that makes physicists very happy: the WFR framework connects naturally to renormalization group ideas. If you’re not familiar with the renormalization group, don’t worry—the intuition is actually straightforward.
The idea is that physical systems often have structure at multiple scales. A turbulent fluid has large eddies containing smaller eddies containing even smaller eddies. An image has global composition, mid-level objects, and fine textures. And crucially, the “rules” at different scales might be different.
In our stacked TopoEncoder architecture, each layer corresponds to a different scale. The WFR metric applies at each scale, but with a scale-dependent parameter \(\lambda\).
For stacked TopoEncoders (Section 7.12), the WFR metric applies recursively with scale-dependent coupling.
Recall the WFR action:
For a hierarchy of layers \(\ell = 0, \ldots, L\):
Definition 56 (Scale-Dependent Teleportation Cost)
where \(\sigma^{(\ell)}\) is the scale factor from Definition Definition 21.
Interpretation:
Layer 0 (Bulk / IR): \(\Pi^{(0)}=1\) sets the reference scale. Jumping is expensive relative to later residual scales when the cumulative factors decrease.
Layer \(L\) (Texture / UV): the cumulative factor \(\Pi^{(L)}\) carries the absolute residual scale; a per-layer ordering must be checked from the measured factors rather than assumed from a single \(\sigma^{(\ell)}\).
Correspondence with Cosmological Constant: If the capacity-constrained metric law is applied separately at each layer, the term \(\Lambda^{(\ell)}G_{ij}\) plays the role of a baseline curvature. The correspondence is:
Bulk (low \(\Lambda\)): Flat, rigid, transport-dominated
Boundary (high \(\Lambda\)): Curved, fluid, reaction-dominated
This gives a useful way to discuss multi-scale representations, provided the measured residual scales support the ordering. The schedule uses the cumulative factors \(\Pi^{(\ell)}\), so a coarse layer is more transport-dominated or a fine layer more reaction-dominated only when those factors and the metric normalization make it so. The words “IR” and “UV” describe the intended interpretation; they do not prove an ordering for arbitrary encoders.
If the factors decrease with depth, a macro representation may carry a larger reaction scale and fine residuals a smaller one. That can make coarse hypotheses harder to replace and fine details easier to revise. Check the learned scales and the WFR residuals before assigning that physical picture to a particular model.
The Cosmological Constant Analogy
The correspondence with the cosmological constant \(\Lambda\) is a controlled analogy. In general relativity, \(\Lambda\) sets a baseline curvature. In this model, the same role can be assigned layer by layer only if the capacity-constrained metric law is applied at each layer and the units and coupling have been declared.
Under that additional schedule, a smaller \(\lambda^{(\ell)}\) corresponds to a larger nominal \(\Lambda^{(\ell)}\) through \(\Lambda^{(\ell)}\sim 1/(\lambda^{(\ell)})^2\). It does not by itself prove that the latent layer is more curved, that reaction is preferred, or that the bulk and boundary have the stated ordering. Those are diagnostics to check against the learned metric and admissible WFR paths.
Connection to Einstein Equations (Section 18)#
We’ve been talking about the geometry of belief space as if it were fixed. Here is the careful version of the connection. Varying the WFR action with the density, velocity, and reaction fields held fixed produces an auxiliary stress tensor. It can be coupled to the metric-law equation only after an additional source identification and unit conversion have been declared.
So the useful analogy with Einstein’s equations is structural: one variational stress tensor can act as a source in a separate metric equation. The WFR action alone does not determine the latent metric, and the metric-law theorem does not automatically identify its Risk Tensor with this WFR tensor.
The WFR dynamics provide an auxiliary variational stress tensor. The metric law’s \(T_{ij}\) is the reward Risk Tensor in Section 18; the WFR tensor is not that source unless an additional coupling and unit conversion are declared.
Theorem 8 (Auxiliary WFR variational stress tensor)
Let the WFR action be
with continuity equation
Define, under the explicit density-fixed convention,
Then
which is a perfect-fluid form under this density-fixed convention, with reaction contributing an additive pressure term \(P_{\mathrm{react}}=\tfrac12\lambda^2\rho r^2\).
Proof sketch. Vary \(\mathcal{S}_{\mathrm{WFR}}\) with respect to \(G^{ij}\) while holding \((\rho,v,r)\) fixed. Use \(\delta\|v\|_G^2=-v_i v_j\,\delta G^{ij}\) and \(\delta d\mu_G=-\tfrac12 G_{ij}\delta G^{ij}d\mu_G\), then collect terms to match \(\delta\mathcal{S}_{\mathrm{WFR}}=-\tfrac12\int T^{\mathrm{WFR}}_{ij}\delta G^{ij}d\mu_G\,ds\). See Appendix C for the full derivation. Holding the belief measure fixed instead removes this pressure term and changes the mass variation, so the convention is part of the statement. \(\square\)
Let me decode this. The auxiliary tensor \(T_{ij}^{\mathrm{WFR}}\) records how belief mass, transport, and reaction enter the variation of the selected WFR action. In relativity, a stress-energy tensor is a source in Einstein’s equations; here the same word describes a diagnostic with a different domain and derivation.
The result has a perfect-fluid-like algebraic form: a density times velocity-squared term and a pressure-like term. Calling it “perfect fluid” describes that algebraic shape; it does not import relativistic matter dynamics.
The reaction contribution is worth noticing. High \(r\) increases the pressure-like term in this auxiliary tensor. It can influence curvature only after a separate metric-law coupling is declared and its hypotheses are checked; the WFR variation alone does not make geometry curve.
What does this mean in practice? Regions of high belief dynamics can be flagged as carrying a larger auxiliary load. Whether they become geometrically different from regions of certainty is a modeling or control response that must be measured, not a consequence of this tensor alone.
Physics Isomorphism: Stress-Energy Tensor
In Physics: The stress-energy tensor \(T_{\mu\nu}\) is derived from the variation of the matter action with respect to the metric: \(T_{\mu\nu} = -\frac{2}{\sqrt{-g}}\frac{\delta S_M}{\delta g^{\mu\nu}}\) [Wald, 1984].
In Implementation: The auxiliary WFR tensor (Theorem Theorem 8) is:
derived from \(\delta \mathcal{S}_{\text{WFR}}/\delta G^{ij}\) under the density-fixed convention. It is an auxiliary perfect-fluid diagnostic, not the Risk Tensor in the capacity metric law, with pressure \(P = \frac{1}{2}\rho(\|v\|_G^2 + \lambda^2 r^2)\).
Correspondence Table:
Field Theory |
Agent (WFR) |
|---|---|
Matter density \(\rho_m\) |
Belief density \(\rho\) |
4-velocity \(u^\mu\) |
Transport velocity \(v^i\) |
Pressure \(p\) |
Reaction pressure \(\frac{\lambda^2}{2}\rho r^2\) |
Rest mass density |
WFR kinetic energy \(\frac{1}{2}\rho|v|_G^2\) |
Implications:
High velocity (\(v\)): Agent moves fast through a region, so the auxiliary stress contribution \(T_{ij}\) can be large under the density-fixed convention. Any resulting curvature response or contraction is conditional on the selected metric-law coupling; the natural-gradient interpretation is a modeling interpretation, not a first-principles consequence of the WFR variation alone.
High reaction (\(r\)): Agent jumps frequently → \(P_{\mathrm{react}}\) increases → capacity stress increases. This triggers the boundary-capacity constraint (Definition Definition 47).
These implications deserve emphasis, with the convention in the theorem kept in view.
The variation says that, under the density-fixed convention, transport and reaction enter an auxiliary stress tensor and that reaction contributes the nonnegative term \(\tfrac12\lambda^2\rho r^2\) to its pressure component. It does not say that high velocity has a universal curvature sign, that the latent space must contract, or that this tensor is already the state-space natural-gradient metric.
A capacity response can be tested after a separate coupling to the metric law has been specified. High reaction may then be a useful load indicator, but it is not by itself a proof of capacity saturation.
Consistency with existing losses:
Existing Loss |
WFR Interpretation |
Status |
|---|---|---|
\(\mathcal{L}_{\mathrm{pred}}\) (Prediction) |
Minimizing transport cost \(\lVert v\rVert_G^2\) |
Compatible |
\(\mathcal{L}_{\mathrm{closure}}\) (Macro closure) |
Penalizing reaction \(r\) in macro channel |
Compatible |
Mass reaction |
\(r<0\) annihilates belief mass; entropy monotonicity requires a separate gradient-flow hypothesis |
Compatible |
Capacity (\(I < C\)) |
Metric curves to keep WFR path within budget |
Compatible |
Why This Matters for Implementation
The compatibility table is a map of possible correspondences, not an identity between losses. A prediction loss can serve as a transport proxy only when its residual and units have been related to \(\|v\|_G^2\). A closure loss can penalize reaction only when the chart model defines reaction through that loss. Entropy monotonicity requires a separate reversible or gradient-flow hypothesis.
Likewise, the capacity row does not say that WFR automatically curves the metric or enforces \(I<C\). That requires the capacity proxy, the metric-law coupling, and the boundary conditions to be specified. WFR can be added as a variational layer or consistency regularizer; whether it replaces existing modules is an implementation choice to be validated.
Comparison: Sasaki vs. WFR#
Let me summarize the comparison between the old approach (Sasaki-like product metrics) and the new approach (WFR). This table tells the whole story.
Feature |
Sasaki (Product Metric) |
WFR (Unbalanced Transport) |
|---|---|---|
State representation |
Fixed point |
Probability mass / belief |
Topology changes |
Manual patching required |
Handled natively via \(r\) |
Path type |
“Walk then Jump” (discontinuous) |
Smooth interpolation |
Optimization |
Combinatorial + Gradient descent |
Convex in \((\rho,\rho v,\rho r)\) with fixed endpoints and metric |
Theoretical consistency |
Ad-hoc construction |
Gradient flow under the stated reversibility hypotheses |
Multi-scale |
Separate metrics per scale |
Unified with scale-dependent \(\lambda\) |
Every row in this table compares a different modelling choice. Let me highlight the useful distinction.
Optimization: A product construction may require an external search over chart sequences. For fixed WFR endpoints, metric, and reaction scale, the flux formulation is convex and avoids that combinatorial enumeration. The claim is about this constrained variational problem; learning the endpoints, chart maps, or metric is still generally nonconvex.
Path type: WFR gives a continuous curve in the metric space of measures when the admissible action is finite. Its transport and reaction components can trade off along that curve, but a coordinate trajectory or chart support need not look smooth, and a numerical scheme can still introduce discontinuities.
Theoretical consistency: WFR is a specified variational metric with a well-defined cone construction under its assumptions. Calling it “unique” without naming the class of metrics is too strong. Its usefulness comes from the stated action, continuity equation, endpoints, and boundary conditions, not from a universal replacement theorem.
Implementation: WFR Consistency Loss#
Now let’s get concrete about how to train models with this framework. The key idea is a consistency loss that measures a time-discretized residual of the unbalanced continuity equation. It is a local diagnostic for the chosen representation and discretization, not the WFR distance itself and not a substitute for solving the full boundary-value problem.
Definition 57 (WFR Consistency Loss / WFRCheck)
The cone-space representation linearizes WFR locally. From \(\partial_s \rho = \rho r - \nabla \cdot (\rho v)\) and \(u = \sqrt{\rho}\), we have \(\partial_s u = \frac{\rho r - \nabla \cdot (\rho v)}{2\sqrt{\rho}}\). Define the consistency loss:
This penalizes deviations from the unbalanced continuity equation.
Practical implementation:
def compute_wfr_consistency_loss(
rho_t: torch.Tensor, # [B, K] belief over charts at time t
rho_t1: torch.Tensor, # [B, K] belief over charts at time t+1
v_t: torch.Tensor, # [B, K, d_n] transport velocity per chart
r_t: torch.Tensor, # [B, K] reaction rate per chart
dt: float = 0.1,
eps: float = 1e-6,
) -> torch.Tensor:
"""
Compute WFR consistency loss (cone-space formulation).
Penalizes violation of unbalanced continuity equation.
"""
sqrt_rho_t = torch.sqrt(rho_t + eps)
sqrt_rho_t1 = torch.sqrt(rho_t1 + eps)
# Approximate divergence term (finite difference)
# In practice, use automatic differentiation if v is differentiable
div_rho_v = torch.zeros_like(rho_t) # Placeholder for nabla . (rho v)
# Predicted change in sqrt(rho) from: d/ds sqrt(rho) = (rho*r - div(rho*v)) / (2*sqrt(rho))
predicted_delta = (dt / (2 * sqrt_rho_t + eps)) * (rho_t * r_t - div_rho_v)
# Actual change
actual_delta = sqrt_rho_t1 - sqrt_rho_t
# L2 loss
loss = ((actual_delta - predicted_delta) ** 2).mean()
return loss
Why work with \(\sqrt{\rho}\) instead of \(\rho\)? This is the “cone-space” change of variables. It gives the local identity used by the residual and often improves numerical conditioning, but it does not remove the need to define a spatial discretization or boundary treatment.
The consistency loss compares the observed one-step change in \(\sqrt{\rho}\) with the change predicted by the current \((v,r)\) fields. A small value means those sampled transitions agree with the discretized equation. It does not prove that the fields are a global WFR minimizer or that the density remains normalized.
The divergence term \(\nabla \cdot (\rho v)\) is the delicate part computationally. Finite differences or automatic differentiation can approximate it only after a grid, mesh, or differentiable spatial representation has been chosen; the placeholder in the example intentionally omits that choice.
Implementation Notes
A few practical considerations:
The \(\epsilon\) stabilizer: We add \(\epsilon = 10^{-6}\) inside the square root to avoid division by zero when \(\rho \approx 0\). This corresponds to a tiny uniform “background” belief.
The divergence term: The placeholder
div_rho_v = torch.zeros_like(rho_t)in the code is a simplification. For a full implementation, you’d need to either:Discretize the divergence using finite differences on a grid
Use automatic differentiation through a neural velocity field
Use a divergence-free parameterization and ignore this term
Batching: The loss is computed per-batch and averaged. You want to see it decrease during training, indicating that the world model’s \((v, r)\) predictions are becoming more consistent with actual belief evolution.
Scaling: The loss magnitude depends on \(dt\) and the scale of \(\rho\). You may need to tune the loss weight relative to other training objectives.
Auxiliary diagnostic: WFRCheck#
Finally, we define a diagnostic node that monitors WFR consistency at runtime. This fits into the larger diagnostic framework described in Section 3.
Following the diagnostic node convention (Section 3.1), we define:
ID |
Name |
Component |
Type |
Interpretation |
Proxy |
Cost |
|---|---|---|---|---|---|---|
aux-WFR |
WFRCheck |
World Model |
Dynamics Consistency |
Transport-Reaction balance? |
\(\mathcal{L}_{\mathrm{WFR}}\) |
\(O(BK)\) |
Trigger conditions:
High \(\mathcal{L}_{\mathrm{WFR}}\): World model’s \((v, r)\) predictions violate continuity
Remedy: Increase training on transitions; check for distribution shift
When should you worry about this diagnostic? A high WFRCheck residual means the sampled one-step prediction disagrees with the discretized continuity equation. Insufficient training and distribution shift are two possibilities, but so are an inaccurate divergence approximation, a mismatched time step, boundary leakage, or an inconsistent density convention.
Treat the number as a warning and inspect those choices before changing the model. A low residual supports local one-step consistency for the samples and discretization used; it does not certify a global WFR geodesic, normalization, or convergence of the learned world model.
Summary: What WFR Buys You
Let me summarize the useful benefits of the WFR framework:
One variational language: Transport and reaction can be priced together for fixed endpoints, metric, and admissible fields. An implementation may still use separate modules for convenience.
Declared jump cost: The parameter \(\lambda\) sets the relative scale of reaction and transport. It can be calibrated from geometry, but it is not derived from geometry without an additional calibration argument.
Convex subproblem: The flux formulation is convex for the fixed-data variational problem. Training a neural world model and learning a metric need not be convex.
Careful physical analogies: Fluid, thermodynamic, and quantum language can suggest useful constructions, while the corresponding physical identifications require their own hypotheses.
A runtime residual: WFRCheck monitors the selected discretization and samples. It is a diagnostic, not a proof of global consistency.
The core picture remains simple: let belief flow and react, charge for both, and check that the numerical path obeys the continuity equation and the declared boundary conditions.