PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Wuyang Zhou, Yuxuan Gu, Giorgos Iacovides, Danilo Mandic

arXiv:2601.21579v2cs.CLcs.LG

TL;DR

HC-related residual mixing can be unstable and difficult to scale because existing manifold-constrained variants either approximate double stochasticity or incur large parameter costs. KromHC uses Kronecker products of smaller doubly stochastic matrices to address both issues, and experiments report consistent improvements over SOTA mHC variants with fewer trainable parameters.

  • Problem

    Existing HC variants face training instability and restricted scalability: mHC may fail to achieve exact double stochasticity, while mHC and mHC-lite have high parameter complexity as n increases.

  • Method

    KromHC parameterizes mHC residual matrices as Kronecker products of smaller doubly stochastic matrices, enforcing exact double stochasticity through tensorized residual mixing.

  • Results

    KromHC yields zero MAE in the residual-matrix stability analysis, whereas standard mHC shows an MAE of around 0.05.

  • Takeaways & Limitations

    KromHC resolves the reported conflict between exact double stochasticity and parameter efficiency while achieving consistent improvements over SOTA mHC variants without customized kernels.

  • Takeaways & Limitations

    mHC still does not guarantee double stochasticity and requires customized kernels for accelerating its 20-iteration Sinkhorn-Knopp projection.

Abstract

from arXiv · show

The success of Hyper-Connections (HC) in neural networks (NN) has also highlighted issues related to 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 exactly doubly stochastic residual matrices; 2) mHC incurs a prohibitive $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, $O \left( nC \cdot n! \right)$. To address both challenges, we propose KromHC, which uses the Kronecker products of smaller doubly stochastic matrices to parametrize the residual matrix in 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 only $O(n^2C)$. Experiments show that KromHC matches or even outperforms other state-of-the-art (SOTA) mHC variants, while requiring significantly fewer trainable parameters. The code is at https://github.com/wz1119/KromHC.

1. Introduction

Hyper-Connections improve feature propagation by widening residual streams, but unconstrained residual mixing can destabilize training and existing manifold-constrained variants trade exact double stochasticity against scalability. KromHC uses Kronecker-structured residual matrices to provide exact double stochasticity with improved parameter efficiency.

  • Motivation: HC expands residual stream width and introduces learnable mixing across streams without increasing FLOPs, enabling more expressive feature propagation.A single HC layer mixes residual streams, aggregates features, and maps outputs back onto the expanded streams.
  • Motivation: Unconstrained HC residual matrices can lose the identity-mapping property when stacked across layers, causing numerical instability in large neural networks.The issue arises from products of residual matrices across multiple layers.
  • Prior approaches: mHC projects residual matrices onto the Birkhoff polytope with Sinkhorn-Knopp, but finite iterations can fail to produce exactly doubly stochastic matrices.The resulting error can accumulate across layers and undermine training stability.
  • Prior approaches: mHC-lite guarantees exact double stochasticity through convex combinations of permutation matrices, but its n! matrix storage causes factorial parameter growth.Its parameter complexity prevents scaling the residual-stream width n effectively.
  • Results: KromHC achieves zero MAE in the residual-matrix stability analysis, while standard mHC exhibits an MAE of around 0.05.The comparison uses an LLM with D = 12 transformer blocks and L = 24 hyper-connection layers.
  • KromHC: KromHC parameterizes residual matrices as Kronecker products of smaller doubly stochastic matrices, guaranteeing exact double stochasticity while improving parameter efficiency.The approach frames residual mixing as a Tucker-structured tensor network and targets the conflict between exact constraints and scalable parameter counts.
  • Results: Extensive LLM-pretraining experiments report consistent improvements over SOTA mHC variants without customized kernels.The paper presents this as evidence of effectiveness and scalability.

2. Related Works

Prior work expands residual-stream width for more expressive feature propagation, then constrains residual mixing to address instability and preserve identity mappings. mHC-lite achieves exact double stochasticity but incurs factorial parameter growth, while tensor networks motivate parameter-efficient factorization.

  • Hyper-Connections expand residual-stream width n without extra FLOPs and introduce learnable mixing across streams for more expressive feature propagation.
  • mHC projects residual matrices onto the Birkhoff polytope, whose doubly stochastic matrices restore the identity mapping across layers.
  • mHC obtains residual mappings from a flattened hidden matrix and applies learned projections, reshaping, sigmoid gating, and 20 Sinkhorn-Knopp iterations.
  • mHC-lite guarantees exact double stochasticity through convex combinations of permutation matrices, but its parameter complexity grows as O(nC · n!).
  • Tensor networks factor higher-order tensors into lower-order cores and factors, providing expressive, parameter-efficient representations that alleviate the curse of dimensionality.

3. Notation and Preliminaries

The paper introduces tensor notation, tensorization, unfolding, and Tucker decomposition for representing higher-order arrays. Tucker decomposition expresses a tensor using a core tensor and factor matrices with specified Tucker ranks.

  • An order-K tensor is a multidimensional array with K modes; vectors and matrices are order-1 and order-2 tensors, respectively.
  • Tensorization reshapes vectors or matrices into higher-order tensors, while unfolding reverses this process through matricization or vectorization.
  • Tucker decomposition parametrizes an order-K tensor as a core tensor multiplied along each mode by K factor matrices.
  • The Tucker ranks are represented by the vector [r1, r2, . . . , rK], and vec(·) converts a tensor into a vector.
  • The mode-k product is defined element-wise by summing over the corresponding core-tensor index and multiplying by entries of Uk.

4. Methodology

KromHC tensorizes the residual stream and parameterizes its residual mixing matrix as a Kronecker product of smaller doubly stochastic factors. This construction guarantees exact double stochasticity while reducing parameter complexity relative to mHC and mHC-lite.

  • Tensorizing the Residual Stream: KromHC represents residual mixing as a Tucker decomposition tensor network whose core tensor is the tensorized residual stream.The residual stream is expanded to n streams, tensorized across K modes, and mixed along those modes.
  • Kronecker-Product Residual Matrices: The KromHC residual matrix is the Kronecker product of factor matrices U_k, each constrained to be doubly stochastic.The method applies residual mixing along each tensor mode using learned factor matrices.
  • Parametrization of KromHC: KromHC parameterizes each smaller factor as a convex combination of permutation matrices, requiring only two 2 × 2 permutation matrices when n is a power of 2 with i_k = 2.Only two scalar coefficients are learned to represent each factor on the Birkhoff polytope.
  • Kronecker-Product Residual Matrices: The Kronecker product of any finite collection of doubly stochastic matrices is also doubly stochastic, so KromHC residual matrices are exact by construction.This adds a Kronecker structural constraint alongside the manifold constraint used in mHC.
  • Kronecker-Product Residual Matrices: KromHC preserves norm preservation and compositional closure because its residual matrices remain exactly doubly stochastic.The residual matrix spectral norm is bounded by 1.

5. Experiments

Experiments evaluate KromHC on language modeling, commonsense reasoning, scaling, and gradient stability. KromHC matches training and validation metrics of other manifold-constrained variants while improving downstream performance and using fewer additional parameters.

  • Experimental Setup: Experiments used approximately 60M- and 186M-parameter models trained on FineWeb-Edu with 6 or 12 transformer blocks.Models used either 4 or 8 NVIDIA RTX PRO 6000 GPUs depending on residual-stream width.
  • Training and Validation Set Metrics: KromHC significantly outperformed SOTA mHC variants on CORE while achieving on-par training loss and validation BPB with mHC and mHC-lite.CORE is a centered accuracy over 22 downstream evaluation tasks.
  • Downstream Task Performances: KromHC achieved the highest commonsense-reasoning averages at 6 blocks (41.1%) and 12 blocks (47.7%).It surpassed second-best scores by up to 2% on COPA and 6.4% on BoolQ.
  • Downstream Task Performances: KromHC achieved the best language-modeling averages at 6 blocks (17.3%) and 12 blocks (24.0%).The reported results cover language-modeling performance during LLM pretraining.
  • Scaling the Width of Residual Stream in KromHC: KromHC benefits from residual-stream widths n ∈ {4, 8, 16}, with larger n producing larger training-loss gaps and consistently improved validation BPB.Additional learnable parameters at each width are recorded in Table 6.
  • Gradient Norm: KromHC consistently achieved the lowest gradient norm among the compared manifold-constrained variants.mHC-lite and KromHC had lower gradient norms than mHC, consistent with their exactly doubly stochastic residual matrices.

6. Conclusion

KromHC is presented as a parameter-efficient framework using Kronecker-product residual matrices to guarantee exact double stochasticity and address scalability limitations. Its main stated limitation concerns residual-stream widths that are large prime numbers.

  • Conclusion: KromHC uses Kronecker-product residual matrices to guarantee exact double stochasticity with parameter-efficient manifold-constrained hyper-connections.The conclusion reports extensive effectiveness experiments in LLM pretraining.
  • Conclusion: Future work aims to apply KromHC to domains such as computer vision.
  • Limitations: KromHC may encounter parameter issues when the residual-stream width n is a large prime number.The paper suggests using a larger n that is a power of 2 or 3 or has only small prime factors.

Impact Statement

The paper frames KromHC as addressing training stability and scalability in manifold-constrained hyper-connections. Its broader impact statement emphasizes more reliable training with fewer parameters for future AI deployment.

  • Impact Statement: KromHC is described as resolving scalability and stability issues in manifold-constrained hyper-connections.
  • Kronecker Products: Kronecker products represent structured linear operators by replacing matrix entries with scaled copies of another matrix.For matrices A ∈ R^m×n and B ∈ R^p×q, A ⊗ B has size (mp) × (nq).
  • Kronecker Products: The Kronecker product extends recursively to multiple matrices, producing a structured operator whose size is determined by the factor dimensions.

B. Proof for Theorem B.1

Theorem B.1 establishes that Kronecker products preserve double stochasticity. The proof verifies nonnegativity and unit row and column sums, then extends the result by induction to finitely many factors.

  • Theorem B.1: The theorem states that the Kronecker product of doubly stochastic matrices is doubly stochastic.The result is formulated for matrices in the corresponding Birkhoff polytopes.
  • Proof: Nonnegative entries remain nonnegative under the Kronecker product because each resulting entry is a product of nonnegative factors.
  • Proof: The proof uses a Kronecker-product identity with all-one vectors to show that every row sum equals one.
  • Proof: The same identity shows that every column sum equals one.
  • Proof: Combining nonnegativity with unit row and column sums proves double stochasticity and extends to finitely many factors by induction.

C. Details of the CORE Tasks

The experiments use CORE as an averaged accuracy metric across 22 diverse tasks, while the paper also describes HC, mHC-lite, and transformer-block training configurations.

  • C. Details of the CORE Tasks: CORE averages linearly scaled accuracy across 22 tasks, with 0 representing random guessing and 1 perfect accuracy.The averaging prevents any single benchmark from dominating the final score.
  • C. Details of the CORE Tasks: CORE covers logical reasoning, factual recall, algorithmic thinking, commonsense inference, and language understanding.Its benchmarks include reasoning, knowledge, BIG-Bench, question answering, commonsense, and other widely used tasks.
  • C. Details of the CORE Tasks: HC can be unstable to train and lacks cross-stream input-dependent residual mixing available in mHC.The paper therefore describes an HC parameterization that remains similar to mHC without using the Sinkhorn-Knopp algorithm.
  • C. Details of the CORE Tasks: mHC-lite constructs dynamic residual mappings from flattened layer inputs using learnable weight matrices, scalars, and permutation matrices.The formulation uses permutation matrices P_m ∈ R^n×n and learnable parameters in each layer.
  • C. Details of the CORE Tasks: Each transformer block has attention and FFN residual connections, and the experiments replace standard or scalar residual variants with mHC variants.Nanochat’s learnable per-layer scalars are initialized to 1 and 0, respectively.
  • C. Details of the CORE Tasks: The experiments use Muon for main-branch attention and MLP weights, and AdamW for hyper-connections, embeddings, and the LM head.The paper reports separate scale-specific hyperparameters for D = 6 and D = 12 transformer blocks.

I. Grad Norm

Figure 8 presents raw gradient-norm dynamics for mHC, mHC-lite, and KromHC during training at D = 12.

  • I. Grad Norm: Figure 8 tracks raw gradient norm across 7000 training steps for mHC, mHC-lite, and KromHC at D = 12.The supplied passage identifies the plotted quantity, methods, training duration, and model depth.

J. System Metrics

The system-metrics evaluation compares manifold-constrained hyper-connections against mHC, while an ablation examines alternative KromHC factorizations.

  • J. System Metrics: KromHC is more computationally efficient than mHC and mHC-lite, with higher throughput, lower wall-clock time per step, and lower GPU memory use.The comparison uses n = 4, 12 transformer blocks, pure PyTorch implementations, and 8 RTX 6000 GPUs.
  • J. System Metrics: Table 9 reports relative changes in token throughput, wall-clock time per step, and GPU memory footprint compared with mHC.A dash indicates no difference.
  • J. System Metrics: At n = 8, the 2 × 2 × 2 KromHC factorization slightly outperforms 4 × 2, with train loss 0.002 lower and validation BPB 0.003 lower.The original mHC-lite configuration at n = 8 is described as unrealistic because of model-size explosion.
Loading 2601.21579v2…