Detailed-Balance Graph Transport Layer

(1) Exact integration point: insert the module after graph attention or message passing and before the next MLP/residual connection. Let each node feature \(h_i\in\mathbb R^d\) produce \(K\) positive transport channels \(\rho_{i,k}=\operatorname{softplus}((W_\rho h_i)_k)+\epsilon\). Normalize each channel over nodes if a fixed channel mass is desired. Construct an edge conductance from the current features, \(c_{ij}=\operatorname{softplus}(u^T[h_i,h_j,|h_i-h_j|])+\epsilon\), and symmetrize it as \(c_{ij}\leftarrow(c_{ij}+c_{ji})/2\). Predict positive equilibrium masses \(\pi_{i,k}\) with a softplus head and normalize them over nodes. (2) Pseudocode: compute \(q_{i,k}=\rho_{i,k}/\pi_{i,k}\); for every undirected edge compute \(J_{ij,k}=c_{ij,k}(q_{i,k}-q_{j,k})\); update \(\rho_{i,k}\leftarrow\rho_{i,k}+\Delta t\sum_j c_{ij,k}(q_{j,k}-q_{i,k})\). For an explicit baseline require \(\Delta t\le\min_i\rho_i/\sum_j c_{ij}/\pi_i\). Prefer an implicit step \((I-\Delta t A)\rho^{new}=\rho^{old}\), where \((A\rho)_i=\sum_j c_{ij}(\rho_j/\pi_j-\rho_i/\pi_i)\), solved with a sparse linear solver. Feed \(\log(\rho^{new}+\epsilon)\) through a learned projection and add or concatenate it to \(h_i\). (3) The mathematics computes conservation, detailed balance, and the dissipation diagnostic; learned conductances, equilibria, step size, and downstream usefulness are estimated empirically. Log total mass, minimum mass, and \(F(\rho)\) after every transport step. (4) First cheap experiment: use a 2-layer GCN or Graph Transformer on Cora and a synthetic stochastic-block-model node-classification dataset. Compare standard residual message passing, ordinary diffusion/GCN propagation, and this transport layer with equal hidden width, parameter count, and propagation FLOPs. Run 5, 10, and 20 refinement steps under both normal and deliberately large residual step sizes. Success requires no negative activations, monotone measured free energy, bounded iterates, and higher validation accuracy or calibration at equal compute; the strongest signal is graceful accuracy degradation when refinement depth increases, unlike the unconstrained residual baseline.