mHC: Manifold-Constrained Hyper-Connections

Overview

mHC (Manifold-Constrained Hyper-Connections) stabilizes deep network training by constraining residual mixing matrices to be doubly stochastic. It provides:

• Stable Training: Lower gradient norm variance via doubly stochastic constraints
• Multiple Streams: Hyper-Connections with learnable mixing across residual streams
• Sinkhorn Projection: Log-space Sinkhorn-Knopp algorithm for doubly stochastic projection
• GPT Integration: Pattern for wrapping attention and MLP layers

Two components:

• HyperConnections Module: Core PyTorch module with H_res, H_pre, H_post matrices
• Sinkhorn-Knopp: Log-space projection to doubly stochastic manifold

Quick Reference

Topic	Reference
Core Concepts & Math	Core Concepts
Sinkhorn Algorithm	Sinkhorn-Knopp
HyperConnections Module	Module Implementation
GPT Integration	GPT Integration
Common Pitfalls	Pitfalls

Installation

# Required packages
pip install torch einops numpy

Minimal Example

import torch
import torch.nn as nn
from einops import rearrange, einsum

def sinkhorn_knopp(logits, num_iters=20, tau=0.05):
    log_alpha = logits / tau
    for _ in range(num_iters):
        log_alpha = log_alpha - torch.logsumexp(log_alpha, dim=-1, keepdim=True)
        log_alpha = log_alpha - torch.logsumexp(log_alpha, dim=-2, keepdim=True)
    return torch.exp(log_alpha)

class HyperConnections(nn.Module):
    def __init__(self, num_streams, dim, branch=None, layer_idx=0):
        super().__init__()
        self.num_streams = num_streams
        self.branch = branch