KL-Transport Condensation Layer
Implementation & benchmark of arXiv:2608.23762 — Transport based embeddings with topological guarantees
Source paper: Transport based embeddings with topological guarantees arXiv:2608.23762 ⓘ · analyzed Aug 29, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Replace a Euclidean embedding bottleneck with a simplex-valued KL transport layer. The encoder maps each input to a positive probability vector, which is compared against learned positive stochastic prototypes through a c-convex log-sum-exp potential; the resulting barycentric or projected representation should suppress nuisance directions while retaining the topology of the data manifold.
Formulas
Mathematical statement
Let Q=(q_{ij}) be an m x n positive stochastic matrix, with row distributions q_{i\bullet} in the probability simplex \Delta_n={p in R^n: p_j>0, \sum_j p_j=1}. For p,q in \Delta_n, the paper uses the KL cost c(p,q)=KL(p\Vert q)=\sum_{j=1}^n p_j\log(p_j/q_j). Its potential is \psi(p)=\log\sum_{i=1}^m\exp[-KL(p\Vert q_{i\bullet})], and the paper proves that \psi is c-convex for this cost. Thus its super-level filtration is compatible, through the associated c-dual transport construction, with the topology of the input filtration. The extracted proof also gives the explicit KL projection formula: if A is a d x n matrix, y is a target moment vector, and p minimizes KL(p\Vert q) subject to Ap=y, then p_j=q_j\exp[-(A^T\lambda(q))_j] for a multiplier \lambda(q) in R^d chosen so that Ap=y; normalization can be included as one row of A. In the neural adaptation, q_{i\bullet} are learned strictly positive prototype distributions, p is the encoder output, and A contains decoder or landmark coordinates. The useful properties are positivity, a closed-form exponential-family projection, and c-convex super-level sets rather than an unstructured Euclidean latent geometry.
Implementation notes
(1) Integration point: insert the layer after an image, graph, or sequence encoder h_theta(x), before the task head. Choose n landmarks or prototypes and emit logits W h_theta(x) in R^n; use p=softmax(logits) with an epsilon floor, p=(1-n epsilon) softmax(logits)+epsilon. Maintain m trainable prototype logits and define q_i=softmax(v_i), also epsilon-floored. The layer returns both the simplex code p and a condensed coordinate r=A p, where A in R^{d x n} contains learned or fixed landmark coordinates.
(2) Pseudocode:
h = encoder(x)
p = floor_softmax(W @ h, eps)
q = floor_softmax(V, eps) # m x n
kl[i] = sum_j p[j] * (log(p[j])-log(q[i,j]))
psi = logsumexp(-kl)
r = A @ p
loss = task_loss(head(r), target) - beta * psi
loss += gamma * topology_proxy_loss(r, batch_neighbors)
backprop(loss)
For an explicit projection variant, set y=A p and solve for lambda by Newton or damped gradient descent on F(lambda)=A[q*exp(-A^T lambda)]-y; then return p_proj=q*exp(-A^T lambda). Include a normalization row in A so the simplex constraint is enforced.
(3) Computed from the paper: the KL cost, c-convex potential, and exponential-family projection. Estimated empirically: prototype distributions V, coordinates A, beta/gamma, and any topology proxy. For topology_proxy_loss, preserve pairwise k-nearest-neighbor edges or match low-dimensional persistent-homology barcodes between input features and r; this is an engineering surrogate, not part of the theorem.
(4) First experiment: train a small ResNet-18 on rotated-MNIST or COIL-20 with a 16-dimensional PCA bottleneck versus the proposed n=64 KL code and a standard MLP bottleneck at equal parameter count. Measure classification accuracy, reconstruction error, kNN trustworthiness, and persistent-homology bottleneck distance as the bottleneck dimension is reduced. A useful signal is equal or better accuracy at 2-4x smaller latent dimension, with fewer spurious connected components or one-cycles and faster convergence than the Euclidean bottleneck.
Verification
Stage 1 · Toy mechanism gate: Failed ✗
Stage 2 · Mechanism transferred to benchmark: Not tested
Stage 2 · Practical benchmark result: Not run
Methodology: Toy-system gate first; the benchmark stage runs only after a pass. How verification works
Stage 1 — Mechanism check agent confidence 9/10
Built a KL-transport condensation layer with floored simplex encoder outputs, learned positive prototypes, KL log-sum-exp potential, condensed coordinates, and an exponential-family projection solver. The projection check succeeded with maximum constraint residual 7.95e-13 and stationarity range 1.0e-15, but the neural mini-experiment showed no win: the KL model reached 52.0% test accuracy versus 97.6% for the Euclidean 8D baseline under the same split and training budget.
- Agent confidence
- 9/10
- Baseline
- Digits test accuracy 0.9756, cross-entropy 0.0816, 1.19 seconds; unconstrained Euclidean 8D bottleneck.
- Idea
- Digits test accuracy 0.5200, cross-entropy 0.9412, 2.45 seconds; 8D simplex code, 12 learned prototypes, learned 8x8 condensation matrix, beta=0.01.
Limitations: Only one fixed digits split and one random seed were tested; no rotated-MNIST, COIL-20, reconstruction, trustworthiness, persistent homology, FLOP-normalized sweep, hyperparameter tuning, or explicit topology-preservation loss was evaluated. The KL architecture and Euclidean baseline were matched on nominal bottleneck dimension but not exhaustively matched on parameter count.
How to run: python3 experiment.py
Verdict computed by deterministic test code from paired-seed statistics — not by the language model.
Artifacts
- experiment.py 4.8 KB View Raw
- report.md 1.4 KB View
- results.json 1.0 KB View Raw