Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

KromHC: Manifold-Constrained Hyper-Connections with Kronecker-Product Residual Matrices

Wuyang Zhou, Yuxuan Gu, Giorgos Iacovides, Danilo Mandic

TL;DR

KromHC addresses instability and non-scalability in manifold-constrained hyper-connections by replacing exact projection via iterative Sinkhorn steps with a Kronecker-product structured residual. By tensorizing the residual stream and enforcing a Tucker/Kronecker construction, the method guarantees exact doubly stochastic residual matrices while achieving $ ext{O}(n^2C)$ parameter complexity, dramatically reducing the parameter burden compared to prior mHC variants. Empirical results on LLM pretraining show KromHC matches or surpasses state-of-the-art mHC methods in training stability and downstream tasks (commonsense reasoning and language modeling) with fewer trainable parameters, and it scales effectively with residual width $n$. The approach leverages tensor-network principles to connect hyper-connection design with efficient, structure-preserving linear operators, offering practical gains for large-scale neural architectures.

Abstract

The success of Hyper-Connections (HC) in neural networks (NN) has also highlighted issues related to its training instability and restricted scalability. The Manifold-Constrained Hyper-Connections (mHC) mitigate these challenges by projecting the residual connection space onto a Birkhoff polytope, however, it faces two issues: 1) its iterative Sinkhorn-Knopp (SK) algorithm does not always yield exact doubly stochastic residual matrices; 2) mHC incurs a prohibitive $\mathcal{O}(n^3C)$ parameter complexity with $n$ as the width of the residual stream and $C$ as the feature dimension. The recently proposed mHC-lite reparametrizes the residual matrix via the Birkhoff-von-Neumann theorem to guarantee double stochasticity, but also faces a factorial explosion in its parameter complexity, $\mathcal{O} \left( nC \cdot n! \right)$. To address both challenges, we propose \textbf{KromHC}, which uses the \underline{Kro}necker products of smaller doubly stochastic matrices to parametrize the residual matrix in \underline{mHC}. By enforcing manifold constraints across the factor residual matrices along each mode of the tensorized residual stream, KromHC guarantees exact double stochasticity of the residual matrices while reducing parameter complexity to $\mathcal{O}(n^2C)$. Comprehensive experiments demonstrate that KromHC matches or even outperforms state-of-the-art (SOTA) mHC variants, while requiring significantly fewer trainable parameters. The code is available at \texttt{https://github.com/wz1119/KromHC}.

KromHC: Manifold-Constrained Hyper-Connections with Kronecker-Product Residual Matrices

TL;DR

KromHC addresses instability and non-scalability in manifold-constrained hyper-connections by replacing exact projection via iterative Sinkhorn steps with a Kronecker-product structured residual. By tensorizing the residual stream and enforcing a Tucker/Kronecker construction, the method guarantees exact doubly stochastic residual matrices while achieving parameter complexity, dramatically reducing the parameter burden compared to prior mHC variants. Empirical results on LLM pretraining show KromHC matches or surpasses state-of-the-art mHC methods in training stability and downstream tasks (commonsense reasoning and language modeling) with fewer trainable parameters, and it scales effectively with residual width . The approach leverages tensor-network principles to connect hyper-connection design with efficient, structure-preserving linear operators, offering practical gains for large-scale neural architectures.

Abstract

The success of Hyper-Connections (HC) in neural networks (NN) has also highlighted issues related to its training instability and restricted scalability. The Manifold-Constrained Hyper-Connections (mHC) mitigate these challenges by projecting the residual connection space onto a Birkhoff polytope, however, it faces two issues: 1) its iterative Sinkhorn-Knopp (SK) algorithm does not always yield exact doubly stochastic residual matrices; 2) mHC incurs a prohibitive parameter complexity with as the width of the residual stream and as the feature dimension. The recently proposed mHC-lite reparametrizes the residual matrix via the Birkhoff-von-Neumann theorem to guarantee double stochasticity, but also faces a factorial explosion in its parameter complexity, . To address both challenges, we propose \textbf{KromHC}, which uses the \underline{Kro}necker products of smaller doubly stochastic matrices to parametrize the residual matrix in \underline{mHC}. By enforcing manifold constraints across the factor residual matrices along each mode of the tensorized residual stream, KromHC guarantees exact double stochasticity of the residual matrices while reducing parameter complexity to . Comprehensive experiments demonstrate that KromHC matches or even outperforms state-of-the-art (SOTA) mHC variants, while requiring significantly fewer trainable parameters. The code is available at \texttt{https://github.com/wz1119/KromHC}.
Paper Structure (32 sections, 3 theorems, 27 equations, 8 figures, 8 tables)

This paper contains 32 sections, 3 theorems, 27 equations, 8 figures, 8 tables.

Table of Contents

  1. Introduction
  2. Related Works
  3. Macro-design of Neural Architecture.
  4. Tensor Networks.
  5. Notation and Preliminaries
  6. Tucker Decomposition Tensor Network.
  7. Methodology
  8. Tensorizing the Residual Stream
  9. Kronecker-Product Residual Matrices
  10. Parametrization of KromHC
  11. Parameter Complexity Analysis.
  12. Experiments
  13. Initialization
  14. Training and Validation Set Metrics
  15. Downstream Task Performances
  16. ...and 17 more sections

Key Result

Theorem 4.1

(Birkhoff-von-Neumann Theorembirkhoff-von-neumann) For any $n \times n$ doubly stochastic matrix, $\mathbf{X}$, there exists a finite collection of permutation matrices $\{\mathbf{P}_k \in \mathbb{R}^{n\times n}\}_{k=1}^{n!}$ and a coefficient vector $\mathbf{a} = (a_1,\ldots,a_{n!}) \in \mathbb{R}^

Figures (8)

  • Figure 1: Illustration of variants of manifold-constrained hyper-connections with a residual stream width $n=8$. (a) mHC: utilizes iterative Sinkhorn-Knopp (SK) algorithm to approximate a doubly stochastic residual matrix; (b) mHC-lite: builds the residual matrix as convex combinations of $n!$ permutation matrices, but becomes infeasible for a large $n$; (c) KromHC (Ours): constructs the residual matrix as the Kronecker products of smaller (e.g., $2 \times 2$) doubly stochastic matrices, thus guaranteeing double stochasticity while remaining parameter efficient.
  • Figure 2: Numerical stability analysis of the products of residual matrices. The plot compares the Mean Absolute Error (MAE) between the column sum of $\prod_{i=0}^{L-1} \mathbf{H}^{\mathrm{res}}_{L-i}$ and $1$ in an LLM with $D=12$ transformer blocks and $L=24$ layers of HC. The standard mHC architecture exhibits a MAE of around $0.05$, indicating potential training instabilities. The mHC-lite and KromHC have exact doubly stochastic matrices, thus yielding zero MAE.
  • Figure 3: The number of learnable parameters against the number of residual streams, $n$, per hyper-connection in mHC, mHC-lite, and KromHC. We assume the feature dimension, $C$, to be 512. Also, $n$ is factored into $\prod_{m=1}^{\log_2(n)} 2$, i.e., $i_1 = i_2=\cdots=i_K=2$.
  • Figure 4: Training loss and validation BPB gaps of KromHC at different widths of the residual stream, $n$, compared to $n=4$. Exponential Moving Average (EMA) is applied to the raw loss before the calculation of the loss gap.
  • Figure 5: Zoomed-in view of gradient norms from 5000 to 7000 steps during training. Trajectories are smoothed using EMA, with shaded regions indicating the EMA variance.
  • ...and 3 more figures

Theorems & Definitions (8)

  • Theorem 4.1
  • Theorem 4.2
  • proof
  • Remark 4.3
  • Remark 4.4
  • Remark 4.5
  • Theorem 2.1
  • proof