JPmHC Dynamical Isometry via Orthogonal Hyper-Connections
Organizations: LLM Suite Team, JP Morgan Chase & Co.
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- 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
Figures & tables
| Variant | Manifold | Key param | FLOPs/module | Norm | Constraint |
|---|---|---|---|---|---|
| Sinkhorn | Birkhoff polytope | , | 576 | RMSNorm | sigmoid |
| Cayley | Stiefel | , | 256 | LayerNorm | softmax |
| Grassmann | , | — | LayerNorm | softmax |
| Metric | Sinkhorn | Cayley | Grassmann ∗ | Cayley/Sink |
|---|---|---|---|---|
| Training Steps | 510,951 | 515,902 | 110,756 | — |
| Eval Accuracy | 86.5% | 86.9% | 82.2% | |
| Exact Accuracy (greedy) | 27.9% | 31.4% | 12.8% | |
| Pass@1 | 36.5% | 40.5% | 27.5% | |
| Pass@2 | 41.7% | 45.4% | 31.2% | |
| Pass@5 | 45.5% | 50.5% | 35.9% |
| Step | Exact Acc. | Pass@1 | Pass@2 | Pass@10 | Pass@100 | Pass@1000 |
|---|---|---|---|---|---|---|
| 68,951 | 10.6% | — | — | — | — | — |
| 81,903 | 12.7% | — | — | — | — | — |
| 94,854 | 15.7% | — | — | — | — | — |
| 107,805 | 17.6% | — | — | — | — | — |
| 161,952 | 23.4% | 35.4% | 40.1% | 47.4% | 52.7% | 54.1% |
| 174,904 | 24.5% | 35.6% | 40.3% | 47.7% | 53.9% | 57.1% |
| Step | Exact Acc. | Pass@1 | Pass@2 | Pass@10 | Pass@100 | Pass@1000 |
|---|---|---|---|---|---|---|
| 52,951 | 1.4% | 5.2% | 7.5% | 12.1% | 14.9% | 16.1% |
| 65,902 | 3.1% | 9.0% | 11.6% | 16.2% | 21.4% | 22.6% |
| 78,854 | 5.0% | 11.6% | 16.2% | 21.1% | 26.8% | 29.5% |
| 99,951 | 8.5% | 19.1% | 22.8% | 29.4% | 36.6% | 39.6% |
| 112,902 | 9.9% | 22.2% | 26.2% | 33.1% | 40.6% | 43.9% |
| 125,853 | 11.2% | 23.7% | 28.0% | 35.5% | 42.8% | 45.5% |
| Step | Exact Acc. | Pass@1 | Pass@2 | Pass@10 | Pass@100 | Pass@1000 |
|---|---|---|---|---|---|---|
| 12,951 | 0.2% | 0.6% | 1.3% | 3.2% | 3.9% | 4.6% |
| 28,951 | 0.8% | 3.6% | 5.5% | 9.9% | 12.4% | 14.0% |
| 41,902 | 2.6% | 8.9% | 12.1% | 16.0% | 18.6% | 21.1% |
| 58,951 | 5.5% | 15.7% | 18.6% | 24.0% | 29.4% | 31.5% |
| 71,902 | 7.4% | 20.0% | 24.1% | 30.4% | 35.7% | 38.4% |
| 84,854 | 9.6% | 22.2% | 27.3% | 34.7% | 40.4% | 43.8% |
| Variant | Forward FLOPs | Backward FLOPs | Total per module |
|---|---|---|---|
| Sinkhorn (implicit) | 576 | ||
| Cayley | 256 | ||
| Grassmann | 72 |
| Feature | HC / mHC | JPmHC (Ours) |
|---|---|---|
| Mixing parameterization | Learned / Sinkhorn | Stiefel, Grassmann, Birkhoff |
| Manifold constraint | Birkhoff polytope only | Stiefel, Grassmann, Birkhoff |
| Implicit differentiation | — | Sinkhorn backward |
| Riemannian optimization | — | Grassmann (Cayley ADAM) |
| Spectral analysis | — | Generating set selection |
| CUDA graph compatible | Not addressed | All variants |
Appendix figures & tables9 assets
Supplementary material from the paper’s appendix.
Appendix
| Matrix size | Feasibility | |||
|---|---|---|---|---|
| 3 | 6 | 15 | s | |
| 4 | 24 | 10,626 | s | |
| 5 | 120 | 8,214,570 | hour | |
| 6 | 720 | Intractable |
| Variant | Total params | Notes | |
|---|---|---|---|
| Sinkhorn | Fused | ||
| Cayley | Fused + LN | ||
| Grassmann | |||
| Perm Mix | + perm indices |
| Iterations | Relative residual | Gradient error (%) |
|---|---|---|
| 1 | 43.2 | |
| 2 | 28.1 | |
| 4 | 11.5 | |
| 8 | 2.8 | |
| 16 | 0.08 | |
| 32 |
| Parameter | Value |
|---|---|
| Total parameters | 7M |
| Hidden dimension | 512 |
| Number of streams | 4 |
| Effective hidden dim | 2048 |
| Unique transformer layers (weight-tied) | 2 |
| Recursive cycles per layer | 6 |
| Hyperparameter | Value |
|---|---|
| Optimizer | AdamAtan2 [ Kunstner et al., 2023 ] |
| Learning rate | |
| Global batch size | 768 |
| Weight decay | 0.1 |
| Gradient clipping | 1.0 |
| LR schedule | Step decay at 80% and 90% of training |
| 2 | 3 | 4 | 5 | 6 | 8 | |
|---|---|---|---|---|---|---|
| 0.500 | 0.667 | 0.750 | 0.800 | 0.833 | 0.875 | |
| Iterations | 7 | 12 | 16 | 21 | 26 | 37 |
| Standard Sinkhorn | Implicit Sinkhorn (Ours) | |
| Autograd nodes | ( K) | (constant) |
| Backward kernels | K microsecond-scale | millisecond-scale |
| DDP overlap | Poor (55% stalls) | Excellent |
| Memory (saved tensors) | (just ) | |
| Gradient accuracy | Exact | (controlled) |
| Property | Doubly-Stochastic ( ) | Orthonormal ( ) |
|---|---|---|
| Norm behavior | Contractive ( ) | Preserving ( ) |
| Entries | Non-negative | Unconstrained sign |
| Row/col sums | Both equal | Not constrained |
| Convex hull | Permutation matrices | Not a convex set |
| Gradient flow | May attenuate | Preserved |
| Variant | Parameters | |
|---|---|---|
| Sinkhorn (full ) | 16 | |
| Cayley (full ) | 16 | |
| Grassmannian ( , ) | 8 |
Hyper-Connections
Hyper-connections expand neural-network residual paths into multiple streams with learned routing and mixing to improve capacity, stability, and efficiency.
See all 18 papers