Paper deep dive
JPmHC Dynamical Isometry via Orthogonal Hyper-Connections
Biswa Sengupta, Jinhua Wang, Leo Brunswic
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 90%
Last extracted: 7/20/2026, 10:11:35 PM
Summary
The paper introduces JPmHC (Jacobian-spectrum Preserving manifold-constrained Hyper-Connections), a deep learning framework designed to address training instability and scalability issues in Hyper-Connections (HC). By replacing identity skips with trainable linear mixers constrained to manifolds like Stiefel or Grassmann, JPmHC ensures dynamical isometry and prevents gradient pathologies. Key contributions include a free-probability analysis for Jacobian spectra, memory-efficient implicit differentiation, and a Stiefel-constrained mixer using Cayley transforms. Empirical results on ARC-AGI show JPmHC achieves faster convergence and higher accuracy than bistochastic baselines.
Entities (10)
Relation Signals (6)
JPmHC → evaluatedon → ARC-AGI
confidence 95% · Empirical evaluations on ARC-AGI demonstrate that JPmHC achieves faster convergence
JPmHC → uses → Cayley Transform
confidence 92% · JPmHC introduces a Stiefel-constrained mixer via Cayley transforms
JPmHC → constrains → Stiefel Manifold
confidence 90% · By constraining the mixer M on operator-norm-bounded manifolds (e.g., bistochastic, Stiefel, Grassmann)
Hyper-Connections → inspired → JPmHC
confidence 90% · As a flexible and scalable extension of HC, JPmHC advances spectrum-aware, stable, and efficient deep learning
JPmHC → improves → Dynamical Isometry
confidence 88% · JPmHC prevents gradient pathologies and enhances stability... JPmHC advances spectrum-aware, stable, and efficient deep learning
JPmHC → outperforms → Bistochastic
confidence 85% · JPmHC achieves faster convergence, higher accuracy, and lower computational cost compared to bistochastic baselines
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:Recent advances in deep learning, exemplified by Hyper-Connections (HC), have expanded the residual connection paradigm by introducing wider residual streams and diverse connectivity patterns. While these innovations yield significant performance gains, they compromise the identity mapping property of residual connections, leading to training instability, limited scalability, and increased memory overhead. To address these challenges, we propose JPmHC (Jacobian-spectrum Preserving manifold-constrained Hyper-Connections), a framework that replaces identity skips with a trainable linear mixer acting on n parallel streams while explicitly controlling gradient conditioning. By constraining the mixer M on operator-norm-bounded manifolds (e.g., bistochastic, Stiefel, Grassmann), JPmHC prevents gradient pathologies and enhances stability. JPmHC introduces three key contributions: (i) a free-probability analysis that predicts Jacobian spectra for structured skips, providing actionable design rules for mixer selection; (ii) memory-efficient implicit differentiation for fixed-point projections, reducing activation memory and synchronization overhead; and (iii) a Stiefel-constrained mixer via Cayley transforms, ensuring orthogonality without post-hoc normalization. Empirical evaluations on ARC-AGI demonstrate that JPmHC achieves faster convergence, higher accuracy, and lower computational cost compared to bistochastic baselines, with a rank-$p$ Grassmannian variant tracking between the two -- consistent with the spectral theory predictions. As a flexible and scalable extension of HC, JPmHC advances spectrum-aware, stable, and efficient deep learning, offering insights into topological architecture design and foundational model evolution. \newline \newline
Tags
Links
- Source: https://arxiv.org/abs/2602.18308v2
- Canonical: https://arxiv.org/abs/2602.18308v2
Trouble viewing inline? Open PDF directly →
Full Text
113,493 characters extracted from source content.
Expand or collapse full text
JPmHC Dynamical Isometry via Orthogonal Hyper-Connections Biswa Sengupta LLM Suite Team, JP Morgan Chase & Co. biswa.sengupta@jpmorgan.com Jinhua Wang LLM Suite Team, JP Morgan Chase & Co. jinhua.wang@jpmorgan.com Leo Brunswic LLM Suite Team, JP Morgan Chase & Co. leo.brunswic@jpmorgan.com February 2026 Abstract Recent advances in deep learning, exemplified by Hyper-Connections (HC), have ex- panded the residual connection paradigm by introducing wider residual streams and diverse connectivity patterns. While these innovations yield significant performance gains, they com- promise the identity mapping property of residual connections, leading to training instability, limited scalability, and increased memory overhead. To address these challenges, we pro- pose JPmHC (Jacobian-spectrum Preserving manifold-constrained Hyper-Connections), a framework that replaces identity skips with a trainable linear mixer acting on n paral- lel streams while explicitly controlling gradient conditioning. By constraining the mixer M on spectrum-controlled manifolds (e.g. Stiefel, Grassmann), JPmHC prevents gradient pathologies and enhances stability. JPmHC introduces three key contributions: (i) a free-probability analysis that predicts Jacobian spectra for structured skips, providing actionable design rules for mixer selection; (i) memory-efficient implicit differentiation for fixed-point projections, reducing activation memory and synchronization overhead; and (i) a Stiefel-constrained mixer via Cayley trans- forms, ensuring orthogonality without post-hoc normalization. Empirical evaluations on ARC-AGI demonstrate that JPmHC achieves faster convergence, higher accuracy, and lower computational cost compared to bistochastic baselines, with a rank-p Grassmannian variant tracking between the two—consistent with the spectral theory predictions. As a flexible and scalable extension of HC, JPmHC advances spectrum-aware, stable, and efficient deep learning, offering insights into topological architecture design and foundational model evo- lution. Disclaimer: This paper was prepared for informational purposes by the LLM Suite group of JP Morgan Chase and its affiliates (‘JPMC’) and is not a product of the Research De- partment of JP Morgan. JP Morgan makes no representation, warranty or undertaking whatsoever and disclaims all liability for the completeness, accuracy or reliability of the information contained herein. This document is not intended as investment research or in- vestment advice, or a recommendation, offer or solicitation for the purchase or sale of any security, financial instrument, financial product or service, or to be used in any way for evaluating the merits of participating in any transaction, and shall not constitute a solicita- tion under any jurisdiction or to any person, if such solicitation under such jurisdiction or to such person would be unlawful 1 arXiv:2602.18308v2 [cs.LG] 4 Mar 2026 © 2026 JP Morgan Chase & Co. All rights reserved. 1 Introduction The residual connection [He et al., 2016]—the per-layer update x l+1 = F (x l ) +x l —is a defining feature of modern deep learning, underpinning Transformers [Vaswani et al., 2017] and virtually every large-scale architecture deployed today. Its variants—Pre-Norm, DeepNorm [Wang et al., 2024]—have enabled training at thousands of layers by smoothing loss landscapes [Li et al., 2018] and stabilizing gradient flow [Pennington et al., 2017, Tarnowski et al., 2019]. However, the identity skip biases layerwise mappings toward the identity, anchoring the function class and limiting expressivity. A natural generalization replaces the identity skip with a learned linear map, x out = H res x + F (x),(1) increasing expressivity but risking gradient instability if the operator norm ∥H res ∥ and the singular spectrum of the end-to-end Jacobian are not controlled. To decouple expressivity from identity anchoring while preserving trainability at scale, Hyper-Connections (HC) [Zhu et al., 2024] split the hidden state into n parallel streams and mix them through a small n×n matrix. Let each stream live in R p and stack the streams so that x ∈ R n ⊗ R p ∼ = R np , with n ≪ p (typically n = 4, p = 512). The HC block takes the form x out = H res (x)⊗ I p x + H post (x)⊗ I p F H pre (x)⊗ I p x ,(2) where H res (x),H pre (x),H post (x) ∈ R n×n are small mixing matrices that depend on the input x. The extra cost scales with n and remains negligible since F is evaluated once per block on a stream mixture and re-distributed. The network learns which information flows where— a strictly richer connection pattern. The gains are most dramatic in the Mixture-of-Experts (MoE) setting, where HC halved the training tokens needed to match baseline on OLMoE and improved BBH and GSM8K by +7 points on DeepSeek’s 27B MoE model [Zhu et al., 2024, Xie et al., 2025]. Both MoE and HC are learnable routing mechanisms—one for tokens across experts, the other for residual streams across layers—and both face the same stability challenge: unconstrained, they diverge (signal gains exceeding 3000× at 27B scale [Xie et al., 2025]). Manifold-Constrained Hyper-Connections (mHC) [Xie et al., 2025] addressed this instability by projecting H res onto the Birkhoff polytope of doubly stochastic matrices via the Sinkhorn– Knopp iteration. Doubly stochastic mixers are appealing because (i) their operator norm is bounded by 1, preventing gradient explosion, and (i) they act as transport plans [Villani, 2003], intuitively preserving information across streams. At 27B-parameter scale, mHC demonstrated strong results with minimal overhead. However, two limitations remain: (1) operator-norm boundedness does not preclude vanishing gradients—a full singular-spectrum analysis of the end-to-end Jacobian is absent; and (2) backpropagating through iterative projections introduces memory and synchronization overhead in distributed training. This is where the argument breaks down. Training deep networks requires that the singular values of the input-output Jacobian J = Q L l=1 Y l remain concentrated near 1—a property called dynamical isometry [Saxe et al., 2014, Pennington et al., 2017]. Without it, expressivity capacity is lost or, worse, gradients may either explode or vanish exponentially. Tarnowski et al. [2019] proved that for scalar skip connections, dynamical isometry is universal: for any activation function, it is achieved when a single condition on the weight variance is met. A generalization of their free probability method to general twisting of the skip-connection seems within reach to go beyond operator-norm boundedness. For clarity-sake, we analyse the spectra for a simplified (2): x l+1 = A l n ⊗ I p x l + φ(W l x l ), N = np,(3) 2 where each A l n ∈ R n×n is a fixed mixing matrix (independent of x) for theoretical analysis. We extend the theory to the operator-valued setting via operator-valued free probabil- ity [Voiculescu, 1995, Dykema, 2007], where the Kronecker structure of (3) collapses the spectral problem from network width N = np to twist dimension n. This reveals two failure modes of doubly stochastic skip connections. The first is eigenvalue contraction: a doubly stochastic matrix has its Perron eigenvalue pinned at one, but generically all others lie strictly inside the unit disk, and deep composition drives|λ| L → 0. The second is eigenspace misalignment: eigen- bases of successive layers are unrelated, so composition scrambles directions and accelerates the collapse beyond what per-layer spectra predict. Together, these produce a partial spec- tral collapse—a growing fraction of the Jacobian’s singular values drifting toward zero—that no reparametrisation of the Birkhoff polytope can escape. Orthogonal matrices eliminate both failure modes: all eigenvalues lie on the unit circle, so no contraction is possible, and group closure under composition prevents misalignment at any depth. We propose replacing the Birkhoff constraint with the orthogonal group O(n), parametrised via the Cayley transform [Li et al., 2020a, Lezcano-Casado, 2019], which maps skew-symmetric matrices to orthogonal ones via (I−S)(I +S) −1 . Beyond spectral preservation, this provides a strictly richer function class (the linear span of O(n) is the full algebra M n (R), dimension n 2 , versus (n−1) 2 + 1 for the Birkhoff polytope) and implicit geometric nonlinearity from the curvature of the orthogonal manifold. Contributions. 1. Spectral diagnosis. We identify eigenvalue contraction and eigenspace misalignment as the mechanisms by which doubly stochastic skip connections break dynamical isometry, and show that this collapse converts to concrete capacity loss in modern training (spectral stalling). 2. Cayley-transform Stiefel projection. We instantiate an orthogonality-preserving mixer by projecting H res onto the Stiefel manifold via a small, fixed number of Cayley iterations (as few as s = 2), yielding norm-preserving mixing with exact gradients and negligible overhead [Li et al., 2020a, Lezcano-Casado, 2019]. 3. Grassmannian subspace mixer. We develop a rank-p variant with O(np) parameters that mixes through a learned p-dimensional subspace, optimized with a Cayley retraction for efficient Riemannian updates. 4. Implicit differentiation for fixed-point projections. We design a custom back- ward pass for iterative normalizations and projections (e.g., Sinkhorn for bistochastic constraints, Cayley for orthogonal constraints), reducing activation memory from O(T ) to O(1) and eliminating distributed data-parallel synchronization stalls, while remaining compatible with CUDA graphs and mixed precision [Eisenberger et al., 2022a]. 5. Operator-valued Dyson pipeline. We develop the first numerical implementation of the full operator-valued free probability pipeline—from the matrix Dyson equation through Dykema’s twisted S-transform multiplicativity to multi-layer spectral densities. 6. Experimental validation. We confirm the spectral predictions against Monte Carlo sim- ulation and validate the practical consequence on a modified Tiny Recursive Model [Jolicoeur- Martineau, 2025] evaluated on ARC-AGI-1 [Chollet, 2019]: orthogonal skip connections (Cayley) converge faster and reach higher accuracy than bistochastic ones (Sinkhorn), while the rank-p Grassmannian variant tracks between the two, consistent with the spec- tral theory predictions. 3 2 Spectral Analysis We now develop the spectral machinery for predicting the singular-value distribution of the end- to-end Jacobian in deep networks with structured skip connections. The key insight is that free probability [Voiculescu, 1991, Nica and Speicher, 2006] reduces the spectral analysis of L-layer compositions to fixed-point equations on order parameters, and the Kronecker structure A n ⊗I p further collapses the problem from network width N = np to twist dimension n. 2.1 Scalar Dyson equation and dynamical isometry Consider a standard residual network with scalar skip connection a and layer-wise update x l+1 = φ(W l x l ) + ax l . The linearized layer map is Y l = D l W l + aI N , where D l = diag(φ ′ (h l i )) and W l ∈ R N×N has i.i.d. Gaussian entries with variance σ 2 w /N. The end-to-end Jacobian is J = Q L l=1 Y l . In the mean-field limit (N → ∞), the activation derivatives D l concentrate around their expectation, making D l W l effectively isotropic [Pennington et al., 2017]. Free probability the- ory [Voiculescu, 1991] then predicts the limiting spectral density of J ⊤ J via the Cauchy transform G(z) := lim N→∞ 1 N ETr(zI − J ⊤ J ) −1 and the S-transform, which linearizes free multiplicative convolution: S J ⊤ J (w) = Q L l=1 S Y l⊤ Y l (w). One can deduce G from S and vice-versa, For scalar skip connections, Tarnowski et al. [Tarnowski et al., 2019] derived a scalar fixed- point equation for the single-layer Cauchy transform. The order parameter m(z) satisfies the scalar Dyson equation: m(z) = a z− σ 2 m(z) , where σ 2 = σ 2 w E[φ ′ ( √ qZ) 2 ],(4) and q is the forward signal variance, Z ∼N (0, 1). The Cauchy transform is G Y ⊤ Y (z) = m(z)/z. For L identical layers, the z 1 -mapping converts S Y ⊤ Y (w) L back to G J ⊤ J (z) without numerically fragile S-transform inversion. A network achieves dynamical isometry when the singular values of its Jacobian J concentrate near 1. For scalar skip connections, regardless of activation function or depth [Tarnowski et al., 2019], dynamical isometry is achieved via suitable scaling of layers weights. This universality breaks down for structured skip connections: when a is replaced by an n× n matrix A n , the scalar trace 1 N Tr averages over A n ’s spectral sectors. Scalar approximation distinguishes bistochastic and orthogonal mixer but predictions are inaccurate for bistochastic (mHC) and general linear (HC) mixers, see figure 1. 2.2 Operator-valued extension: Kronecker collapse Hyper-Connections [Zhu et al., 2024] replace the scalar skip with a Kronecker product A = A n ⊗I p , where A n ∈ R n×n mixes n parallel streams of dimension p, and N = np. The layer-wise Jacobian becomes Y l = (A l n ⊗ I p ) + D l W l ,(5) where W l ∈ R N×N remains isotropic Gaussian, but the skip structure is now block-diagonal with n× n blocks. Why scalar theory fails. The scalar Cauchy transform G(z) = 1 N Tr(zI N − M ) −1 computes an average over all N eigenvalues. When M = (A n ⊗ I p ) + noise, this trace averages over the n spectral sectors induced by A n , collapsing eigenvalue structure that is critical for gradient flow. For instance, if A n is bistochastic with eigenvalues1,λ 2 ,...,λ n where|λ i | < 1 for i≥ 2, the scalar theory sees only the average behavior, not the sector-wise contraction that drives vanishing gradients. 4 Operator-valued free probability. The solution is to work over the base algebra B = M n (C) [Voiculescu, 1995, Speicher, 1998]. We promote the order parameter from a scalar m(z) ∈ C to a matrix M (z) ∈ M n (C), and the Cauchy transform to a B-valued functional. The Kronecker structure A n ⊗ I p ensures that the self-consistent equation for M (z) depends only on A n and the noise variance σ 2 , not on the stream dimension p. Critically, the S-transform multiplicativity rule becomes twisted in the operator-valued set- ting [Dykema, 2007]: S XY (B) = S Y (B)S X S Y (B) −1 BS Y (B) , B ∈B.(6) This conjugation by S Y (B) encodes the eigenspace rotation between successive layers when A l n do not commute—precisely the misalignment effect absent in scalar theory. Proposition 2.1 (Kronecker collapse). Under the Kronecker structure Y l = (A l n ⊗ I p ) + D l W l and mean-field isotropy, the B-valued order parameter M (z) ∈ M n (C) defined for z ∈ M n (C) satisfies the matrix fixed-point equation M (z) = A h (M )· zI n − A h (z) ⊤ A h (z) −1 ,(7) where A h (z) := A n + σ 2 z is the dressed matrix. The scalar Cauchy transform is recovered by G(z) = 1 n Tr n zI n − A h (z) ⊤ A h (z) −1 . When n = 1, this reduces to (4). Proof sketch. The key steps are: (i) the Kronecker structure implies E[W l (A l n ⊗ I p )] = 0 by isotropy; (i) the self-energy Σ(z) has the block form Σ n ⊗ I p where Σ n ∈ M n (C); (i) inserting the ansatz M (z) = M n (z)⊗ I p into the Dyson-Schwinger equation and tracing over the p× p blocks yields (7). Computational complexity. Solving (7) requires Newton iteration in C n 2 at cost O(n 6 ) per step (matrix inversion dominates). Since n ≪ p (typically n = 4, p = 512), this collapses the spectral problem from O(N 3 ) = O((np) 3 ) to O(n 6 ), a reduction of factor (p/n) 3 ≈ 10 5 at typical scales. This makes exhaustive spectral analysis tractable for networks of arbitrary width N. 2.3 Numerical pipeline We now describe the computational methods for solving the scalar and operator-valued Dyson equations and extracting spectral densities. Scalar solver. For fixed z ∈ C + (upper half-plane), equation (4) is solved by Newton’s method with the iteration m (k+1) = m (k) − m (k) − a z−σ 2 m (k) 1 + aσ 2 (z−σ 2 m (k) ) 2 .(8) We sweep a grid of z-values from large |z| to small |z|, using each solution to seed the next (branch continuation). This ensures convergence even near the spectral edges where the Cauchy transform has poles. Convergence is typically achieved in 3–5 iterations with tolerance 10 −12 . Multi-layer: the z 1 -mapping. For L identical layers, the S-transform multiplicative prop- erty gives S J ⊤ J (w) = S Y ⊤ Y (w) L . The z 1 -mapping inverts this directly: given G Y ⊤ Y (z 1 ), we solve for w such that χ(w) := w+1 wS Y ⊤ Y (w) = z (the G ↔ S relation), then compute S J ⊤ J (w) = S Y ⊤ Y (w) L , and finally solve w+1 wS J ⊤ J (w) = z ′ to obtain G J ⊤ J (z ′ ). This avoids numerical S- transform inversion, which is ill-conditioned near w = 0. 5 Matrix Dyson solver. For Kronecker skip connections with fixed A n across layers, equa- tion (7) is a coupled system of n 2 complex equations. We vectorize M ∈ M n (C) to m ∈ C n 2 and apply Newton’s method: m (k+1) = m (k) − J F (m (k) ) −1 F (m (k) ),(9) where F (m) = m − vec A h (M )(zI n − A h (M ) ⊤ A h (M )) −1 and J F is the n 2 × n 2 Jacobian computed via automatic differentiation. The cost is O(n 6 ) per iteration due to the matrix inverse in (7). Branch continuation from large to small |z| remains essential for stability. Operator-valued multi-layer pipeline. For heterogeneous layers (A l 1 n ̸= A l 2 n ), the twisted S- transform multiplicativity (6) requires iterating the composition S J ⊤ J (B) = S Y 1⊤ Y 1 (B ′ )S Y 2⊤ Y 2 (B ′ ) with conjugation updates B ′ = S Y 2 (B) −1 BS Y 2 (B). Each conjugation requires solving the operator-valued G↔ S relation, itself a matrix fixed-point problem. The full pipeline has com- plexity O(Ln 10 ) per z-point (nested matrix inversions). In practice, for n≤ 4 and L≤ 100, this completes in < 1 second per z-point on a CPU. Validation. Figure 1 compares the theoretical predictions to Monte Carlo histograms of sin- gular values sampled from finite networks (n = 4, p = 25, 500 samples). At L = 1, the scalar theory (red curves) matches Monte Carlo perfectly for all mixer types. At L = 10, the scalar theory fails for bistochastic and Gaussian mixers, which develop spectral mass near zero (eigen- value contraction), while orthogonal mixers maintain dynamical isometry. The operator-valued theory correctly predicts the sector-wise collapse for bistochastic matrices. Spectral density extraction. The spectral density ρ(x) is recovered from the imaginary part of the Cauchy transform via the Stieltjes inversion formula: ρ(x) =− 1 π lim ε→0 + ImG(x + iε).(10) We evaluate G(x + iε) for small ε ≈ 0.01 on a dense real grid x ∈ [λ min ,λ max ] and extract ρ(x) =−ImG(x + iε)/π. Numerical integration confirms normalization R ρ(x)dx = 1 to within 10 −3 . 3 Cayley Twisted Skip-Connections The spectral analysis of Section 2 reveals two failure modes of doubly stochastic skip connections: eigenvalue contraction and eigenspace misalignment. Orthogonal matrices eliminate both—all eigenvalues lie on the unit circle and group closure prevents misalignment at any depth. We therefore constrain the residual mixer H res to the orthogonal group O(n) via an iterative Cayley transform [Li et al., 2020b,a, Lezcano-Casado, 2019]. 3.1 Iterative Cayley Projection The Cayley transform maps a skew-symmetric matrix W =−W ⊤ to an orthogonal matrix via (I−W/2)(I +W/2) −1 . The closed-form requires a matrix inverse that is expensive for batched, per-token computation. Following Li et al. [2020b], we replace the inverse with a fixed-point iteration that converges to the same retraction. Given an unconstrained parameter matrix ̃ H∈ R n×n : 1. Skew-symmetrize: W = ̃ H− ̃ H ⊤ , guaranteeing W∈so(n). 2. Initialize: Y 0 = I n + αW, with step-size α > 0 (default 0.1). 6 Figure 1: Scalar and OV theories vs. Monte Carlo singular value densities. Panels show four skip-connection types (n = 4 streams, p = 25 per stream, c 2 L = 0.05, η = 0.02, 500 samples) at depths L ∈ 1, 2, 10 (rows) for mixers A n ∈ Identity, Bistochastic, Orthogonal, (columns). The scalar Dyson prediction (dotted orange curve) matches Monte Carlo histograms at L = 1 for all cases. At L ∈ 2, 10, bistochastic and Gaussian mixers develop mass near zero (spectral collapse), while orthogonal mixers preserve dynamical isometry. Scalar theory fails while Operator-value theory is able to fully catch the spectrum, regularization parameter η used to reduce numerical instabilities smooth out the distribution: it pushes it away from zero and reduces the spikes thus increases the mass allocated to the 1.0 mode. The scaling Lc 2 = const ensures weights W l ∼N (0,σ 2 w /L) maintain constant forward signal variance. This normalization is shown to be accurate as spectra have a main mode bounded away from 0 an infinity. 3. Iterate s times: Y i+1 = I n + α 2 W (I n + Y i ), i = 0,...,s−1.(11) In practice s = 2 iterations suffice, achieving ∥Y ⊤ Y− I∥ max < 10 −3 (Section I). Each step is a single fused multiply-add (baddbmm), and all matrices are n× n with n = 4, so the overhead relative to the p-dimensional sub-layer F is negligible. 7 3.2 Layer Architecture A single linear projection produces three unconstrained n × n matrices per token from the flattened stream representation x flat ∈ R np : [ ̃ H pre | ̃ H post | ̃ H res ] = W fused LayerNorm(x flat ), W fused ∈ R 3n 2 ×np .(12) Each matrix is then projected onto its respective constraint manifold: H pre = softmax( ̃ H pre /τ, dim =−1), H post = softmax( ̃ H post /τ, dim =−2), H res = Cay s ( ̃ H res ). (13) The pre-mixer H pre is row-stochastic (aggregates streams), the post-mixer H post is column- stochastic (fans output back), and the residual mixer H res is orthogonal (norm-preserving skip). The forward pass implements (2) as: x in = H pre x streams ,y = F ( ̄ x in ),x out = H res x streams + H post ·(y⊗ 1 n ),(14) where ̄ x in denotes the stream-averaged input (a mean over the n streams) and F is the sub-layer (multi-head attention or feed-forward network), evaluated once on a single p-dimensional vector. The final combination is fused into a single baddbmm call. HC-Attention Residual Block HC-FeedForward Residual Block Input x∈ R N , N =qp Twisted Skip H res (x)⊗ I p · x Stream Project H proj (x)·x q→1 stream p dims Multi-Head Attention Layer(·) width p Stream Expand (·)⊗ 1 1→q streams back to N = qp Gate H out (x)·(·) + Layer Norm q parallel streams Twisted Skip H res (x)⊗ I p · x Stream Project H proj (x)·x q→1 stream p dims Feed-Forward Network Layer(·) width p Stream Expand (·)⊗ 1 1→q streams back to N = qp Gate H out (x)·(·) + Layer Norm q parallel streams Output x out ∈ R N Figure 2: Hyper-Connected Transformer Encoder Block. Block 1 (Multi-Head Attention, bot- tom) feeds into Block 2 (Feed-Forward, top). Within each block the input forks into a twisted skip path (thick arrows, H res ) and a compute path (thin arrows) that projects q→1 streams, applies the layer, expands 1→q, and gates via H out ; both paths merge at the + node before Layer Norm. 4 Experimental Setup 4.1 Task: ARC-AGI We evaluate JPmHC on the Abstraction and Reasoning Corpus (ARC-AGI) [Chollet, 2019], a benchmark designed to measure general fluid intelligence. Each task presents a small number of demonstration input–output grid pairs and one or more test inputs; the solver must infer the latent transformation rule and produce the exact output grid (Figure 3). Grids are rectangular 8 (a) Task 1: Input(b) Task 1: Output(c) Task 2: Input(d) Task 2: Output Figure 3: Two representative ARC-AGI tasks. Each task requires discovering a latent rule (here, pattern tiling and region filling) from a few demonstrations, then applying it to a novel test input. matrices of integers 0–9 (visualized as colors), with dimensions up to 30× 30. A task is solved only when every test output is reproduced cell-for-cell, including its dimensions. ARC-AGI is particularly suited to stress-test our spectral claims for two reasons. First, each task has a unique underlying rule, so the benchmark is resistant to memorization and demands systematic generalization—precisely the regime where gradient conditioning determines whether a model can learn compositional abstractions. Second, the all-or-nothing exact-match criterion amplifies the practical consequence of partial spectral collapse: even a small fraction of vanishing singular values can corrupt a few output cells and turn a near-correct grid into a failure. For all experiments, we use the full ARC-AGI-1 corpus of 1000 tasks, split evenly between training and evaluation (400 training, 400 evaluation, plus 200 for ablation and validation). This ensures that models are evaluated on held-out tasks with unseen rules, and that the results reflect true generalization rather than memorization. 4.2 Model and Training We adapt the Tiny Recursive Model (TRM) [Jolicoeur-Martineau, 2025], a 7M-parameter recursive transformer, by expanding each transformer block with n = 4 parallel streams. The attention and FFN residual sub-blocks are each wrapped by a JPmHC module (Figure 2): a read mapping H pre aggregates streams into a single sublayer input, a write mapping H post fans the output back to all streams, and a residual mixer H res ∈ R n×n —constrained to a chosen manifold—mixes the original streams before addition. Hidden dim per stream is d=512 (effective dim nd=2048). Two unique weight-tied blocks are each applied 6 times (12 total recursive passes) with Adaptive Computation Time (ACT) halting, yielding 4 unique JPmHC modules reused across all 12 recursions. This architecture is a stringent test bed for mixer design: the 12-fold weight-tied recur- sion means the same mixing matrix is composed with itself repeatedly, directly exposing the eigenvalue contraction and eigenspace misalignment phenomena analyzed in Section 1. All variants share identical training hyperparameters: AdamAtan2 optimizer [Kunstner et al., 2023] with lr 10 −4 , global batch size 768, bfloat16 mixed precision, and PyTorch DDP with torch.compile on 8× NVIDIA B200 GPUs. Full configuration details are provided in Section G. 4.3 Ablated Variants We compare four JPmHC mixer constraints, summarized in Table 1. The key design axes are the manifold constraint on the residual mixer H res and the resulting computational cost per module. 9 Table 1: JPmHC variant configurations and per-module compute cost. VariantManifoldKey param FLOPs/moduleNormConstraint SinkhornBirkhoff polytope T =20, k=16576RMSNorm sigmoid CayleyStiefel O(n) s=2, α=0.1256LayerNorm softmax GrassmannGr(n,p)p=2, α=0.01—LayerNorm softmax 4.4 Evaluation Metrics We report three complementary metrics: • Exact accuracy: fraction of tasks where the greedy prediction matches the ground-truth grid cell-for-cell—the strictest measure, directly sensitive to spectral health. • Pass@k (k ∈ 1, 2, 5, 10, 100, 1000): probability that at least one of k i.i.d. samples is correct, estimated via pass@k = E h 1− N−c k / N k i . • Eval LM loss: stablemax cross-entropy on output grid tokens, providing a smooth proxy for per-token prediction quality. 5 Results The Cayley and Sinkhorn variants have completed training (∼516K and ∼511K steps respec- tively, with near-identical compute budgets). The Grassmann variant has recently started train- ing and has completed ∼111K steps; its results are preliminary but already informative. Table 2: ARC-AGI concept evaluation metrics (best observed per metric). Cayley and Sinkhorn have completed training (∼516K and∼511K steps respectively); Grassmann training is ongoing (∼111K steps). ∗ Grassmann numbers are preliminary. MetricSinkhorn Cayley Grassmann ∗ Cayley/Sink Training Steps510,951 515,902110,756— Eval Accuracy86.5% 86.9%82.2%1.00× Exact Accuracy (greedy) 27.9% 31.4%12.8%1.13× Pass@136.5% 40.5%27.5%1.11× Pass@241.7% 45.4%31.2%1.09× Pass@545.5% 50.5%35.9%1.11× Pass@1046.8% 53.1%39.1%1.14× Pass@10051.9% 59.2%46.0%1.14× Pass@100056.1% 62.7%49.0%1.12× Eval LM Loss0.817 0.6431.0671.27× 5.1 Analysis Key Observations. 1. Cayley leads across all metrics at convergence. With both Cayley and Sinkhorn now trained to comparable step counts (∼516K and ∼511K), the Cayley variant achieves 40.5% pass@1 (best checkpoint at step 380,755) versus Sinkhorn’s 36.5% at step 406,903— a 1.11× gap. The advantage is more pronounced in exact-match accuracy (1.13×: 31.4% vs. 27.9%), confirming that Cayley produces more consistently correct full-grid predictions. 10 Table 3: Cayley evaluation metrics at each checkpoint. Bold indicates the best value for each metric. Training ran for ∼516K steps. Step Exact Acc. Pass@1 Pass@2 Pass@10 Pass@100 Pass@1000 68,95110.6%— 81,90312.7%— 94,85415.7%— 107,80517.6%— 161,95223.4%35.4% 40.1% 47.4%52.7%54.1% 174,90424.5%35.6% 40.3% 47.7%53.9%57.1% 188,95125.1%36.3% 39.6% 45.4%51.2%52.9% 201,90225.4%36.8% 40.7% 46.1%53.8%56.9% 214,85326.3%36.3% 40.4% 47.5%54.9%58.0% 231,95126.0%36.9% 39.8% 47.4%52.5%53.8% 244,90226.4%37.3% 41.1% 48.9%55.5%57.6% 257,85327.8%38.5% 41.6% 50.0%56.2%58.5% 270,80427.0%37.7% 41.7% 50.5%57.4%59.5% 283,75526.9%38.5% 43.1% 52.0%58.0%60.5% 296,70728.1%38.7% 43.8% 51.9% 59.2%62.0% 309,65928.3%39.0% 43.6% 52.6% 59.2%62.7% 328,95128.9%36.4% 40.3% 46.8%52.9%54.4% 341,90227.3%37.9% 42.5% 49.6%56.4%58.1% 354,85328.6%38.1% 43.8% 52.2%58.4%60.1% 367,80429.0%39.5% 43.5% 52.7%58.6%60.6% 380,75529.6% 40.5% 43.6% 53.1%58.9%61.4% 392,95229.3%37.7% 42.1% 49.6%55.1%56.6% 405,90329.2%38.7% 43.8% 51.7%57.7%60.5% 418,854 31.4%38.2% 45.4% 52.5%58.7%61.6% 431,80631.0%38.9% 45.1% 52.1%59.0%62.6% 456,95230.9%38.5% 43.4% 48.4%55.0%55.3% 469,90431.0%38.6% 43.3% 49.6%55.9%57.9% 482,85531.4%38.6% 43.9% 50.0%56.4%58.4% 502,95131.3%36.4% 41.5% 47.5%53.1%53.3% 515,90231.4%36.9% 42.5% 49.3%54.0%55.6% 2. Scaling with k. All three variants benefit from increased sampling budget with dimin- ishing returns. For Cayley: pass@1 → pass@1000 improves by +22.2 percentage points (p). For Sinkhorn: +19.6 p. For Grassmann (preliminary): +21.5 p. The pass@k ratio between Cayley and Sinkhorn narrows from 1.11× at k = 1 to 1.12× at k = 1000, indicating comparable prediction variance at convergence. 3. Sinkhorn narrowed the gap but did not close it. At the earlier reporting point (∼349K for Sinkhorn, ∼419K for Cayley), the pass@1 ratio was 1.19× and exact-accuracy ratio 1.41×. With Sinkhorn now extended to∼511K steps (matched compute), the pass@1 gap narrowed to 1.11× and exact-accuracy to 1.13×. Importantly, Sinkhorn’s best pass@1 (36.5%) and pass@1000 (56.1%) were achieved at steps ∼407K and ∼242K respectively; later checkpoints showed declining pass@k despite continued exact-accuracy gains, sug- gesting overfitting to greedy decoding at the expense of sampling diversity. 4. High token accuracy, divergent task accuracy. All variants exceed 82% per-token eval accuracy (Cayley 86.9%, Sinkhorn 86.5%, Grassmann 82.2%), yet task-level pass@1 11 k=1k=2k=5k=10k=100k=1000 0 20 40 60 40 . 5 45 . 4 50 . 5 53 . 1 59 . 2 62 . 7 36 . 5 41 . 7 45 . 5 46 . 8 51 . 9 56 . 1 27 . 5 31 . 2 35 . 9 39 . 1 46 49 Sampling Budget Pass@ k (%) Cayley (best ckpt) Sinkhorn (best ckpt) Grassmann ∗ (111K steps) Figure 4: Pass@k scaling comparison. Cayley consistently outperforms Sinkhorn across all sampling budgets k. Grassmann results are preliminary (∼111K steps vs. ∼500K+ for others); at matched step counts, Grassmann tracks ahead of Sinkhorn’s early trajectory. The gap between Cayley and Sinkhorn narrows at higher k, indicating Sinkhorn has higher prediction variance. ranges from 27.5% to 40.5%. This confirms that the advantage of orthogonal mixing lies in producing coherent complete solutions, not merely better per-token predictions. 5. Computational efficiency. The Cayley JPmHC module requires ∼2.25× fewer FLOPs than Sinkhorn (Table 6), achieving higher accuracy with lower per-module cost—a clear Pareto improvement. 6. Lower evaluation loss. The Cayley variant achieves a best evaluation LM loss of 0.643 compared to Sinkhorn’s 0.817—a 21% reduction (Figure 6). This 1.27× loss ratio exceeds the pass@1 ratio, suggesting that the accuracy advantage is driven by fundamentally better language modeling on the ARC grid tokens. Both variants show mild loss increase after their respective optima, consistent with slight overfitting at late training stages. 7. Faster convergence. The Cayley variant surpassed Sinkhorn’s final best exact-match accuracy (27.9%) at approximately step ∼297K—58% of Sinkhorn’s training budget. For pass@1, Cayley exceeded Sinkhorn’s best (36.5%) at step∼202K—only 40% of Sinkhorn’s budget—demonstrating substantially higher sample efficiency. 8. Grassmann: promising early trajectory. Despite only ∼111K steps of training, the Grassmann variant already achieves 27.5% pass@1 and 12.8% exact-match accuracy. At the comparable step count (∼113K), Sinkhorn had 22.2% pass@1 and 9.9% exact-match, while Cayley (at ∼108K) had 17.6% exact-match (pass@1 not yet measured). This places Grassmann’s early convergence rate between the two completed variants, consistent with its status as a rank-p orthogonal projection—a middle ground between full orthogonal mixing (Cayley) and full bistochastic mixing (Sinkhorn). 5.2 Compute-Accuracy Tradeoff † Grassmann backward includes Cayley retraction onn×p frame; values forn=4,p=2. 6 Related Work Hyper-Connections and structured skip connections Beyond the references on which our work is buit [Zhu et al., 2024, Xie et al., 2025], two concurrent works explore related direc- 12 0100200300400500 50 60 70 80 90 Training Steps (×10 3 ) Eval Accuracy (%) (a) Per-Token Accuracy Cayley Sinkhorn Grassmann ∗ 0100200300400500 0 10 20 30 Training Steps (×10 3 ) Exact Accuracy (%) (b) Exact-Match (Full Grid) Cayley Sinkhorn Grassmann ∗ Figure 5: Evaluation accuracy curves. (a) Per-token accuracy shows Cayley and Sinkhorn both exceeding 86% at convergence, with Grassmann tracking a steep early trajectory. (b) Exact- match accuracy reveals a persistent gap: Cayley plateaus at ∼31% while Sinkhorn saturates at ∼28%. Grassmann at 111K steps (12.8%) tracks ahead of Sinkhorn’s comparable point (9.9% at 113K). ∗ Grassmann training is ongoing. tions: Yang and Gao [2026] parametrises doubly stochastic matrices as Kronecker products of smaller bistochastic factors via the Birkhoff-von Neumann theorem, and Alonso [2026] proposes an operator-constrained framework. All of these remain within the Birkhoff polytope. Our work departs from this line by showing that the polytope’s contractive geometry causes spectral collapse, and proposes the orthogonal group as the correct constraint manifold. Optimization on Matrix Manifolds. Optimization over structured matrix sets has a rich history. The Cayley transform for parameterizing orthogonal matrices dates to Cayley [1846], with modern applications in neural networks [Li et al., 2020b]. Absil et al. [2008] provides principled gradient descent on curved spaces. Ablin et al. [2024] study the soft landing approach. The Sinkhorn operator and its differentiable variants have been extensively studied for optimal transport [Sinkhorn, 1967, Eisenberger et al., 2022b]. Orthogonal constraints in neural networks Orthogonal and unitary weight constraints have a rich history in recurrent networks, where they prevent gradient decay across time steps [Ar- jovsky et al., 2016, Wisdom et al., 2016]. Lezcano-Casado [2019] developed a general framework for gradient-based optimisation on matrix manifolds via trivializations (exponential map, Cay- ley transform), enabling efficient training with hard orthogonal constraints. We apply these parametrisation techniques not to weight matrices but to the skip connection matrices A q in Hyper-Connections, motivated by our spectral analysis showing that orthogonality of A q is the decisive property for preserving dynamical isometry. Signal propagation and mean-field theory The mean-field theory of deep networks Schoen- holz et al. [2017] studies the propagation of pre-activation moments, identifying the edge-of-chaos phase transition. Yang and Schoenholz [2017] extended this to residual networks. Pennington et al. [2017] connected signal propagation to the Jacobian’s singular value distribution, intro- ducing dynamical isometry for nonlinear networks. These works characterise signal propagation through scalar order parameters; our work computes the full spectral density, revealing failure modes—such as partial spectral collapse in bistochastic skip connections—that are invisible to mean-field analysis. To the best of our knowledge, operator valued free probability calculus 13 050100150200250300350400450500 0.5 1 1.5 2 2.5 Training Steps (×10 3 ) Eval LM Loss (per token) Cayley Sinkhorn Grassmann ∗ Figure 6: Evaluation LM loss (per-token cross-entropy, lower is better). The Cayley variant achieves the lowest loss (0.643 at 419K steps) with a 1.27× advantage over Sinkhorn’s best (0.817). Grassmann’s steep descent suggests it may approach Sinkhorn-level loss with continued training. Both Cayley and Sinkhorn show mild loss increase after their respective optima, suggesting slight overfitting at late training stages. for Neural Networks Jabobian spectrum comnputation has been introduced in Yang [2020]. In particular, Operator-valued freeness is proved there. Recursive Reasoning Models. The Hierarchical Reasoning Model (HRM) Wang et al. [2025] introduced recursive multi-step reasoning with deep supervision for puzzle-solving tasks. The Tiny Recursive Model (TRM) Jolicoeur-Martineau [2025] simplified HRM to a single 2-layer network with weight-tied recursion, achieving state-of-the-art results on ARC-AGI-1 (45% ac- curacy) with only 7M parameters. We adopt the TRM architecture as our evaluation platform, extending it with structured multi-stream mixing. 7 Discussion 7.1 Why Does Cayley Outperform Sinkhorn? With both variants now trained to convergence at matched compute budgets (∼516K and∼511K steps), the results show a consistent advantage for the Cayley JPmHC variant over implicit Sinkhorn, with the gap most pronounced in exact-match accuracy (1.13×) and evaluation LM loss (1.27×). We identify two contributing factors supported by both theory and empirical evidence: Bistochastic Induces Spectral Stalling. Spectral Stalling is the phenomenon by which directions associated with small singular values are ignored during the gradient descent: there is a hard cutoff of the spectrum. Effectively, the spectrum acts as a filter on the parameter space, the model is trained only on the subspace associated with singular values above the threshold. As shown in Section 2, Orthogonal and Bistochastic skip-connections show substantial differences in Jacobian singular spectrum. On the one hand, orthogonal skip-connection is undistinguishable from identity skip-connection, therefore dynamical isometry is achieved, the whole spectrum is above threshold. On the other hand, bistochastic skip-connection shows that more than 75% of the spectral mass is concentrated around 0, which suggest the same fraction of the weights are ignored: the model capacity is reduced. 14 Empirical Gradient Evidence. The gradient statistics from training corroborate the spec- tral stalling prediction. Despite achieving worse evaluation loss, the Sinkhorn variant exhibits ∼4× larger gradient norms than Cayley throughout training (average dense gradient norm: 0.39 vs. 0.10). The Grassmann variant, at its earlier training stage, shows even larger norms (0.84). This pattern—larger gradients with worse loss reduction—is consistent with a significant frac- tion of gradient energy being directed into spectral sectors with near-zero Jacobian singular values, where parameter updates produce little functional change. In contrast, Cayley’s smaller but more efficient gradients are concentrated in the full-rank spectral region, producing more effective parameter updates per step. The per-layer gradient statistics further support this: Sinkhorn’s maximum per-layer gradient norm (0.21 avg) is 4.2× larger than Cayley’s (0.05 avg), indicating that gradient energy in the Sinkhorn variant is not only larger in total but also more heterogeneously distributed across layers. Orthogonal Has Full Mixing Expressivity The respective intrinsic dimension of Orthog- onal matrices and Bistochastic matrices are q(q− 1)/2 and q 2 − 2q suggesting that Bistochastic matrices are more expressive. However, the latter form a polytope (a linear object) while the former form a spherical domain (non-linear). The span of Orthogonal matrices is the whole space of mixing matrices with dimension q 2 while the span of Bistochastic matrices has dimen- sion q 2 − 2q. The non-linear structure of orthogonal matrices has full mixing expressivity, while bistochastic matrices do not. Computational Efficiency. The Cayley JPmHC variant requires ∼2.25× fewer FLOPs per module (Table 6), enabling more optimization steps per unit of wall-clock time. 7.2 Grassmann: A Middle Ground The preliminary Grassmann results (∼111K steps) reveal an intriguing pattern. The Grassmann variant uses a rank-p orthogonal projector U ⊤ (Section J), which shares the orthogonality structure of Cayley but with reduced rank. At matched step counts, Grassmann tracks ahead of Sinkhorn’s early trajectory (27.5% vs. 22.2% pass@1 at∼111K steps) but behind Cayley’s. This ordering—Cayley > Grassmann > Sinkhorn at matched steps—is consistent with the spectral theory prediction: orthogonal projections (full-rank or rank-p) preserve more of the gradient spectrum than bistochastic matrices, with the full-rank Cayley variant preserving the most. The Grassmann variant also offers the lowest per-module FLOPs (72 vs. 256 for Cayley), making it a potentially attractive efficiency–accuracy trade-off. Whether Grassmann’s asymp- totic performance at convergence matches or exceeds Sinkhorn’s will be determined as training continues. 7.3 Implicit Differentiation: Correctness and Efficiency The custom backward pass for Sinkhorn (Section H) achieves two goals simultaneously: 1. Memory reduction: From O(T ) intermediate tensors to O(1) (only the output P). 2. DDP compatibility: Elimination of 128K autograd nodes that caused synchronization stalls in distributed training. The key insight is that Sinkhorn’s fixed-point structure admits a closed-form implicit deriva- tive. The Jacobian-vector product of the Sinkhorn operator at its fixed point can be expressed as a linear system involving only the fixed point P itself, bypassing the need to unroll through T iterations. 15 Self-Stabilization. An important property of the implicit gradient formula (23) is self-stabilization: the Hadamard product P⊙ (·) automatically zeros out gradient contributions to entries where P ij ≈ 0, preventing gradient flow through near-zero mixing weights. This is analogous to the “straight-through” behavior of hard attention, but arises naturally from the fixed-point structure. 7.4 Comparison with Hyper-Connections and mHC Our JPmHC framework extends both the original HC [Zhu et al., 2025] and the concurrent mHC [Xie et al., 2025] in several dimensions: 7.5 Limitations Late-Stage Overfitting. Both Cayley and Sinkhorn show increasing evaluation LM loss after their respective best checkpoints (Cayley after ∼310K, Sinkhorn after ∼420K), while exact- match accuracy continues to improve. This divergence between loss and accuracy, combined with near-zero training loss (< 0.002), suggests mild overfitting that manifests as reduced sampling diversity (declining pass@k) despite improved greedy performance. Pre/Post Architecture Confound. As noted, the Cayley and Sinkhorn variants differ in pre/post normalization and mapping architecture, making it impossible to attribute the entire performance gap to the manifold choice alone. Single Architecture. All experiments use the 7M-parameter TRM on ARC-AGI. General- ization to larger models, different architectures, and other tasks (language modeling, vision) remains to be validated. Incomplete Grassmann Training. The Grassmann variant has completed only ∼111K of the planned∼500K+ steps. While early results are promising, conclusions about its asymptotic performance relative to Cayley and Sinkhorn are premature. Small n. With n = 4 streams, the n× n mixing matrices are small enough for exact spectral analysis. Scaling to n ≥ 8 may require approximate methods for the operator-valued Dyson pipeline. 8 Conclusion We have presented JPmHC, a unified framework of manifold-constrained mixing strategies for multi-stream residual architectures, extending the mHC framework [Xie et al., 2025] with novel projection methods and efficient differentiation. Our contributions include implicit Sinkhorn differentiation, Cayley transform projection, and Grassmannian subspace optimiza- tion—addressing complementary challenges: the first eliminates DDP synchronization stalls, the second provides norm-preserving orthogonal mixing with exact gradients, and the third offers parameter-efficient subspace mixing via Riemannian optimization. 8.1 Summary of Key Results • The Cayley JPmHC variant achieves 40.5% pass@1 and 31.4% exact-match accuracy—a per- sistent 1.11×/1.13× advantage over Sinkhorn at matched compute budgets (∼500K+ steps each). 16 • The Cayley variant reaches a 21% lower evaluation LM loss (0.643 vs. 0.817) and surpasses Sinkhorn’s final best pass@1 (36.5%) at only 40% of Sinkhorn’s training budget, demonstrating superior sample efficiency. • The Cayley JPmHC module requires 2.25× fewer FLOPs than Sinkhorn while achieving higher accuracy—a Pareto improvement in both compute and quality. • The Sinkhorn variant reaches 36.5% pass@1 and 27.9% exact-match at ∼511K steps, signifi- cantly improving from earlier checkpoints but unable to close the gap to Cayley. • The Grassmann variant, at only ∼111K steps, already achieves 27.5% pass@1—exceeding Sinkhorn’s performance at matched step counts and offering the lowest per-module FLOPs (72 vs. 256 for Cayley). • Empirical gradient statistics corroborate the spectral stalling theory: the Sinkhorn variant exhibits 4× larger gradient norms than Cayley despite achieving worse loss, consistent with gradient energy dissipating in near-zero spectral sectors. • All JPmHC variants exceed 82% per-token accuracy, confirming the viability of structured mixing for recursive reasoning. 8.2 Broader Impact This work demonstrates that geometric structure—manifold constraints, group-theoretic anal- ysis, implicit differentiation—can be profitably applied to architectural components typically treated as unconstrained parameters. By restricting mixing matrices to well-understood math- ematical objects (orthogonal matrices, doubly-stochastic matrices, Grassmannians), we obtain models that are more computationally efficient and more effective. This approach is orthogo- nal to advances in attention mechanisms, normalization, and activation functions, suggesting potential for broader adoption in multi-stream architectures. 8.3 Future Work • Complete Grassmann training: Extend the Grassmann run to∼500K+ steps to determine its asymptotic performance relative to Cayley and Sinkhorn. • Pre/post ablation: Isolate the contribution of manifold choice from pre/post architecture differences. • Scale experiments: Larger models (n ≥ 8 streams, d ≥ 1024) and additional benchmarks (language modeling, ARC-AGI-2 [Chollet et al., 2025]). • Adaptive variant selection: Learn which mixing strategy to apply at each layer during training. • Overfitting mitigation: Investigate regularization strategies to prevent the late-stage loss/accuracy divergence observed in both Cayley and Sinkhorn. References Pierre Ablin, Simon Vary, Bin Gao, and Pierre-Antoine Absil. Infeasible deterministic, stochas- tic, and variance-reduction algorithms for optimization under orthogonality constraints. Jour- nal of Machine Learning Research, 25(389):1–38, 2024. URL http://jmlr.org/papers/v25/ 23-0451.html. 17 Pierre-Antoine Absil, Robert Mahony, and Rodolphe Sepulchre. Optimization Algorithms on Matrix Manifolds. Princeton University Press, 2008. A. Noguer I Alonso. Operator-constrained residual connections. Technical Report 6048614, SSRN, 2026. D. G. Anderson. Iterative procedures for nonlinear integral equations. Journal of the ACM, 12 (4):547–560, 1965. G. W. Anderson, A. Guionnet, and O. Zeitouni. An Introduction to Random Matrices. Cam- bridge Studies in Advanced Mathematics. Cambridge University Press, 2010. M. Arjovsky, A. Shah, and Y. Bengio. Unitary evolution recurrent neural networks. In Proceed- ings of the International Conference on Machine Learning, pages 1120–1128, 2016. L. Armijo. Minimization of functions having Lipschitz continuous first partial derivatives. Pacific Journal of Mathematics, 16(1):1–3, 1966. Z. Bai and J. W. Silverstein. Spectral Analysis of Large Dimensional Random Matrices. Springer Series in Statistics. Springer, New York, second edition, 2010. S. T. Belinschi and H. Bercovici. A new approach to subordination results in free probability. Journal d’Analyse Mathématique, 101:357–365, 2007. S. T. Belinschi, R. Speicher, J. Treilhard, and C. Vargas. Operator-valued free multiplicative convolution: analytic subordination theory and applications to random matrix theory. Inter- national Mathematics Research Notices, 2015(14):5933–5958, 2015. arXiv:1209.3508. S. T. Belinschi, T. Mai, and R. Speicher. Analytic subordination theory of operator-valued free additive convolution and the solution of a general random matrix problem. Journal of the European Mathematical Society, 19(8):2241–2312, 2017. H. Bercovici and D. Voiculescu. Free convolution of measures with unbounded support. Indiana University Mathematics Journal, 42(3):733–773, 1993. P. Biane. Processes with free increments. Mathematische Zeitschrift, 227:143–174, 1998. Z. Burda, A. Jarosz, G. Livan, M. A. Nowak, and A. Swiech. Eigenvalues and singular values of products of rectangular Gaussian random matrices. Physical Review E, 82(6):061114, 2010. Arthur Cayley. Sur quelques propriétés des déterminants gauches. Journal für die reine und angewandte Mathematik, 32:119–123, 1846. François Chollet. On the measure of intelligence. arXiv preprint arXiv:1911.01547, 2019. François Chollet, Mike Knoop, Greg Kamradt, Bryan Landers, and Hansueli Pinkard. ARC- AGI-2: A new challenge for frontier AI reasoning systems. arXiv preprint arXiv:2505.11831, 2025. K. Dykema. On the S-transform over a Banach algebra. Journal of Functional Analysis, 231(1): 90–110, 2006. arXiv:math/0501083v2. K. Dykema. Multilinear function series and transforms in free probability. Advances in Mathe- matics, 208(1):351–407, 2007. M. Eisenberger, A. Toker, L. Leal-Taixé, F. Bernard, and D. Cremers. A unified framework for implicit Sinkhorn differentiation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 509–519, 2022a. arXiv:2205.06688. 18 Marvin Eisenberger, Aysim Toker, Laura Leal-Taixé, Florian Bernard, and Daniel Cremers. A unified framework for implicit Sinkhorn differentiation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2022b. G. H. Golub and J. H. Welsch. Calculation of Gauss quadrature rules. Mathematics of Compu- tation, 23(106):221–230, 1969. U. Haagerup and S. Möller. The law of large numbers for the free multiplicative convolution. Operator Theory: Advances and Applications, 149:157–186, 2005. Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recog- nition (CVPR), pages 770–778, 2016. J. W. Helton, R. Far, and R. Speicher. Operator-valued semicircular elements: solving a quadratic matrix equation with positivity constraints. International Mathematics Research Notices, 2007(22):rnm086, 2007. Alexia Jolicoeur-Martineau. Less is more: Recursive reasoning with tiny networks. arXiv preprint arXiv:2510.04871, 2025. D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015. arXiv:1412.6980. Frederik Kunstner, Jacques Chen, Jonathan Wilder Lavington, and Mark Schmidt. Noise is not the main factor behind the gap between SGD and Adam on transformers, but sign descent might be. arXiv preprint arXiv:2304.13960, 2023. M. Ledoux. The Concentration of Measure Phenomenon, volume 89 of Mathematical Surveys and Monographs. American Mathematical Society, 2001. M. Lezcano-Casado. Trivializations for gradient-based optimization on manifolds. In Advances in Neural Information Processing Systems, volume 32, pages 9154–9164, 2019. Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss landscape of neural nets. Advances in neural information processing systems, 31, 2018. J. Li, F. Li, and S. Todorovic. Efficient Riemannian optimization on the Stiefel manifold via the Cayley transform. In International Conference on Learning Representations, 2020a. arXiv:2002.01113. Jun Li, Fuxin Li, and Sinisa Todorovic. Efficient riemannian optimization on the Stiefel manifold via the Cayley transform. In International Conference on Learning Representations (ICLR), 2020b. V. A. Marchenko and L. A. Pastur. Distribution of eigenvalues in certain sets of random matrices. Matematicheskii Sbornik, 72(4):507–536, 1967. R. A. Meyer, C. Musco, C. Musco, and D. P. Woodruff. Hutch++: Optimal stochastic trace estimation. Proceedings of the Symposium on Simplicity in Algorithms (SOSA), pages 142–155, 2021. arXiv:2010.09649. Y. Nakatsukasa, O. Sète, and L. N. Trefethen. The A algorithm for rational approximation. SIAM Journal on Scientific Computing, 40(3):A1494–A1522, 2018. Alexandru Nica and Roland Speicher. Lectures on the combinatorics of free probability, vol- ume 13. Cambridge University Press, 2006. 19 J. Nocedal and S. J. Wright. Numerical Optimization. Springer Series in Operations Research. Springer, New York, second edition, 2006. J. Pennington and P. Worah. Nonlinear random matrix theory for deep learning. In Advances in Neural Information Processing Systems, volume 30, 2017. J. Pennington, S. Schoenholz, and S. Ganguli. Resurrecting the sigmoid in deep learning through dynamical isometry: theory and practice. In Advances in Neural Information Processing Systems, volume 30, 2017. P. Pulay. Convergence acceleration of iterative sequences. The case of SCF iteration. Chemical Physics Letters, 73(2):393–398, 1980. N. R. Rao and R. Speicher. Multiplication of free random variables and the S-transform: the case of vanishing mean. Electronic Communications in Probability, 12:248–258, 2007. L. F. Richardson. The approximate arithmetical solution by finite differences of physical prob- lems involving differential equations, with an application to the stresses in a masonry dam. Philosophical Transactions of the Royal Society A, 210:307–357, 1911. A. M. Saxe, J. L. McClelland, and S. Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. arXiv preprint arXiv:1312.6120, 2014. S. S. Schoenholz, J. Gilmer, S. Ganguli, and J. Sohl-Dickstein. Deep information propagation. In International Conference on Learning Representations, 2017. Richard Sinkhorn. Diagonal equivalence to matrices with prescribed row and column sums. The American Mathematical Monthly, 74(4):402–405, 1967. R. Speicher. Combinatorial theory of the free product with amalgamation and operator-valued free probability theory. Memoirs of the American Mathematical Society, 132(627), 1998. Wojciech Tarnowski, Piotr Warchoł, Stanisław Jastrzębski, Jacek Tabor, and Maciej Nowak. Dynamical isometry is achieved in residual networks in a universal way for any activation function. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 2221–2230. PMLR, 2019. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Infor- mation Processing Systems, volume 30, 2017. C. Villani. Topics in Optimal Transport, volume 58 of Graduate Studies in Mathematics. Amer- ican Mathematical Society, Providence, RI, 2003. D. Voiculescu. Operations on certain non-commutative operator-valued random variables. Astérisque, 232:243–275, 1995. D. Voiculescu, K. J. Dykema, and A. Nica. Free Random Variables, volume 1 of CRM Monograph Series. American Mathematical Society, Providence, RI, 1992. Dan Voiculescu. Limit laws for random matrices and free products. Inventiones mathematicae, 104(1):201–220, 1991. H. F. Walker and P. Ni. Anderson acceleration for fixed-point iterations. SIAM Journal on Numerical Analysis, 49(4):1715–1735, 2011. 20 Guoxin Wang, Jiaqi Li, Yifan Sun, Xiang Chen, Chang Liu, Yang Wu, Ming Lu, Shuqiang Song, and Yasin Abbasi Yadkori. Hierarchical reasoning model. arXiv preprint arXiv:2506.21734, 2025. H. Wang, S. Ma, L. Dong, S. Huang, D. Zhang, and F. Wei. DeepNet: Scaling transformers to 1,000 layers. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2024. arXiv:2203.00555. S. Wisdom, T. Powers, J. Hershey, J. Le Roux, and L. Atlas. Full-capacity unitary recurrent neural networks. In Advances in Neural Information Processing Systems, volume 29, 2016. M. A. Woodbury. Inverting modified matrices. Memorandum Report, 42, 1950. Zhenda Xie, Yixuan Wei, Huanqi Cao, Chenggang Zhao, Chengqi Deng, Jiashi Li, Damai Dai, Huazuo Gao, Jiang Chang, Kuai Yu, Liang Zhao, Shangyan Zhou, Zhean Xu, Zhengyan Zhang, Wangding Zeng, Shengding Hu, Yuqing Wang, Jingyang Yuan, Lean Wang, and Wenfeng Liang. mHC: Manifold-constrained hyper-connections. arXiv preprint arXiv:2512.24880, 2025. G. Yang and S. Schoenholz. Mean field residual networks: on the edge of chaos. In Advances in Neural Information Processing Systems, volume 30, 2017. Greg Yang. Tensor programs i: Neural matrix laws. arXiv preprint arXiv:2009.10685, 2020. X. Yang and X. Gao. KromHC: Manifold-constrained hyper-connections with Kronecker-product residual matrices. arXiv preprint arXiv:2601.21579, 2026. Defa Zhu, Hongzhi Huang, Zihao Huang, Yutao Zeng, Yunyao Mao, Banggu Wu, Qiyang Min, and Xun Zhou. Hyper-connections. arXiv preprint arXiv:2409.19606, 2025. v2: 2512.24880. Z. Zhu, Y. He, L. Liu, J. Xu, Z. Xie, A. Dai, and D. Dai. Hyper-connections. arXiv preprint arXiv:2409.19606, 2024. Accepted at ICLR 2025. A Matrices Manifolds A.1 Doubly-Stochastic Matrices and the Birkhoff Polytope Definition A.1 (Doubly-Stochastic Matrix). A matrix P∈ R n×n is doubly stochastic if P ij ≥ 0 for all i,j, P1 = 1, and P ⊤ 1 = 1. The set of all n × n doubly-stochastic matrices forms the Birkhoff polytope B n . By the Birkhoff-von Neumann theorem, B n is the convex hull of the n! permutation matrices: B n = convP σ : σ ∈ S n ,(15) where (P σ ) ij = 1[σ(j) = i]. A.2 Sinkhorn-Knopp Algorithm The Sinkhorn-Knopp algorithm Sinkhorn [1967] projects an arbitrary non-negative matrix onto the Birkhoff polytope via alternating row and column normalization. In log-space (for numerical stability): 21 Algorithm 1 Sinkhorn-Knopp Projection (Log-Space) Require: Unconstrained logit matrix X∈ R n×n , iterations T Ensure: Doubly-stochastic matrix P∈B n 1: log M← clamp(X,−10, 10) 2: for t = 1,...,T do 3:log M← log M− LSE row (log M)▷ Row normalize 4:log M← log M− LSE col (log M)▷ Column normalize 5: end for 6: P← exp(log M) 7: return P where LSE row (A) ij = log P k exp(A ik ) broadcasts along rows. A.3 The Stiefel Manifold Definition A.2 (Stiefel Manifold). The Stiefel manifold St(n,p) is the set of n× p matrices with orthonormal columns: St(n,p) =U∈ R n×p : U ⊤ U = I p .(16) When p = n, St(n,n) = O(n) is the orthogonal group. Points on St(n,p) can be parametrized via the Cayley transform of skew-symmetric matrices. A.4 The Grassmann Manifold Definition A.3 (Grassmann Manifold). The Grassmann manifold Gr(n,p) is the set of p- dimensional subspaces of R n : Gr(n,p) = St(n,p)/O(p),(17) where O(p) acts by right multiplication. Two matrices U,V∈ St(n,p) represent the same point on Gr(n,p) if V = UQ for some Q∈ O(p). The canonical representation of a Grassmannian point is the orthogonal projector P = U ⊤ , which is invariant to the O(p) fiber action. A.5 Cayley Transform The Cayley transform maps skew-symmetric matrices to orthogonal matrices: Cay(W) = (I + W/2)(I− W/2) −1 ,(18) where W =−W ⊤ is skew-symmetric. This mapping is a diffeomorphism from the space of skew- symmetric matrices to the connected component of O(n) containing the identity (i.e., det = +1), minus the set where I− W/2 is singular. 22 B Detailed Algorithm Pseudocode B.1 Complete Sinkhorn Implicit Backward Algorithm 2 Complete Sinkhorn Implicit Backward Pass Require: Saved output P∈ R B×n×n , upstream gradient ∂ℓ ∂P Require: Number of Gauss-Seidel iterations k (default: 4n = 16) Ensure: Gradient ∂ℓ ∂M 1: G← ∂ℓ ∂P ▷ Upstream gradient 2: H← P⊙ G▷ Element-wise product 3: h row ← H· 1▷ Row sums: (B,n, 1) 4: h col ← H ⊤ · 1▷ Col sums: (B,n, 1) 5: v← 0▷ Initialize dual variable 6: for i = 1,...,k do 7: u← h row − P· v ⊤ ▷ Update u from v 8: v← h col − P ⊤ · u ⊤ ▷ Update v from u 9: end for 10: ∂ℓ ∂M ← H− u· P− v· P▷ Gradient w.r.t. logit matrix 11: return ∂ℓ ∂M B.2 Complete Cayley Transform Algorithm 3 Iterative Cayley Transform Require: Input matrix H∈ R B×n×n , step size α, iterations s Ensure: Orthogonal matrix Q∈ O(n) 1: ̃ H← H.view(B,n,n) 2: W← ̃ H− ̃ H ⊤ ▷ Skew-symmetrize 3: Y ← I + αW▷ Initialize: Y 0 = I + αW 4: for i = 1,...,s do 5: Y ← I + α 2 W(I + Y)▷ Fixed-point iteration via baddbmm 6: end for 7: Q← Y 8: return Q B.3 Grassmannian Riemannian Step Algorithm 4 Riemannian Gradient Step on Gr(n,p) Require: Basis U∈ St(n,p), Euclidean gradient ∇ E , momentum M Require: Step size α, momentum coefficient β 1 Ensure: Updated basis U ′ ∈ St(n,p) 1: ∇ hor ← (I− U ⊤ )∇ E ▷ Horizontal (Riemannian) gradient 2: M← β 1 M + (1− β 1 )∇ hor ▷ Momentum update 3: W← MU ⊤ − UM ⊤ ▷ Skew-symmetric generator 4: U ′ ← U▷ Initialize Cayley retraction 5: for i = 1,...,s do 6: U ′ ← U + α 2 W(U + U ′ )▷ Cayley retraction iteration 7: end for 8: return U ′ 23 C Spectral Gap Computation Details C.1 Algorithm for Exhaustive Spectral Gap Search Algorithm 5 Exhaustive Search for Maximum Spectral Gap Generating Set Require: Group order n, number of generators K Ensure: Generating set S ∗ ⊆ S n with |S ∗ | = K and maximum spectral gap 1: G← S n ▷ Symmetric group on n elements 2: elements← list(G), N ←|G| = n! 3: ∆ ∗ ← 0, S ∗ ←∅ 4: for each S ⊆ G with |S| = K do 5: if S does not generate G via BFS closure then 6:continue 7: end if 8: T← 0 N×N ▷ Transition matrix 9: for g ∈ G, σ ∈ S do 10:T[g,gσ] += 1/K▷ Right-multiplication walk 11: end for 12: λ 1 ,λ 2 ,...← eigenvalues(T), sorted by |λ| 13:∆←|λ 1 |−|λ 2 | 14: if ∆ > ∆ ∗ then 15:∆ ∗ ← ∆, S ∗ ← S 16: end if 17: end for 18: return S ∗ , ∆ ∗ C.2 Complexity Analysis For n streams and K generators: • Number of candidate subsets: n! K • BFS closure check: O(n!· K) per subset • Eigendecomposition: O((n!) 3 ) per valid generating set • Total: O n! K · (n!) 3 D JPmHC Module Parameter Counts With 4 unique JPmHC modules per model (2 per unique layer, weight-tied across 6 recursive cy- cles), the total JPmHC parameter overhead ranges from∼115K (Perm Mix) to∼410K (Cayley), which is 1.6–5.8% of the ∼7M total model parameters. E Convergence of Gauss-Seidel Solver The convergence rate of the Gauss-Seidel solver for the implicit Sinkhorn backward depends on the spectral radius ρ of the iteration matrix. For a doubly-stochastic matrix P with entries bounded away from 0 and 1, the spectral radius satisfies ρ < 1, ensuring convergence. Our default of k = 4n = 16 iterations achieves < 0.1% gradient error, which is sufficient for stable training with standard learning rates. 24 F Derivation of Dyson Equation for mHC (Scalar case) We derive the generalized Green’s function governing the singular value spectrum of random matrices of the form Y = A + X, where A is deterministic and X is an isotropic random matrix. Using a block linearization and a resolvent expansion, we prove a matrix Dyson equation (Schwinger–Dyson / Pastur equation) for a deterministic equivalent resolvent. Isotropy forces the self-energy to collapse to scalar multiples of the identity, reducing the spectral analysis to a small set of order parameters. F.1 Problem Setup Let Y = A + X ∈ R N×N , where A is deterministic and X is centered random. We study singular values of Y via S := Y Y ⊤ . For z ∈ C + define the Stieltjes transform G S (z) = 1 N Tr(zI − S) −1 . We will access G S through a block linearization. F.2 Block Linearization Define the 2N × 2N matrix L(z) = zI −Y −Y ⊤ I . Schur complement. The (1, 1) block of L(z) −1 equals the desired resolvent: (L(z) −1 ) 11 = (zI − Y Y ⊤ ) −1 . Hence G S (z) = 1 N Tr (L(z) −1 ) 11 . F.3 Generalized Green’s Function For a block matrix M = M 11 M 12 M 21 M 22 ∈ R 2N×2N , define the block trace bTr(M ) = 1 N Tr(M 11 ) 1 N Tr(M 12 ) 1 N Tr(M 21 ) 1 N Tr(M 22 ) . Define the generalized Green’s function G(z) = bTr E[L(z) −1 ] = g 11 g 12 g 21 g 22 . Then G S (z) = g 11 (z). 25 F.4 Isotropy Assumption We assume X is isotropic in the second-moment sense: there exists σ 2 > 0 such that E[XBX ⊤ ] = σ 2 N Tr(B)Ifor every deterministic B ∈ R N×N . This holds for i.i.d. Gaussian X ij ∼ N (0,σ 2 /N ) and more generally for left-right orthogonally invariant ensembles with the same covariance. F.5 Resolvent, Resolvent Identity, and the Dyson Equation This section expands the resolvent machinery and proves the Dyson equation used later. F.5.1 Definition of the resolvent For any (square) matrix H ∈ R m×m and any w ∈ C not in the spectrum of H, the resolvent is R H (w) := (wI m − H) −1 . In our setting, the primary object is the block resolvent of L(z): R(z) :=L(z) −1 . When we split L(z) into a deterministic part and a random perturbation, we also use the bare resolvent R 0 (z) :=L 0 (z) −1 . F.5.2 Splitting into deterministic and random parts Write L(z) = zI −A −A ⊤ I | z =: L 0 (z) − 0 X X ⊤ 0 |z =: X . Thus L =L 0 −X. F.5.3 Resolvent identity (exact) The following is a standard identity. Lemma (resolvent identity). If B is invertible and C is arbitrary such that B − C is invertible, then (B− C) −1 = B −1 + B −1 C(B− C) −1 . Apply this with B =L 0 and C =X: R =R 0 +R 0 XR.(19) Taking expectation (and using that R 0 is deterministic), E[R] =R 0 +R 0 E[XR].(20) Equation (20) is exact but not closed because E[XR] depends on correlations between X and R. 26 F.5.4 Dyson equation via a self-energy (Gaussian / Wick closure) To close (20), one introduces the self-energy operator Σ[·]. For Gaussian (or Wick-type) ensem- bles, the closure is governed by second moments. We state a standard large-N closure (often proved with Gaussian integration by parts / Stein’s lemma, or with cumulant expansions and planar diagrammatics). We use it here as the key computational step. Assumption (Gaussian/Wick closure). Assume X has entries with variance σ 2 /N and satisfies Wick’s rule (e.g. i.i.d. Gaussian). Then for resolvents R, the leading-order contribution to E[XR] can be written as E[XR] = Σ[E[R]] E[R] + o N (1), in an entrywise or normalized-trace sense (depending on the regularity assumptions). Under this closure, (20) becomes E[R]≈R 0 +R 0 Σ[E[R]] E[R]. Rearranging, R −1 0 − Σ[E[R]] E[R]≈ I, hence E[R]≈ L 0 − Σ[E[R]] −1 . This motivates defining a deterministic equivalentM(z) as the solution to the matrix Dyson equation: M(z) = L 0 (z)− Σ[M(z)] −1 . (21) In many models (including i.i.d. Gaussian), one can show E[R(z)]−M(z) → 0 in normalized trace, uniformly for Im(z)≥ η > 0. F.5.5 Computing the self-energy under isotropy Let M = M 11 M 12 M 21 M 22 . Because X is off-diagonal, Σ[M] is also off-diagonal at leading order. Under isotropy, the contraction identities imply: E[XM 21 X ⊤ ] = σ 2 N Tr(M 21 )I,E[X ⊤ M 12 X] = σ 2 N Tr(M 12 )I. This yields the self-energy map Σ[M] = 0 σ 2 m 12 I σ 2 m 21 I0 , m 12 := 1 N Tr(M 12 ), m 21 := 1 N Tr(M 21 ). Therefore, even for general deterministic A, isotropy collapses the random correction to scalar multiples of the identity, reducing the random-matrix effect to two scalar order parameters. 27 F.6 Closed Dyson Equation and Order Parameters Combining (21) with the isotropic self-energy gives M −1 = zI −A −A ⊤ I − 0 σ 2 m 12 I σ 2 m 21 I0 , i.e. M −1 = zI−(A + σ 2 m 12 I) −(A ⊤ + σ 2 m 21 I)I . The order parameters satisfy the self-consistency equations m 12 = 1 N Tr(M 12 ), m 21 = 1 N Tr(M 21 ). Finally, G S (z)≈ 1 N TrM 11 (z). F.7 Special Case: Scalar Skip Connection If A = aI, then the Dyson equation collapses from 2N× 2N to a 2× 2 system because all blocks commute and M ij are scalar multiples of I. F.8 Conceptual Takeaway The key mechanism is: 1. linearize Y Y ⊤ into a 2N × 2N operator, 2. write the exact resolvent identity, 3. close it via Wick contractions into a Dyson equation, 4. use isotropy to reduce the self-energy to scalars ×I. This identifies the minimal set of order parameters controlling the singular spectrum of Y = A + X. G Training and Architecture Details G.1 TRM Architecture Configuration G.2 Training Hyperparameters AdamAtan2 Optimizer. AdamAtan2 replaces the standard Adam update m t /( √ v t +ε) with atan2(m t , √ v t ), providing more stable gradient scaling and eliminating sensitivity to the ε hy- perparameter. For the Grassmann variant, we additionally employ a GrassmannianOptimizer wrapper that applies Riemannian gradient steps to the subspace basis parameters U. DeepSpeed compatibility. Although our framework supports DeepSpeed ZeRO-2/ZeRO-3, profiling showed that for 7M parameters, vanilla DDP with torch.compile is faster due to lower communication overhead. All JPmHC variants are validated compatible with both backends. 28 G.3 Permutation Basis Construction For the Perm Mix variant with n = 4 streams and K = 6 permutations, the default permutation basis is: 1. Identity e: (0, 1, 2, 3) 2. Adjacent transpositions: (1, 0, 2, 3), (0, 1, 3, 2) 3. Cyclic shifts: (1, 2, 3, 0), (3, 0, 1, 2) 4. Random fill (seed 42): remaining permutations sampled to reach K = 6 This basis includes generators of S 4 (adjacent transpositions generate the full symmetric group) while adding cyclic structure for efficient mixing. H Sinkhorn Variant with Implicit Differentiation The Sinkhorn variant projects the residual mixing matrix H res onto the Birkhoff polytope B n of doubly-stochastic matrices. Our key contribution is a custom backward pass that eliminates the autograd graph explosion of the standard Sinkhorn-Knopp implementation. H.1 Problem: Autograd Graph Explosion The standard implementation records all T Sinkhorn iterations in the PyTorch autograd graph. Each iteration involves two log-sum-exp operations, each of which PyTorch decomposes into multiple elementary ops (exp, sum, log, subtract). For T = 20 iterations on n× n = 4× 4 matrices, this creates approximately 128,000 backward nodes. These nodes produce only microsecond-scale GPU kernels—far too small to overlap with DDP AllReduce communication. Our profiling on NVIDIA B200 GPUs revealed that 55% of the total backward pass time was spent on DDP synchronization stalls, waiting for AllReduce operations to complete because no substantial compute was available to overlap with them. H.2 Solution: Custom Autograd Function with Implicit Differentiation We implement a custom torch.autograd.Function that decouples the forward and backward passes: Forward Pass. The Sinkhorn iterations run under torch.no_grad(), so no autograd graph is recorded. Only the final doubly-stochastic matrix P is saved for the backward pass: Algorithm 6 Implicit Sinkhorn — Forward Pass Require: Logit matrix X∈ R n×n , iterations T Ensure: P∈B n (doubly-stochastic) 1: with torch.no_grad():▷ No autograd recording 2:log M← clamp(X,−10, 10) 3: for t = 1,...,T do 4:log M← log M− LSE row (log M) 5:log M← log M− LSE col (log M) 6: end for 7: P← exp(log M) 8: save P for backward 9: return P 29 Backward Pass. Gradients are computed via implicit differentiation of the fixed-point con- ditions, following the framework of Eisenberger et al. Eisenberger et al. [2022b]. At the doubly- stochastic fixed point, the constraints are: P1 = 1 and P ⊤ 1 = 1.(22) Proposition H.1 (Implicit Sinkhorn Gradient Eisenberger et al. [2022b]). Let P = Π B (X) be the Sinkhorn projection of X onto the Birkhoff polytope. Given the upstream gradient ∇ P L, the gradient with respect to the input is: ∂L ∂X = P⊙ ∇ P L− u1 ⊤ − 1v ⊤ ,(23) where u∈ R n×1 and v∈ R 1×n solve the coupled linear system: Pv ⊤ + u = h row ,(24) P ⊤ u + v ⊤ = h col ,(25) with h row = (P⊙∇ P L)1, h col = 1 ⊤ (P⊙∇ P L).(26) H.3 Gauss-Seidel Solver for the Coupled System The coupled system (24)–(25) can be solved efficiently via Gauss-Seidel iteration. Starting from v (0) = 0: Algorithm 7 Implicit Sinkhorn — Backward Pass (Gauss-Seidel) Require: Saved P∈B n , upstream gradient ∇ P L Ensure: ∇ X L 1: H← P⊙∇ P L▷ Elementwise product 2: h row ← H1▷ Row sums, shape (...,n, 1) 3: h col ← 1 ⊤ H▷ Column sums, shape (..., 1,n) 4: k ← GaussSeidelIters(n)▷ See Proposition H.2 5: v← 0▷ Shape (..., 1,n) 6: for i = 1,...,k do 7: u← h row − Pv ⊤ ▷ Shape (...,n, 1) 8: v← h col − (P ⊤ u) ⊤ ▷ Shape (..., 1,n) 9: end for 10: ∇ X L← H− u⊙ P− v⊙ P 11: return ∇ X L H.4 Convergence Analysis Proposition H.2 (Gauss-Seidel Convergence Rate). For the coupled linear system arising from an n× n doubly-stochastic matrix P, the spectral radius of the Gauss-Seidel iteration matrix is bounded by: ρ≤ 1− 1 n .(27) To achieve a relative residual ε≤ 0.01, the required number of iterations is: k ≥ logε logρ = log 0.01 log(1− 1/n) .(28) This yields the following iteration counts for practical values of n: For our production configuration with n = 4 streams, k = 16 Gauss-Seidel iterations suffice. We enforce bounds k ∈ [10, 50] as a safety net. 30 H.5 Self-Stabilization Property An important property of the implicit gradient formula (23) is its self-stabilization: the factor P⊙ (·) ensures that entries where P ij ≈ 0 automatically produce near-zero gradients, regardless of the accuracy of u and v. This means no clamping of P is needed in the backward pass, and the formula is robust to numerical precision issues in the Sinkhorn forward iterations. Remark H.3. Empirically, we verified that clamped (P ij ≥ 10 −8 ) and unclamped versions produce identical gradients (cosine similarity difference < 10 −6 ) across all tested gradient mag- nitudes. H.6 Fused Projection Optimization All three mapping projections (φ pre ,φ post ,φ res ) are computed via a single fused linear layer: [φ pre |φ post |φ res ] = W fused · RMSNorm(x flat ),(29) where W fused ∈ R (n+n+n 2 )×nd . This reads the input tensor once instead of three times, reducing memory bandwidth by ∼ 3× and cutting kernel launch overhead from 3 to 1. H.7 Complexity Comparison I Cayley Transform Variant The Cayley variant replaces the doubly-stochastic constraint (Birkhoff polytope) with the or- thonormality constraint (Stiefel manifold). Instead of requiring H ⊤ res 1 = 1 and H res 1 = 1, we require H ⊤ res H res = I n , which provides norm-preserving stream mixing: ∥H res x∥ =∥x∥. I.1 Mathematical Formulation The Cayley transform (18) maps skew-symmetric matrices W =−W ⊤ to orthogonal matrices. However, the closed-form requires a matrix inverse (I − W/2) −1 , which is computationally expensive for batched computation and not amenable to GPU parallelism. Following Li et al. Li et al. [2020b], we use an iterative approximation. I.2 Iterative Cayley Transform Given an unconstrained matrix ̃ H∈ R n×n , the projection proceeds in three steps: Step 1: Skew-Symmetrization. Ensure the input lies in the Lie algebra of O(n): W = ̃ H− ̃ H ⊤ .(30) This guarantees W =−W ⊤ , which is necessary for the Cayley transform to map to O(n). Step 2: Initialization. Since we start from the identity (X = I n ), the initial estimate sim- plifies: Y 0 = I n + αW,(31) where α > 0 is a step-size parameter (default α = 0.1). This saves one matrix multiplication compared to the general case Y 0 = X + αWX. 31 Step 3: Fixed-Point Iteration. The iterate converges to the Cayley retraction: Y i+1 = I n + α 2 W(I n + Y i ), i = 0, 1,...,s− 1.(32) The full algorithm is: Algorithm 8 Iterative Cayley Transform Projection Require: Unconstrained matrix ̃ H∈ R n×n , step-size α, iterations s Ensure: Approximately orthonormal matrix Y ∈ R n×n (Y ⊤ Y ≈ I) 1: W← ̃ H− ̃ H ⊤ ▷ Skew-symmetrize 2: Y ← I n + αW▷ Initialize (saves one matmul since X = I) 3: for i = 1,...,s do 4: Y ← I n + α 2 W(I n + Y)▷ Fixed-point step via baddbmm 5: end for 6: return Y I.3 Properties of the Cayley Projection Proposition I.1 (Norm Preservation). For Y produced by Algorithm 8, ∥Yx∥ ≈ ∥x∥ for all x∈ R n , with the approximation improving with more iterations s. Proposition I.2 (Determinant). The determinant satisfies| det(Y)|≈ 1, approaching exactness as s→∞. In practice, s = 2 iterations suffice for deep learning applications Li et al. [2020b], achieving orthonormality deviation ∥Y ⊤ Y− I∥ max < 10 −3 . I.4 CUDA Graph Compatibility A critical implementation detail is pre-allocation of the identity matrix. When using torch.compile(mode=’reduce-overhead’), PyTorch captures CUDA graphs that require static tensor shapes. Dynamic calls to torch.eye() or torch.ones() during forward pass break graph capture. We solve this by registering the identity matrix as a persistent buffer during module initial- ization: self.register_buffer("_identity", torch.eye(n), persistent=False)(33) For batched operation (when the batch size B· L is known), we pre-expand: I batch = I n .unsqueeze(0).expand(B· L,n,n).contiguous().(34) This ensures the same memory is reused across forward passes, enabling CUDA graph cap- ture. I.5 Layer Architecture The full Cayley JPmHC layer computes three mapping matrices per token: 1. Fused projection: A single linear layer produces all three unconstrained matrices: [ ̃ H pre | ̃ H post | ̃ H res ] = W fused · LayerNorm(x flat ),(35) where W fused ∈ R 3n 2 ×nd and the output is split and reshaped to (B· L,n,n) for each. 32 2. Constraint projection: H pre = softmax( ̃ H pre /τ, dim =−1),(36) H post = softmax( ̃ H post /τ, dim =−2),(37) H res = Cay iter ( ̃ H res ).(38) Note that H pre and H post use softmax (row-stochastic and column-stochastic respectively, with temperature τ) rather than sigmoid, since the Cayley variant operates on full n× n mixing matrices for pre/post. 3. Forward computation: x pre = H pre · x streams ,(39) x in = mean(x pre , dim = stream),(40) y = F (x in ),(41) y streams = s· H post,sum ⊙ y,(42) x out = H res · x streams + y streams .(43) The final combination uses torch.baddbmm for a fused residual-plus-write operation, avoiding a separate bmm followed by addition. I.6 Comparison: Doubly-Stochastic vs. Orthonormal Mixing The norm-preserving property of orthonormal mixing is particularly beneficial for recursive models where the same layer is applied multiple times—it prevents representation collapse or explosion across recursion steps. J Grassmannian Variant The Grassmannian variant provides a parameter-efficient alternative to both Sinkhorn and Cayley by learning a rank-p subspace projection instead of a full n× n mixing matrix. The residual mapping is represented as: H res = U ⊤ , U∈ St(n,p), U ⊤ U = I p ,(44) where p ≤ n (default p = ⌊n/2⌋). This projector is idempotent (H 2 res = H res ), symmetric (H res = H ⊤ res ), and has rank exactly p. J.1 Parameter Efficiency The key advantage is the reduction in parameters for the residual mixing: While the savings are modest for n = 4, they become significant for larger n and when accumulated across all layers in a recursive model. J.2 Riemannian Optimization via Cayley ADAM Standard gradient descent in Euclidean space does not preserve the Stiefel constraint U ⊤ U = I p . We employ a full Riemannian optimization scheme that combines horizontal projection, momentum, and Cayley retraction. 33 J.2.1 Step 1: Horizontal Projection On the Grassmann manifold Gr(n,p) = St(n,p)/O(p), the tangent space at U decomposes into: T U St(n,p) = T hor U |z horizontal ⊕ T ver U |z vertical (fiber) .(45) The vertical space corresponds to rotations within the subspace (i.e., right multiplication by O(p)), which do not change the projector U ⊤ . The horizontal projection removes this component: ∇ hor = (I n − U ⊤ )∇ eucl =∇ eucl − U(U ⊤ ∇ eucl ).(46) This is essential for Grassmannian (as opposed to Stiefel) optimization: it ensures we move in directions that actually change the subspace, not just rotate the basis within it. J.2.2 Step 2: Cayley ADAM Momentum We maintain exponential moving averages of the horizontally-projected gradient: M t+1 = β 1 M t + (1− β 1 )∇ hor,t ,(47) v t+1 = β 2 v t + (1− β 2 )∇ 2 hor,t ,(48) where β 1 = 0.9, β 2 = 0.999, and ∇ 2 hor,t denotes elementwise squaring. With optional adaptive learning rate (analogous to Adam Kingma and Ba [2015]), the scaled momentum is: ˆ M t = M t /(1− β t 1 ) p v t /(1− β t 2 ) + ε .(49) Without adaptive scaling (the default), we simply use ˆ M t = M t . J.2.3 Step 3: Skew-Symmetric Direction The Cayley retraction requires a skew-symmetric direction matrix: W = ˆ M t U ⊤ − U ˆ M ⊤ t ∈ R n×n , W =−W ⊤ .(50) This W defines a curve on the Stiefel manifold through the current point U. J.2.4 Step 4: Iterative Cayley Retraction The Cayley retraction maps the tangent vector back to the manifold without an explicit matrix inverse: Y 0 = U + α ˆ M t ,(51) Y i+1 = U + α 2 W(U + Y i ), i = 0,...,s− 1.(52) After s iterations (typically s = 2), we update U← Y s . J.2.5 Step 5: QR Retraction (Post-Step Correction) For robustness under DDP gradient synchronization and mixed-precision training, we also apply a QR retraction after each optimizer step: Q,R = QR(U),(53) U← Q· diag(sign(diag(R))).(54) 34 The sign correction ensures a unique representative on the Stiefel manifold (the Q-factor from QR is unique up to column sign flips). The complete Riemannian optimization step is given in Algorithm 9. Algorithm 9 Cayley ADAM Riemannian Optimization Step Require: Current basis U∈ St(n,p), Euclidean gradient ∇ U L Require: Momentum state M,v,t; hyperparameters α,β 1 ,β 2 ,ε Ensure: Updated U∈ St(n,p) 1: t← t + 1 2:▷ Step 1: Horizontal projection 3: ∇ hor ←∇ U L− U(U ⊤ ∇ U L) 4:▷ Step 2: Momentum update 5: M← β 1 M + (1− β 1 )∇ hor 6: v← β 2 v + (1− β 2 )∇ 2 hor ▷ Optional adaptive LR 7: if adaptive LR then 8: ˆ M← M/(1−β t 1 ) √ v/(1−β t 2 )+ε 9: else 10: ˆ M← M 11: end if 12:▷ Step 3: Skew-symmetric direction 13: W← ˆ MU ⊤ − U ˆ M ⊤ 14:▷ Step 4: Iterative Cayley retraction 15: Y ← U + α ˆ M 16: for i = 1,...,s do 17: Y ← U + α 2 W(U + Y) 18: end for 19: U← Y 20:▷ Step 5: QR retraction (optional, for DDP compatibility) 21: Q,R← QR(U) 22: U← Q· diag(sign(diag(R))) J.3 DDP/DeepSpeed Integration The Grassmannian optimization is implemented as a post-step optimizer wrapper (GrassmannianOptimizer) that: 1. Automatically discovers all GrassmannianProjection modules in the model (handling DDP/DeepSpeed wrapper unwrapping). 2. After each standard optimizer step (optimizer.step()), applies QR retraction to project U back onto St(n,p). 3. Maintains its own state dict for momentum (M), second moment (v), and time step (t), supporting full checkpoint save/load. 4. All tensor operations use in-place ops and avoid CPU-GPU synchronization (the time step t is kept as a GPU tensor to avoid .item() calls that trigger sync). J.4 Geometric Interpretation The Grassmannian projector H res = U ⊤ filters the n-dimensional residual stream through a learned p-dimensional subspace. Geometrically: • Components of x within the subspace col(U) are preserved. 35 • Components orthogonal to the subspace are projected to zero. • The sublayer output is added back, potentially reintroducing orthogonal components. This provides an implicit form of information bottleneck in the residual stream, where the model learns which p directions carry the most useful information across layers. K Computational Pipeline Overview The spectral density of the end-to-end Jacobian J = Y L ·Y 1 is computed through a modular pipeline comprising four stages: 1. Signal propagation. A forward recursion determines the per-layer operating point (pre- activation variance q ℓ , post-activation variance v ℓ , and self-energy e ℓ ), following the mean- field theory of Schoenholz et al. [2017], Pennington et al. [2017]. 2. Per-layer Dyson equation. For each layer ℓ, the Stieltjes transform G ℓ (z) of the squared singular-value distribution of Y ℓ = A ℓ + q σ 2 ℓ X ℓ is obtained from a self-consistent subordi- nation equation Tarnowski et al. [2019], Voiculescu [1991]. In the scalar case (q = 1) this re- duces to a single complex equation; for Kronecker-structured skip connections (A = A q ⊗I p ) it becomes a q× q matrix equation, following the operator-valued framework of Speicher [1998], Belinschi et al. [2017], Helton et al. [2007]. 3. Free multiplicative convolution. The per-layer Stieltjes transforms are composed to obtain the L-layer transform G L (z). Three composition strategies are available: the z 1 - mapping for identical layers (derived from the S-transform of Voiculescu [1991], Voiculescu et al. [1992]), the Belinschi–Speicher subordination iteration for heterogeneous scalar lay- ers Belinschi et al. [2015], and the Dykema twisted multiplicativity theorem for non- commuting operator-valued layers Dykema [2006]. 4. Spectral inversion. The density is recovered via the Stieltjes inversion formula, option- ally refined by Richardson extrapolation Richardson [1911] or A rational approxima- tion Nakatsukasa et al. [2018]. The appropriate solver is selected automatically based on the twist dimension q, depth L, layer homogeneity, and available hardware. L Activation Functions and Signal Propagation The activation function φ and its derivative φ ′ enter the theory through three Gaussian moments, computed by Gauss–Hermite quadrature Golub and Welsch [1969] with 100 nodes: ψ(v) = E Z∼N(0,1) φ ′ ( √ v Z) 2 ,(55) κ(v) = E Z∼N(0,1) φ( √ v Z) 2 ,(56) φ(v) = E Z∼N(0,1) φ( √ v Z) .(57) The function ψ controls the self-energy (effective noise variance) and κ the second moment of the layer output. The effective cumulant appearing in the Dyson equation is c 2 = σ 2 w ψ(q). These moments are well-defined for standard activation functions including ReLU, tanh, sigmoid, and leaky-ReLU; see Pennington and Worah [2017] for the general framework and Tarnowski et al. [2019] for the residual-network specialisation. 36 The operating point of each layer is determined by a forward recursion over the pre-activation variance q ℓ , the post-activation variance v ℓ , and the self-energy e ℓ Schoenholz et al. [2017], Yang and Schoenholz [2017], Tarnowski et al. [2019]: q ℓ = σ 2 w,ℓ (v ℓ−1 + e ℓ−1 ),(58) e ℓ = ψ(q ℓ ),(59) v ℓ = κ(q ℓ ) + ∥A ℓ ∥ 2 F p v ℓ−1 ,(60) with initial conditions v 0 (the input variance) and e 0 = ψ(σ 2 w,1 v 0 ). The per-layer Dyson self- energy is then σ 2 ℓ = σ 2 w,ℓ ψ(q ℓ ). The recursion accepts either a shared or per-layer skip matrix A ℓ and weight standard deviation σ w,ℓ . M Scalar Dyson Equation Solver M.1 Subordination form Under the isotropy assumption, the Stieltjes transform G(ζ) = 1 p Tr (ζI − Y T Y ) −1 of a single layer Y = A 0 + √ σ 2 X satisfies the subordination equation Tarnowski et al. [2019], Voiculescu [1991], Marchenko and Pastur [1967] G(ζ) = u 1 p p X i=1 1 ω− s i , u = 1− σ 2 G(ζ), ω = ζ u 2 ,(61) wheres i p i=1 are the eigenvalues of A T 0 A 0 , precomputed once in O(p 3 ) time. This is an instance of the general subordination phenomenon in free probability Biane [1998], Belinschi and Bercovici [2007]. M.2 Newton iteration Defining G free (ω) = 1 p P i (ω− s i ) −1 , the fixed-point residual is F (G,ζ) = uG free (ω)− G.(62) The analytical Jacobian is ∂F ∂G =−σ 2 G free (ω) + 2σ 2 ζ u 2 1 p p X i=1 1 (ω− s i ) 2 − 1.(63) Newton’s method is applied with Armijo backtracking Armijo [1966], Nocedal and Wright [2006] using step sizes α∈1, 1 2 , 1 4 , 1 8 to prevent divergence near spectral edges. M.3 Batch solver The batch solver uses a two-pass strategy: 1. Sequential continuation pass: sweep the z-grid from right to left (large to small |z|), using the converged solution at point j as the initial guess for point j +1. This continuation technique Nocedal and Wright [2006] provides a good warm start in approximately 5 Newton iterations per point. 2. Vectorised refinement: apply Newton iterations to all grid points simultaneously in a batched array operation until global convergence. 37 M.4 Woodbury acceleration When A 0 = aI +UV T has low-rank perturbation of rank r, the Woodbury matrix identity Wood- bury [1950] allows precomputation of the eigenvalues of A T 0 A 0 in O(pr 2 ) time, reducing setup cost from O(p 3 ) for large p. N Operator-Valued Dyson Solver For Kronecker-structured skip connections A = A q ⊗I p with A q ∈ R q×q and N = qp, the condi- tional expectation E B = id q ⊗ tr p reduces the Dyson equation to a q×q matrix problem within the operator-valued free probability framework Speicher [1998], Voiculescu [1995], Belinschi et al. [2017]. N.1 Scalar subordination path Because the noise is isotropic in the p-directions, the self-energy is proportional to I q : Σ = σ 2 G scalar I q , where G scalar = 1 q Tr(G (B) ). The q× q Green’s function is then G (B) (z) = u (ωI q − A T q A q ) −1 , G scalar = 1 q Tr G (B) ,(64) with u = 1−σ 2 G scalar and ω = z u 2 . This is a scalar subordination equation in G scalar with the same structure as (61), but summing over the q eigenvalues of A T q A q . N.2 Matrix spectral parameter For general b∈ M q (C) (needed by the Ψ-inversion in the operator-valued S-transform, cf. §O.2), a 2-scalar Schur complement iteration solves for the self-energies (g 11 ,g 22 ): G (B) (b) = b− σ 2 g 22 I q − u −1 A T q A q −1 ,(65) u = 1− σ 2 g 11 , where g 11 = 1 q Tr(G (B) ) and g 22 is similarly defined. This Schur reduction, inspired by the fixed- point characterisation of Helton et al. [2007], costs O(q 3 ) per iteration (dominated by the q× q matrix inverse), compared to O(q 6 ) for a full q 2 -dimensional Newton. Typically 5–30 iterations with adaptive damping suffice. O S-Transform and Free Multiplicative Convolution O.1 Scalar S-transform The S-transform, introduced by Voiculescu [1991], Voiculescu et al. [1992], is defined implicitly by S z G(z)− 1 = G(z) z G(z)− 1 ,(66) and linearises free multiplicative convolution Bercovici and Voiculescu [1993], Nica and Speicher [2006]: S AB (w) = S A (w)S B (w).(67) The computational procedure involves three steps: 1. Cauchy-to-S: compute S(w) from G(z) by solving w = zG(z)− 1 for z via Newton’s method with continuation threading. 38 2. S-product: form the pointwise product Q ℓ S ℓ (w). 3. S-to-Cauchy: recover G(z) from S(w) by solving the implicit equation z S(w)w−w−1 = 0 for w given z. See Rao and Speicher [2007], Haagerup and Möller [2005] for analytic properties of the S- transform relevant to the vanishing-mean case. O.2 Operator-valued S-transform (Dykema) For the subalgebra B = M q (C), the operator-valued S-transform Dykema [2006], Voiculescu [1995] is defined via the Ψ-inversion problem: given W ∈ M q , find b∈ M q such that bG (B) (b)− I q = W, S (B) (W ) = G (B) (b)W −1 .(68) Dykema’s twisted multiplicativity theorem (Dykema [2006], Theorem 1.1) gives the composition rule for non-commuting layers: S (B) xy (W ) = S (B) y (W ) S (B) x S (B) y (W ) −1 W S (B) y (W ) .(69) For L layers, the cumulative S-transform is built by folding from layer L (innermost) to layer 1 (outermost): S [L] (W ) = S L (W ), S [k...1] (W ) = S [k+1...1] (W ) S k S [k+1...1] (W ) −1 W S [k+1...1] (W ) , k = L−1,..., 1.(70) This twisted fold is the operator-valued analogue of the scalar product (67) and reduces to it when all A q,ℓ commute. The non-commutativity of the twisted product is essential for capturing the eigenspace misalignment effect described in the main text. P Multi-Layer Composition: Identical Layers When all L layers are identical (same A and σ 2 ), the L-layer Stieltjes transform G L (z L ) is related to the single-layer G 1 (z 1 ) by the subordination (“z 1 -mapping”) formula, derived from the S-transform identity (67) Burda et al. [2010], Tarnowski et al. [2019]: G L (z L ) = G 1 (z 1 ) L (z 1 G 1 (z 1 )− 1) L−1 ,(71) z L = z 1 z 1 − 1/G 1 (z 1 ) L−1 .(72) This avoids the numerically fragile S-transform round-trip G→ S → S L → G. 39 Algorithm 10 z 1 -mapping for identical layers Require: Single-layer solver for G 1 (z), depth L, grid z j n j=1 with Im(z j ) > 0 Ensure: G L (z j ) for j = 1,...,n 1: z 1,prev ← None 2: for j = 1,...,n do 3:Construct multi-start guesses: z 1,prev , geometric range z 1/L j ,...,z j , fixed heuristics 4: for each guess z (0) 1 do 5:if L≥ 5 then 6:Set w ← log(z 1 )▷ Log parameterisation 7:Newton on f log (w) = w + (L−1) logω− logz j 8:else 9:Newton on f (z 1 ) = z 1 ω L−1 − z j ▷ ω = z 1 − 1/G 1 (z 1 ) 10:end if 11:Armijo backtracking: halve step until |f| <|f prev | and Im(z 1 ) > 0 12: end for 13:Select z ∗ 1 with smallest residual 14: G L (z j )← G 1 (z ∗ 1 ) L / (z ∗ 1 G 1 (z ∗ 1 )− 1) L−1 15: z 1,prev ← z ∗ 1 ▷ Continuation seed 16: end for 17: Interpolate isolated unconverged points (linear in Re/Im) The analytical Jacobian df/dz 1 is obtained via implicit differentiation on the Dyson equation: dG 1 dz 1 =− ∂F/∂z ∂F/∂G ,(73) where F (G,z) = 0 is the Dyson residual (62), eliminating finite-difference approximations. Q Multi-Layer Composition: Heterogeneous Layers Q.1 Subordination iteration (Belinschi–Speicher) For heterogeneous layers with distinct A ℓ and σ 2 ℓ , the composition uses the subordination it- eration of Belinschi, Speicher, Treilhard, and Vargas Belinschi et al. [2015]. For each spectral parameter z, the algorithm iterates in the w-domain: z L (w) = w + 1 w Q L ℓ=1 S ℓ (w) ,(74) where S ℓ (w) = G ℓ (z ℓ )/w with z ℓ satisfying z ℓ G ℓ (z ℓ )− 1 = w. Newton’s method on w drives z L (w)→ z. 40 Algorithm 11 Subordination for heterogeneous layers Require: Per-layer solvers G ℓ (·) L ℓ=1 , grid z j Ensure: G L (z j ) for each j 1: Sweep z-grid from large to small (continuation) 2: for each z = z j do 3:Initialise w from previous converged point or w ← z G 1 (z)− 1 4: for iter = 1,...,N max do 5:for ℓ = 1,...,L do 6:Solve z ℓ G ℓ (z ℓ )− 1 = w for z ℓ ▷ Scalar Newton 7:S ℓ (w)← G ℓ (z ℓ )/w 8:end for 9:z L (w)← (w + 1)/(w Q ℓ S ℓ (w)) 10:if |z L (w)− z| < tol· max(|z|, 1) then 11:G L (z)← (w + 1)/z; break 12:end if 13:Newton update: w ← w− (z L (w)− z)/(dz L /dw)▷ FD derivative 14: end for 15: end for Convergence is probed on a small subset (10 uniformly-spaced points); if ≥ 80% converge, the subordination iteration is run on the full grid. Otherwise, the algorithm falls back to the S-transform round-trip G ℓ → S ℓ → Q S ℓ → G L . This fallback exploits the analyticity properties established in Belinschi and Bercovici [2007], Bercovici and Voiculescu [1993]. Q.2 Anderson/DIIS acceleration Fixed-point iterations for subordination are optionally accelerated by Anderson mixing Ander- son [1965], Walker and Ni [2011], also known as DIIS (Direct Inversion in the Iterative Sub- space) Pulay [1980] in the computational chemistry literature. Given a history of iterates x k and residuals r k with r k = g(x k )− x k , the accelerated estimate minimises ∥ P i c i r i ∥ 2 subject to P i c i = 1, solved via an (m+1)× (m+1) linear system with Lagrange multiplier (history depth m = 5). R Operator-Valued Multi-Layer Pipeline For Kronecker-structured heterogeneous layers with non-commuting A q,ℓ , the full operator- valued pipeline computes G (B) L (z)∈ M q (C) via a triple-nested Newton structure: 1. Outer Newton (over W ∈ M q , q 2 complex unknowns): solve F (W ) := (W + I q ) S prod (W )W −1 − z I q = 0,(75) where S prod (W ) is the twisted S-product (70). 2. Middle loop: evaluates the twisted fold (70) over L layers, applying Dykema’s composi- tion rule (69) at each step. 3. Inner Newton (per-layer Ψ-inversion): for each layer, solves the Ψ-inversion problem (68) bG (B) (b)−I q −W twisted = 0 for b∈ M q via Newton’s method with multiple initial guesses. Each iteration requires a matrix Dyson solve (§N). Convergence is managed via: 41 1. Analytical Jacobians (where feasible) with finite-difference fallback (step h = 10 −7 ); see Nocedal and Wright [2006] for the general theory of inexact Newton methods. 2. Damped Newton with Armijo backtracking Armijo [1966]: α∈1, 0.5, 0.25, 0.1, 0.05, 0.02, 0.01. 3. Continuation threading: solutions (W,b ℓ ,M ℓ ) from the previous z-point warm-start the next, reducing per-layer Newton iterations from ∼20 (cold) to ∼3–5 (warm). 4. Per-layer caching: each layer carries its converged (b,M ) pair through all nesting levels, enabling efficient re-use across the middle-loop evaluations. For large spectral grids, the z-points are sorted in descending order and partitioned into chunks, with a sequential pre-sweep seeding every K-th chunk boundary. Chunks are then processed in parallel with work-stealing scheduling. S Spectral Density Recovery S.1 Stieltjes inversion The spectral density is recovered from the Stieltjes transform via the standard inversion for- mula Anderson et al. [2010], Bai and Silverstein [2010]: ρ(x) =− 1 π Im G(x + iη) ,(76) followed by clipping to [0,∞) and normalisation to unit integral. S.2 Richardson extrapolation To sharpen spectral edges without reducing η to the point of numerical instability, a Neville- tableau Richardson extrapolation Richardson [1911] is applied. Densities are computed at η k = η base · 2 k for k = 0,...,n levels − 1, then combined via T k,j = 4 j T k,j−1 − T k−1,j−1 4 j − 1 , T k,0 = ρ η k (x),(77) assuming O(η 2 ) error scaling. The final estimate is T n levels −1,n levels −1 . S.3 A rational approximation An alternative high-accuracy path fits a barycentric rational approximant to G(z) via the A algorithm Nakatsukasa et al. [2018] on Chebyshev nodes, then evaluates the approximant on the full grid. This avoids grid-based artifacts and provides uniform accuracy near spectral edges with typically O(20–50) evaluations of G. T GPU-Accelerated Solvers All solvers described above admit natural GPU parallelisation over the spectral grid z j , since each z-point involves an independent fixed-point problem. The GPU implementation provides: 1. Batched scalar Dyson solver: the subordination equation (61), Newton updates, and Armijo backtracking are expressed as batched tensor operations, processing all z-points simultaneously. 2. Batched operator-valued Dyson solver: the Schur complement iteration (65) is par- allelised over z-points with batched q× q matrix inversions. 42 3. Multi-layer pipeline: the two-pass strategy (continuation + refinement) for identical layers, and the w-domain Newton for heterogeneous layers, are adapted with per-chunk GPU batching and work-stealing. Memory management uses heuristic batch sizing based on available GPU memory, with automatic fallback to smaller batches upon memory exhaustion. U Monte Carlo Validation Theoretical predictions are validated against direct Monte Carlo simulation of the random matrix product: 1. Construct n samples realisations of J = Y L ·Y 1 with Y ℓ = A ℓ + D ℓ W ℓ , where W ℓ has i.i.d. standard Gaussian entries and D ℓ = σ ℓ I. 2. Compute eigenvalues of J T J for each realisation. 3. Pool all eigenvalues and estimate the density via Gaussian KDE or histogram. 4. Compare to the theoretical prediction using L 1 , L 2 , L ∞ , and Kolmogorov–Smirnov metrics. The GPU variant uses batched eigendecomposition with automatic batch-size estimation based on available VRAM. As the network width N → ∞, the empirical spectral distribution converges to the theoretical prediction by the concentration of measure phenomenon Ledoux [2001], Anderson et al. [2010]. V Operator-Valued Spectral Density Panel The validation panel (Figure 1) displays a 3× 4 grid comparing the theoretical spectral density with Monte Carlo simulation for Kronecker-structured skip connections A = A q ⊗ I p across depths L ∈ 1, 2, 10 and four families of q× q twist matrices: identity, random bistochastic, random Haar-orthogonal, and normalised Gaussian. The full pipeline is detailed in Algorithm 12. All panels use q = 4, p = 25 (so N = qp = 100), and a fixed noise budget L· σ 2 = 0.05, giving per-layer self-energy σ 2 ℓ = 0.05/L. The imaginary regularisation is η = 0.02. Monte Carlo sample counts are n samples ∈300, 200, 100 for L∈1, 2, 10 respectively. V.1 Operator-valued z 1 -mapping For identical layers with L > 1, the OV multi-layer composition exploits a key structural prop- erty: because the noise is isotropic in the p-directions, the subordination variable b 1 remains on the scalar manifold b 1 = z 1 · I q . The scalar subordination equation z L = z 1 z 1 − 1/G scalar (z 1 ) L−1 , G scalar = 1 q Tr G (B) 1 (z 1 ) ,(78) is solved for z 1 by multi-start Newton iteration (as in Algorithm 10, using the analytical deriva- tives from the OV Dyson solver §N), and the L-layer q×q Green’s function is reconstructed via eigendecomposition: G (B) L (z L ) = G (B) 1 (z 1 ) L z 1 G (B) 1 (z 1 )− I q −(L−1) ,(79) where the matrix powers are computed via eigendecomposition of the q× q matrices G (B) 1 and z 1 G (B) 1 − I q . 43 V.2 Panel generation algorithm For clarity, Algorithm 12 describes the computation for a single twist matrix A q and depth L with L identical (homogeneous) layers sharing the same A q and σ 2 . The full 3×4 panel is obtained by repeating this procedure over L∈1, 2, 10 and A q ∈identity, bistochastic, orthogonal, Gaussian. Algorithm 12 OV spectral density: single cell (homogeneous layers) Require: Twist matrix A q (q× q), depth L, width p, self-energy σ 2 = c/L, regularisation η Ensure: Singular-value density plot with theory vs. MC overlay 1: A full ← A q ⊗ I p ▷ N × N Kronecker expansion — Monte Carlo phase — 2: for s = 1,...,n samples do 3: Y ← Q 1 ℓ=L (A full + √ σ 2 W s,ℓ )▷ W s,ℓ : i.i.d. N (0, 1/N ) entries 4:Record eigenvalues of Y T Y 5: end for 6: Pool all eigenvalues → empirical distribution — Adaptive grid calibration — 7: x max ← max 1.5· max(MC eigenvalues), 10 8: x-grid ← uniform [0.01, x max ], max(400, ⌊40x max ⌋) points — Theory phase (OV Dyson + z 1 -mapping) — 9: Precompute eigenvalues of A T q A q ▷ O(q 3 ), once 10: z-grid ← x-grid +iη 11: if L = 1 then 12: G scalar,j ← OV-Dyson(A q , σ 2 , z-grid)▷ §N 13: else 14: G scalar,j ← OV-z 1 -Mapping(A q , σ 2 , L, z-grid)▷ Alg. 10 / (79) 15: end if — Spectral density recovery — 16: ρ(x j )←− 1 π Im G scalar,j ; clip to [0,∞); normalise R ρ = 1 — Change of variables — 17: ρ sv (σ)← 2σ ρ(σ 2 ); renormalise — Validation — 18: W 1 ← Wasserstein-1(theory, MC) in eigenvalue domain 19: Plot: MC histogram + theory curve ρ sv ; annotate W 1 V.3 Heterogeneous layers When layers have distinct (non-commuting) twist matrices A q,1 ,...,A q,L with per-layer self- energies σ 2 1 ,...,σ 2 L , the z 1 -mapping of Algorithm 12 no longer applies. Two operator-valued composition methods are available, selected automatically by the dispatcher (§X): OV subordination (default). The scalar subordination variable w still lives in C (isotropic noise), so the per-layer S-products remain scalar. Algorithm 13 details the procedure. Twisted S-transform (fallback). For validation or when the subordination Newton diverges, the full Dykema twisted fold S (B) [L] (W ) = S L (W )· S L−1 (S −1 L WS L )· is used, with an outer Newton over W ∈ M q (C) (q 2 complex unknowns). 44 Algorithm 13 OV spectral density: single cell (heterogeneous layers) Require: Twist matrices A q,1 ,...,A q,L (q× q), per-layer self-energies σ 2 1 ,...,σ 2 L , width p, reg- ularisation η Ensure: Singular-value density plot with theory vs. MC overlay 1: A full,ℓ ← A q,ℓ ⊗ I p for ℓ = 1,...,L — Monte Carlo phase — 2: for s = 1,...,n samples do 3: Y ← Q 1 ℓ=L (A full,ℓ + q σ 2 ℓ W s,ℓ ) 4:Record eigenvalues of Y T Y 5: end for 6: Pool all eigenvalues → empirical distribution — Adaptive grid & theory phase — 7: Calibrate x-grid from MC support (as in Alg. 12, lines 8–9) 8: Build per-layer solvers: scalar D ℓ = DysonSolver(A q,ℓ , σ 2 ℓ ), matrix M ℓ = MatrixDysonSolver(A q,ℓ , σ 2 ℓ ) — OV subordination (sweep large z to small) — 9: for each z j in z-grid (descending, with continuation) do 10:Initialise w from previous point (or bootstrap w = z j G 1 (z j )− 1) 11: repeat 12:for ℓ = 1,...,L do 13:Solve z ℓ G ℓ (z ℓ )− 1 = w for z ℓ ▷ Newton 14:S ℓ ← G ℓ (z ℓ )/w 15:end for 16:z L (w)← (w + 1)/ (w· Q ℓ S ℓ ) 17:Newton update: w ← w− z L (w)− z j / ∂z L /∂w ▷ FD derivative 18: until |z L (w)− z j | < tol 19: G scalar,j ← (w + 1)/z j 20: end for — Spectral density recovery & plotting — 21: ρ, ρ sv , W 1 , plot (as in Alg. 12, lines 16–19) When q = 1, the pipeline reduces to the scalar z 1 -mapping of §P. W Stochastic Trace Estimation For large-dimensional problems where explicit eigendecomposition is prohibitive, the Hutch++ algorithm Meyer et al. [2021] provides stochastic trace estimation: Tr(A)≈ Tr(Q T AQ) + 1 s s X i=1 g T i (I − Q T )A (I − Q T )g i ,(80) where Q is obtained from the QR decomposition of AΩ for a random Gaussian matrix Ω∈ R n×k , andg i are i.i.d. complex Gaussian vectors (normalised to unit norm). Applied to the resolvent A = (zI−M ) −1 , this yields stochastic estimates of Tr (zI−M ) −1 = N G(z), enabling Stieltjes transform computation without full diagonalisation. X Solver Selection The computational pipeline automatically selects the appropriate solver based on problem pa- rameters: 45 ConditionMethod q = 1, L = 1Scalar Dyson (§M) q = 1, L > 1, identical layersz 1 -mapping (§P) q = 1, L > 1, heterogeneous layers Subordination iteration (§Q) q > 1, L = 1Operator-valued Dyson (§N) q > 1, L > 1OV multi-layer pipeline (§R) GPU availableGPU-accelerated variants (§T) Y Numerical Stability Techniques Several techniques are employed throughout the pipeline to ensure numerical stability: 1. Continuation threading: every grid-based solver seeds each new z-point with the con- verged solution from the previous point Nocedal and Wright [2006], reducing Newton iterations from ∼20 to ∼3–5. 2. Multi-start initialisation: critical solvers try multiple initial guesses (continuation seed, heuristic estimates, geometric range, fixed fallbacks) and select the solution with smallest residual, mitigating the basin-of-attraction problem inherent in Newton’s method. 3. Armijo backtracking Armijo [1966]: Newton steps are damped by halving the step size until the residual decreases, preventing divergence near spectral edges and singular points. 4. Logarithmic parameterisation: for L ≥ 5, the z 1 -mapping uses w = log(z 1 ) to com- press O(e L ) dynamic range to O(L), avoiding overflow in (z 1 − 1/G 1 (z 1 )) L−1 . 5. Tikhonov regularisation: when the Jacobian matrix in the operator-valued Newton has condition number > 10 10 , a regularisation term λI with λ = 10 −8 ∥J∥ F is added. 6. Non-finite value guarding: solver outputs are checked for non-finite values; any detected NaN or infinity is replaced by a safe fallback, preventing propagation through the pipeline. 7. Post-processing interpolation: isolated unconverged grid points (< 10% of total) are linearly interpolated in Re/Im parts from neighbouring converged values. 46 Table 4: Sinkhorn evaluation metrics at each checkpoint. Bold indicates the best value for each metric. Training ran for ∼511K steps; Sinkhorn continued to improve in exact-match accuracy through the final checkpoint. Step Exact Acc. Pass@1 Pass@2 Pass@10 Pass@100 Pass@1000 52,9511.4%5.2%7.5%12.1%14.9%16.1% 65,9023.1%9.0% 11.6% 16.2%21.4%22.6% 78,8545.0%11.6% 16.2% 21.1%26.8%29.5% 99,9518.5%19.1% 22.8% 29.4%36.6%39.6% 112,9029.9%22.2% 26.2% 33.1%40.6%43.9% 125,85311.2%23.7% 28.0% 35.5%42.8%45.5% 138,80513.0%26.6% 30.3% 38.6%44.9%48.8% 151,75714.3%28.6% 31.7% 40.3%46.0%49.8% 164,70915.9%29.0% 33.0% 41.2%47.4%51.5% 177,66116.2%30.0% 34.6% 41.6%47.5%51.7% 190,61217.4%30.8% 34.7% 42.3%48.5%53.0% 203,56317.1%31.0% 35.5% 42.8%48.8%53.8% 216,51418.2%31.3% 35.5% 42.4%50.5%54.1% 229,46518.6%31.7% 36.0% 43.0%50.9%55.4% 242,41618.7%31.2% 35.6% 43.4%51.1%56.1% 258,95119.3%30.8% 33.4% 42.3%47.5%48.0% 271,90219.4%31.9% 34.5% 44.1%48.9%51.0% 284,85320.3%33.6% 37.6% 44.7%50.1%52.2% 297,80520.1%34.1% 38.0% 45.5%51.4%54.3% 322,95221.2%31.5% 35.0% 42.5%48.1%49.1% 335,90322.0%33.6% 37.6% 44.6%50.1%51.4% 348,85422.2%34.0% 38.5% 45.6%50.2%53.0% 361,80522.6%34.7% 39.5% 46.4%51.7%54.3% 374,75623.0%36.3% 40.5% 45.9% 51.9%54.5% 393,95222.9%33.8% 38.5% 43.9%48.8%49.8% 406,90324.0% 36.5% 41.1% 46.3%51.1%52.6% 419,85426.0%36.0% 41.7% 46.8%51.9%53.4% 434,95226.8%35.7% 40.1% 45.8%49.5%49.8% 459,95226.7%35.0% 40.4% 43.6%48.0%48.2% 472,90427.7%36.1% 41.2% 45.6%48.9%49.4% 485,85527.6%35.9% 41.2% 46.4%49.3%50.6% 510,951 27.9%36.4% 40.6% 45.1%48.9%49.5% 47 Table 5: Grassmann evaluation metrics at each checkpoint (training ongoing at ∼111K steps). Bold indicates the best value for each metric. At matched step counts, Grassmann tracks ahead of Sinkhorn’s early trajectory (cf. Sinkhorn at 113K: 9.9% exact, 22.2% pass@1). Step Exact Acc. Pass@1 Pass@2 Pass@10 Pass@100 Pass@1000 12,9510.2%0.6%1.3%3.2%3.9%4.6% 28,9510.8%3.6%5.5%9.9%12.4%14.0% 41,9022.6%8.9% 12.1% 16.0%18.6%21.1% 58,9515.5%15.7% 18.6% 24.0%29.4%31.5% 71,9027.4%20.0% 24.1% 30.4%35.7%38.4% 84,8549.6%22.2% 27.3% 34.7%40.4%43.8% 97,80510.9%26.1% 29.1% 36.5%43.0%45.9% 110,756 12.8% 27.5% 31.2% 39.1% 46.0%49.0% Table 6: Forward and backward pass compute estimates per JPmHC module. VariantForward FLOPs Backward FLOPs Total per module Sinkhorn (implicit) O(Tn 2 ) = 320 O(kn 2 ) = 256576 CayleyO(sn 3 ) = 128 O(sn 3 ) = 128256 GrassmannO(np) = 8 O(n 3 ) † = 6472 Table 7: Feature comparison with HC and mHC. FeatureHC / mHCJPmHC (Ours) Mixing parameterization Learned n× n / Sinkhorn Stiefel, Grassmann, Birkhoff Manifold constraintBirkhoff polytope only Stiefel, Grassmann, Birkhoff Implicit differentiation—Sinkhorn backward Riemannian optimization—Grassmann (Cayley ADAM) Spectral analysis—Generating set selection CUDA graph compatibleNot addressedAll variants Distributed trainingDualPipe (mHC)DDP + DeepSpeed ZeRO Table 8: Tractability of exhaustive spectral gap search. n n! n! 4 Matrix size Feasibility 3 6156× 6 < 0.01s 4 2410,62624× 24 ∼ 0.3s 5 120 8,214,570120× 120 ∼ 1 hour 6 720 ∼ 10 10 720× 720 Intractable Table 9: Parameter count per JPmHC module (n = 4, d = 512, nd = 2048). VariantD res Total params Notes Sinkhorn n 2 = 16 (n + n + 16)× nd + 36 = 49,188 Fused φ Cayley3n 2 = 4848× nd + LayerNorm = 102,400 Fused φ + LN Grassmann np = 8(n + n + 8)× nd + 28 = 32,796 p = 2 Perm MixK = 6(n + n + 6)× nd + 22 = 28,694 + perm indices 48 Table 10: Gauss-Seidel convergence for n = 4 with typical P matrices. Iterations k Relative residual ∥r∥/∥r 0 ∥ Gradient error (%) 17.5× 10 −1 43.2 25.6× 10 −1 28.1 43.2× 10 −1 11.5 81.0× 10 −1 2.8 161.0× 10 −2 0.08 321.0× 10 −4 < 0.001 Table 11: Full TRM architecture configuration. ParameterValue Total parameters∼7M Hidden dimension d512 Number of streams n4 Effective hidden dim nd2048 Unique transformer layers (weight-tied) 2 Recursive cycles per layer6 Total recursive applications12 Attention heads8 FFN expansion ratio4× Halting mechanismAdaptive Computation Time (ACT) ACT max recursion depth16 ACT exploration probability0.1 JPmHC modules per layer2 (pre-attention, pre-FFN) Total JPmHC modules4 (shared via weight tying) Positional encodingRoPE Flash AttentionEnabled (with SDPA fallback) Table 12: Training hyperparameters (shared across all variants). HyperparameterValue OptimizerAdamAtan2 [Kunstner et al., 2023] Learning rate1× 10 −4 Global batch size768 Weight decay0.1 Gradient clipping1.0 LR scheduleStep decay at 80% and 90% of training LR decay factors0.316× and 0.1× Warmup steps2000 Precisionbfloat16 (mixed precision) FrameworkPyTorch DDP + torch.compile Compile mode default HardwareNVIDIA B200 192GB GPUs (×8) Puzzle embedding optimizer SignSGD (lr=10 −2 ) Puzzle embedding dim512 (× 16 tokens) 49 Table 13: Required Gauss-Seidel iterations for ε = 0.01 convergence. n234568 ρ = 1− 1/n 0.500 0.667 0.750 0.800 0.833 0.875 Iterations k71216212637 Table 14: Backward pass complexity: standard vs. implicit Sinkhorn. Standard SinkhornImplicit Sinkhorn (Ours) Autograd nodesO(T · n 2 ) (∼ 128K)O(1) (constant) Backward kernels ∼ 128K microsecond-scale ∼ 20 millisecond-scale DDP overlapPoor (55% stalls)Excellent Memory (saved tensors)O(T · n 2 )O(n 2 ) (just P) Gradient accuracyExactε≤ 0.01 (controlled) Table 15: Properties of doubly-stochastic (Sinkhorn) vs. orthonormal (Cayley) mixing. PropertyDoubly-Stochastic (B n )Orthonormal (O(n)) Norm behavior Contractive (∥Hx∥≤∥x∥) Preserving (∥Hx∥ =∥x∥) EntriesNon-negativeUnconstrained sign Row/col sumsBoth equal 1Not constrained det(H)∈ [0, 1]±1 Convex hullPermutation matricesNot a convex set Gradient flowMay attenuatePreserved Backward cost O(T ) or implicitStandard autograd Table 16: Parameter count for residual mixing matrix (n = 4). VariantParameters n = 4 Sinkhorn (full n× n)n 2 16 Cayley (full n× n)n 2 16 Grassmannian (n× p, p = n/2)np8 50