Track papers across labs, topics, and trends.
Search papers, labs, and topics across Lattice.
The paper introduces go-$m$HC, a novel and efficient parameterization of doubly stochastic matrices using generalized orthostochastic matrices, enabling learned mixing across residual streams in neural networks. This parameterization scales as $\mathcal{O}(d^3)$ and allows for continuous interpolation between a computationally efficient boundary and the fully expressive Birkhoff polytope via a single hyperparameter. Experiments on synthetic tasks and a 30M parameter GPT-style language model demonstrate that go-$m$HC achieves better performance and faster convergence compared to Kronecker-factorized methods, while more completely filling the Birkhoff polytope.
Forget Kronecker factorization: go-$m$HC offers an exact, $\mathcal{O}(d^3)$ parameterization of doubly stochastic matrices that unlocks a new dimension of model capacity by efficiently scaling the number of residual streams.
Doubly stochastic matrices enable learned mixing across residual streams, but parameterizing the set of doubly stochastic matrices (the Birkhoff polytope) exactly and efficiently remains an open challenge. Existing exact methods scale factorially with the number of streams ($d$), while Kronecker-factorized approaches are efficient but expressivity-limited. We introduce a novel exact parameterization grounded in the theory of generalized orthostochastic matrices, which scales as $\mathcal{O}(d^3)$ and exposes a single hyperparameter $s$ which continuously interpolates between a computationally efficient boundary and the fully expressive Birkhoff polytope. Building on Manifold-Constrained Hyper-Connections ($m$HC), a framework for learned dynamic layer connectivity, we instantiate this parameterization in go-$m$HC. Our method composes naturally with Kronecker-factorized methods, substantially recovering expressivity at similar FLOP costs. Spectral analysis indicates that go-$m$HC fills the Birkhoff polytope far more completely than Kronecker-factorized baselines. On synthetic stream-mixing tasks, go-$m$HC achieves the minimum theoretical loss while converging up to $10\times$ faster. We validate our approach in a 30M parameter GPT-style language model. The expressivity, efficiency, and exactness of go-$m$HC offer a practical avenue for scaling $d$ as a new dimension of model capacity.