Transport-Signature Density Decoder

Integrate this as the output head of a conditional density model, VAE decoder, diffusion auxiliary model, or neural operator. Choose a fixed reference point cloud \(x_i\) with normalized weights \(w_i\) approximating \(\rho_0\), and precompute \(r\) basis functions \(\psi_j\) using the lowest eigenvectors of a finite-element or graph discretization of \(\mathcal L_{\rho_0}\). The network receives condition \(y\) and outputs \(a\in\mathbb R^r\). Compute \(\phi=\bar\phi+B a\) on the reference points, differentiate it spatially to obtain \(z_i=\nabla\hat\phi(x_i)\), and represent the output as weighted particles \((z_i,w_i)\). For a grid output, use a differentiable Gaussian splat or kernel histogram. Pseudocode: a = net(y); phi = phi_bar + B @ a; T = gradient(phi, x_ref); loss = sinkhorn(T, w, target_samples) + alpha * convexity_penalty; backpropagate(loss). The paper's mathematics supplies the pushforward decoder and Wasserstein surrogate; automatic differentiation, Sinkhorn approximation, and convexity penalties are engineering choices. Ensure approximate convexity by using a strongly convex \(\bar\phi\), limiting coefficient magnitudes, or penalizing negative eigenvalues of the Hessian. First test on a two-dimensional translated and deformed Gaussian or continuity-equation dataset. Compare against a pixelwise MLP decoder with equal parameter count and evaluate held-out Wasserstein error, negative predicted mass, total-mass error, and wall-clock training. Success means lower Wasserstein error with fewer latent coefficients and near-zero mass violation.