Sparse-Cost Temperature Calibration
Implementation & benchmark of arXiv:2609.01278 — An exact and fast solution of the inverse Regularized Optimal Transport problem
Source paper: An exact and fast solution of the inverse Regularized Optimal Transport problem arXiv:2609.01278 ⓘ · analyzed Sep 2, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Use a small set of known or trusted pairwise costs to estimate the effective entropic temperature of a Sinkhorn attention or mixture-of-experts routing layer directly from its observed transport plan. This provides a calibration controller that can detect over-concentrated routing and adjust epsilon without backpropagating through a costly temperature search.
Formulas
Mathematical statement
Let W be the observed transport plan, delta a positive numerical floor, and B = -H_m log(W plus delta) H_n. If a subset Omega of cost entries is known, the inverse-transport relation implies B_ij = C_centered,ij divided by epsilon, where C_centered = H_m C H_n. The least-squares temperature estimate minimizes the squared discrepancy between epsilon B_ij and C_centered,ij over Omega. Its closed form is epsilon_hat = sum over Omega of B_ij C_centered,ij divided by sum over Omega of B_ij squared. A normalized residual R measures whether the observed plan is consistent with the entropic model. A low R indicates reliable calibration; a high R indicates model mismatch, insufficient Sinkhorn convergence, or noisy cost references.
Implementation notes
Apply this controller to the temperature of a Sinkhorn attention head or an MoE routing matrix. Trusted costs can come from synthetic teacher geometry, distances between token metadata, class relations, or a small routing table; only a sparse index set Omega is required. Every T training steps, compute W, form B = -H_m log(W plus delta) H_n, gather B and centered reference costs on Omega, and calculate epsilon_hat using the closed-form regression formula. Compute R as the normalized residual. If R is below a threshold, update the live temperature with the damped geometric rule and clip it to safe limits. If R is high, freeze epsilon, increase Sinkhorn iterations, and log a model-mismatch warning rather than applying an unreliable update. The closed-form estimator and residual are the paper mechanism; damping, clipping, threshold, and update frequency are empirical design choices. First test on a small Transformer for synthetic sequence classification in which pairwise costs are Euclidean distances and 5 percent of costs are exposed. Compare fixed epsilon, gradient-learned epsilon, and inverse-OT calibration. The prediction is epsilon_hat divided by true epsilon within 10 percent when R is below 0.01. With independent observation noise, estimation error should decrease approximately as the inverse square root of the number of trusted entries. R should increase before routing entropy collapses when the learned attention no longer follows an entropic OT model.
Verification
This idea has not been verified yet.
Verification happens in two stages: Stage 1 — a mechanism check on a toy system confirms the claimed mathematical phenomenon reproduces; Stage 2 — a benchmark implements the idea on a real (small) neural network task and compares it against a tuned baseline over 8 paired seeds with a permutation test.
Artifacts
Artifacts unavailable.