Rethinking Dense Linear Transformations: Stagewise Pairwise Mixing (SPM) for Near-Linear Training in Neural Networks

TL;DR

Dense linear layers dominate compute and parameter budgets due to their $O(n^2)$ complexity. Stagewise Pairwise Mixers (SPM) replace dense maps with a sequence of sparse 2×2, pairwise mixing stages, achieving near-linear time and parameter growth $O(nL)$ while preserving exact forward and backward computations. The approach introduces a compositional inductive bias, improves gradient stability, and offers theoretical justifications for better generalization on structured tasks. Empirical results on synthetic compositional teachers, AG News with hashed features, and Shakespeare language modeling demonstrate substantial compute savings and competitive or improved accuracy, especially at large widths, with a flexible framework that accommodates GRUs and attention mechanisms.

Abstract

Dense linear layers are a dominant source of computational and parametric cost in modern machine learning models, despite their quadratic complexity and often being misaligned with the compositional structure of learned representations. We introduce Stagewise Pairwise Mixers (SPM), a structured linear operator that replaces dense matrices with a composition of sparse pairwise-mixing stages. An SPM layer implements a global linear transformation in $O(nL)$ time with $O(nL)$ parameters, where $L$ is typically constant or $log_2n$, and admits exact closed-form forward and backward computations. SPM is designed as a drop-in replacement for dense linear layers in feedforward networks, recurrent architectures, attention mechanisms, etc. We derive complete forward and backward expressions for two parameterizations: an orthogonal norm-preserving rotation-based variant and a fully general $2 \times 2$ mixing variant. Beyond computational savings, the stagewise structure of SPM induces an explicit compositional inductive bias that constrains model capacity and improves generalization when aligned with task structure. We present proof-of-concept experiments demonstrating substantial reductions in wall-clock cost and improved accuracy on structured learning problems, while retaining competitive performance on real-world benchmarks.