Hyper-connections widen the residual stream of a Transformer to n parallel streams. Their manifold-constrained version (mHC) mixes the streams at each layer with a doubly stochastic matrix, which it computes by Sinkhorn normalization of exponentiated logits. We give a geometric theory of this design on the Birkhoff polytope. First, a doubly stochastic mixer splits the stream into a mean channel, on which mHC is exactly a residual network, and a difference channel, which each layer contracts by its second singular value σ2≤1−nminijHij. Thus the extra width is a fading memory with a horizon of 1/(1−σ2) layers, and among nonnegative mixers only the permutations do not collapse. Second, the Sinkhorn-logit map is a global chart, and its logit gradient is exactly the Fisher-Rao gradient. Thus logit gradient flow follows a squared Fisher-Rao metric, and the straight-through update is exactly entropic mirror descent. Third, under logit gradient flow the logarithm of each entry moves at a rate of at most 4n3∥∇f∥∞ε, where ε is the distance to the nearest permutation. Thus gradient flow approaches and leaves the vertices only at rate 1/t, but mirror descent moves at an exponential rate. Fourth, the local convergence factor of Sinkhorn is σ22, so a fixed iteration budget limits the horizon. Experiments confirm the predicted rates.
Figures & tables
Figure 1: The three effects of the theory, all computed. (a) Four streams under one mHC mixer H=SK(exp(4A)) at every layer. The stream mean (dashed) does not change, and the differences decay inside the envelope σ2(H)l∥x0⊥∥ of Theorem 3.2 . (b) The same start under the isometric mixer A∈O1(n) : the differences keep their norm (Proposition 7.1 ). (c) The Birkhoff polytope B3 under a linear map that sends the six permutations to a regular hexagon and J/n to its center. The loss ⟨C,H⟩ has its minimum at I . From one start, the straight-through update reaches I by t=8 ( ε=0.003 ). Logit gradient flow moves toward the vertex (12) and is near it at t=8 ; it needs t≈400 to reach ε=0.001 at I (Theorems 5.3 and 5.4 ).
Figure 2: Numerical checks with n=4 . (a) Second singular value of the composite of L projected random mixers (median of 30 draws), and the bound exp(−n∑lδ(Hl)) of Theorem 3.2 (b) (dotted). (b) Vertex proximity under three updates, from the uniform matrix. The dashed line is proportional to 1/t . (c) Time to leave the identity vertex, as a function of the initial proximity ε0 . The dashed line is proportional to 1/ε0 . (d) Sinkhorn sweeps that give a column-sum error of 10−3 , as a function of the horizon τ of the target.
Appendix figures & tables5 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 3: Dependency graph of the results. An arrow from X to Y means that the proof of Y cites X . We omit an arrow if other arrows already give a path from X to Y . Bold frames mark the main results, and dashed frames mark the known results of Appendix D . A script (fig_dependency.py) builds the graph from the references in the structured proofs, so the graph and the proofs agree.
Figure 4: Collapse of composite mixers (Theorem 3.2 ). Left: σ2 of the composite of independent projected random logits of scale M (median of 30 draws, solid), and the bound exp(−n∑lδ(Hl)) (dotted). Right: largest entrywise distance of the composite to J/n . Even for M=4 , the difference channel loses six orders of magnitude within 25 layers.
Figure 5: Vertex viscosity (Theorems 5.3 and 5.4 ). Left: vertex proximity for the loss ⟨C,H⟩ , C=J−I+0.1Ξ , from the uniform matrix. The updates are exact logit gradient descent (step 0.2 ), Adam on the logits (step 0.02 ), and the straight-through update (step 0.2 ). The dashed line is proportional to 1/t . Right: time to leave the identity vertex for a transposition vertex, as a function of the initial proximity ε0 , under the same three updates. The start is W=log(1/ε0)I , and the escape time is the first time at which every entry on the transposition is larger than 0.5 and every other entry is smaller than 0.25 . The dashed line is proportional to 1/ε0 .
Figure 6: Truncated Sinkhorn and the horizon (Proposition 6.1 ). Left: column-sum error of the row-normalized iterate against the number of sweeps, for targets (1−τ−1)Pπ+τ−1J/n with horizon τ and a 4-cycle π . Each run starts from a random diagonal scaling of the target (solid). The dotted lines show the local law ∝σ22k , and the vertical line is tmax=20 . Right: median number of sweeps (20 random scalings) for a column-sum error of 10−3 , as a function of τ .
Figure 7: Two channels (Theorem 3.1 and Proposition 7.1 ). Fraction of the energy of the residual-path signal HL:0x in the difference channel, for a random Gaussian input x∈R4×64 . The mixers are projected random logits ( M=1,3 ) or random isometric mixers in O1(n) (median of 10 draws). Under doubly stochastic mixing the energy of the differences decays geometrically. Under isometric mean-preserving mixing it stays constant, because every factor is an isometry of both channels.
Hyper-Connections (HC) improve residual networks by introducing learnable mixing across multiple residual streams, but unconstrained mixing leads to training instability. Manifold-Constrained Hyper-Connections (mHC) address this by enforcing approximate double stochasticity via Sinkhorn normalization, while mHC-lite ensures exact constraints through convex combinations of permutation matrices at the cost of factorial complexity. KromHC reduces this cost using Kronecker-product parameterizations, but restricts the mixing matrices to a structured submanifold of the Birkhoff polytope . We propose Transportation Birkhoff Polytope (TBP) parameterizations and their Recursive variants (RTBP), which construct exactly doubly stochastic mixing matrices with (n−1)2 degrees of freedom. Our approach avoids iterative normalization and combinatorial explosion while preserving full expressivity of the Birkhoff polytope. Empirical results on language model pre-training' demonstrate competitive performance with improved stability and scalability.
Manifold-constrained hyper-connections (mHCs) have recently been proposed as a principled extension of hyper-connections, where the residual mixing matrices are constrained to be doubly stochastic via projection onto the Birkhoff polytope. In practical mHC implementations, this constraint is enforced by Sinkhorn-Knopp iterations, and the backward pass relies on unrolling the iterative solver. This design introduces substantial computation and memory overhead, and may also yield inaccurate projections when the algorithm converges slowly on challenging inputs, undermining the intended norm-control and stability guarantees of mHCs. In this work, we focus on the practically important 4x4 Birkhoff projection setting and develop an end-to-end acceleration framework. By leveraging the dual formulation, we reduce the problem to a three-dimensional unconstrained convex problem and solve it with Newton's method, achieving fast convergence and high accuracy. For the backward pass, we replace the unrolled differentiation with implicit differentiation, yielding exact gradients without storing intermediate states. To exploit massive parallelism, we design a warp-level CUDA kernel that uses only register-level primitives, avoiding global and shared memory I/O. Extensive experiments against representative open-source baselines demonstrate that the proposed solver yields substantially more reliable doubly stochastic projections -- especially when the input magnitude is large -- and achieves significant end-to-end speedups (including the backward pass), reaching over 20x acceleration at large batch sizes while maintaining orders of magnitude smaller marginal errors.
Chenrui Wang, Yixuan Qiu
RenminSchoolUniversityof Statisticsof China · School of Statistics and Data Science & Institute of Big Data Research · Shanghai University of Finance and Economics
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.