How fixed graphs shape gradients in sparse transport layers
Support Topology and Gradient Mixing in Sinkhorn Layers
Artificial IntelligenceMachine Learning
Summary
Gradient flow in machine learning layers that move data between tokens can be controlled by restricting connections using fixed graphs. The authors analyze how these graphs influence the way error signals travel backward through the computations that adjust those connections. They find specific conditions on the graph and data that ensure these error signals shrink steadily, helping with stable learning. This work offers guidelines for designing transport layers that are both efficient and mathematically well-behaved during training.
Sinkhorn algorithmtransport layergradient propagationfixed support graphrow-stochastic operatorreverse-mode differentiationtransportation polytopecontraction coefficientDobrushin contractiondifferentiable transport
Authors
Dylan Forde
Abstract
Sparse Sinkhorn layers use a fixed support graph to restrict transport between tokens. How does this graph control gradient propagation through the scaling iterations. We develop a fixed-support calculus showing that each row-column cycle induces a row-stochastic operator on column-potential perturbations modulo constants. Its transpose propagates zero-mass reverse-mode cotangents. The finite-cycle operator uses two distinct half-step transport plans; at a balanced fixed point it reduces to a two-step walk determined by a single plan. We derive the accompanying score and marginal source terms and use Dobrushin contraction and minorization to bound homogeneous and source-driven tail cotangents. Our main result characterizes when support and marginals guarantee one-step contraction uniformly over finite scores: every feasible face of the transportation polytope must have pairwise two-hop column overlap. Otherwise, suitable score directions make the contraction coefficient arbitrarily close to one. We extend this analysis to ordered support schedules and derive certificates for partition heat-bath layers, coordinate sweeps, forced shared mass, and register-augmented supports. These results provide mathematical criteria for support design in differentiable transport layers, with guarantees restricted to the fixed-support quotient-gradient component.