Source-linked AI summary
A Mathematical Theory of Reusable Neural Bases for Network Compression
Binshuai Wang, Peng Wei
TL;DR
Large-model memory requirements motivate a parameter-efficient architecture that retains stable training. LRNBA constructs blocks from reusable neural bases and position-specific coefficients, and experiments report reduced parameters with competitive or improved training behavior. The approach supports wider and deeper networks under limited memory, though dynamic block-weight computation increases training time and energy.
Problem
Large AI models impose substantial memory, computational, and hardware requirements, motivating parameter reuse that avoids classical RNN optimization difficulties.
Method
LRNBA represents network blocks as linear combinations of shared reusable neural bases, extending reuse to attention, projection, and sub-neural components.
Results
Experiments on FFN-based ResNets and decoder-only Transformers report significantly reduced parameter counts while maintaining stable training and competitive performance.
Takeaways & Limitations
LRNBA enables substantially wider and deeper networks under limited memory budgets, with further basis, width, or depth expansion improving convergence and final accuracy in several settings.
Takeaways & Limitations
Dynamic computation and updating of block-specific weights increases training time and energy consumption, although caching effective weights can eliminate this cost during inference.
Abstract
from arXiv · showhide
As large AI models become increasingly prevalent across a wide range of applications, memory cost has become a critical bottleneck in both training and inference. To mitigate this issue, we introduce the Linear Reusable Neural Bases Architecture (LRNBA), a novel framework aimed at improving parameter efficiency and reducing memory cost. Inspired by recurrent neural network (RNN) designs, the core idea of our approach is to represent each network block as a linear combination of a shared set of neural bases, thereby enjoying highly network compression rate while maintaining stable training. The proposed architecture allows for the construction of significantly wider and deeper networks under the same parameter budget. Extensive experiments demonstrate that our model achieves comparable or even faster convergence and lower loss than classical architectures, while maintaining stable training dynamics.
1 INTRODUCTION
Large models improve capability but impose severe memory and hardware demands. LRNBA addresses this bottleneck by reusing neural bases across blocks, reducing parameters while preserving stable training and competitive performance.
- Large AI models can require hundreds or thousands of GPUs, making memory and hardware requirements a major development and deployment bottleneck.
- Parameter reuse, as in RNNs, can reduce distinct parameters and memory cost, but classical RNNs may suffer vanishing or exploding gradients.
- LRNBA represents each network block as a linear combination of reusable neural bases with different coefficients across positions.
- Compared with classical architectures, LRNBA reports comparable or faster convergence, lower final loss, stable training, and substantially reduced parameter count.
- LRNBA extends neural bases to attention and projection components and introduces sub-neural reuse to improve parameter efficiency.
- Unlike classical models, LRNBA's parameter count is almost decoupled from depth and depends primarily on width, enabling wider and deeper networks under limited memory.
2 RELATED WORK
Prior parameter-sharing methods include recurrent, factorized, structured, and Transformer-based approaches. LRNBA differs by reusing nonlinear features at the neuron level across depths and components.
- RNNs reuse transformations across time, while LSTMs and GRUs use gating to mitigate vanishing and exploding gradients.
- Low-rank factorizations, auxiliary parameter-generating networks, tensorized networks, and structured matrices reduce storage or parameter costs through alternative parameterizations.
- Transformer methods share FFN, attention, transition, or looped blocks across depth to reduce parameters or support iterative computation.
- LRNBA reuses nonlinear features at the neuron level across depths, including attention and FFN components, allowing more flexible combinations than identical block reuse.
3 BACKGROUND
Deep composition improves representational ability but increases computational and memory costs. Residual blocks address stable deep training and can be interpreted as combinations of nonlinear functions.
- Sequentially composed deep networks can represent complex functions more efficiently than shallow networks and often have better approximation ability.
- Deep architectures require more computation and memory because parameters and intermediate activations must be retained across sequential components.
- ResNet shortcut connections pass input information directly to deeper layers, alleviating vanishing gradients and enabling stable training of very deep networks.
- A residual feedforward block uses an input affine transformation, nonlinear activation, and output linear transformation, with x as the shortcut and R(x) as the residual component.
- The residual component can be viewed as a linear combination of shared nonlinear functions whose input directions and offsets come from (S, b), while P determines output coefficients.
4 METHODOLOGY
LRNBA builds blocks from reusable neural bases and position-specific coefficients, extending this construction to broader components and finer-grained reuse. Its training procedure rescales reused gradients, while multi-sensor neurons further reduce parameters.
- 4 METHODOLOGY: LRNBA reuses a shared collection of neural bases while assigning different coefficients to construct blocks at different positions.
- 4.1 THE NEURAL BASES: A neural basis is an atomic nonlinear function associated with sensor parameters (s_j, b_j) and response parameter p_j.
- 4.1 THE NEURAL BASES: Each neural basis defines a vector field, and a residual feedforward block is the superposition of many such fields.
- 4.1 THE NEURAL BASES: The framework extends neural bases to attention and projection blocks, treating different network blocks as linear combinations of corresponding bases, including in Transformers.
- 4.2 LINEAR REUSABLE NEURAL BASES ARCHITECTURE: The architecture contains reusable bases and a coefficient table whose scalar entries determine each basis's contribution to each block.
- 4.2 LINEAR REUSABLE NEURAL BASES ARCHITECTURE: The coefficient formulation assumes real-valued independent coefficients, while nonnegativity or binary constraints are left for future investigation.
- 4.2 LINEAR REUSABLE NEURAL BASES ARCHITECTURE: The implementation averages accumulated gradients of reused bases across their reuse depth and performs forward computation through shared sensors, activations, coefficients, and responses.
- 4.2 LINEAR REUSABLE NEURAL BASES ARCHITECTURE: Sub-neural reuse lets multiple neurons share a response vector, forming multi-sensor neurons that can reduce parameters and alleviate dead-neuron issues.
5 DISCUSSION
LRNBA compresses networks by reusing neural bases across depths, making parameter count primarily width-dependent while enabling wider and deeper architectures. The discussion also describes coefficient energy as a pruning criterion and identifies training-time overhead as a limitation.
- Architecture shape: LRNBA’s parameter count is almost decoupled from depth and depends primarily on width, unlike classical networks.This supports substantially wider and deeper networks under a limited memory budget.
- Network compression: LRNBA introduces compression through reusable neural bases across depths, adding component-level and topological connection compression.Component reuse captures recurrent or cyclical patterns, while wider and deeper networks provide more opportunities to discover essential neuron interconnections.
- Limitations and training considerations: Simultaneously zero-initializing coefficient tables and response vectors should be avoided because their bilinear connection would make gradients vanish entirely.Other initialization schemes remain possible.
- Coefficient energy and pruning: The energy of a neural basis reflects its importance or frequency relative to the dataset and architecture and can guide pruning.The energy is defined from the basis coefficients across depths.
- Limitations and training considerations: LRNBA increases training time and energy consumption because block-specific weights are dynamically computed and updated during backpropagation.The overhead can be eliminated during inference by precomputing and caching effective block weights.
6 EXPERIMENTS
Experiments evaluate LRNBA on nonlinear regression with an FFN-based ResNet and language modeling with a decoder-only Transformer. Across both settings, LRNBA maintains competitive training outcomes with fewer parameters, while dynamic construction adds training-time overhead.
- Experimental settings: The experiments use an FFN-based ResNet for function approximation and a decoder-only Transformer for language modeling.The evaluations measure loss, training time, and parameter count, with test metrics also reported for the regression setting.
- FFN-based ResNet: 1/7 parameters: LRNBA matches standard ResNet convergence and final loss on nonlinear regression while substantially outperforming parameter-matched alternatives.The alternatives include a classical small model and a repeatedly shared residual-block baseline.
- FFN-based ResNet: 4L uses 1/6 the standard ResNet parameters while converging faster and reaching lower final loss.This outperforms the reported width expansion, where 4m uses half as many parameters with comparable final loss.
- FFN-based ResNet: LRNBA incurs a modest per-epoch training-time increase because residual blocks are dynamically constructed from reusable bases and coefficients.Precomputing and caching effective weights can remove this overhead during inference.
- Decoder-only Transformer: 6h6m achieves the fastest convergence and lowest training loss on WikiText-2 while using approximately 3/4 as many parameters as the standard Transformer.Increasing reusable attention and FFN bases consistently accelerates convergence and reduces final training loss.
- Decoder-only Transformer: Transformer training time increases approximately linearly with the number of reusable bases despite parameter savings.The caption relates this growth partly to limited GPU on-chip SRAM capacity and suggests larger on-chip memory or improved memory management as possible mitigations.
7 CONCLUSION
The paper introduces LRNBA to reduce parameter and memory requirements through reusable neural bases and linear coefficients. Experiments on ResNets and decoder-only Transformers report reduced parameter counts with stable training and competitive performance under limited memory budgets.
- Conclusion: LRNBA decomposes feedforward neural networks into reusable neural bases and linear coefficients to improve parameter efficiency and reduce memory requirements.The paper also gives this decomposition a geometric interpretation through vector fields.
- Conclusion: Experiments on FFN-based ResNets and decoder-only Transformers show substantially reduced parameter counts with stable training and competitive performance.Expanding reusable bases or network width and depth further improves convergence and final accuracy in several settings.
- Conclusion: LRNBA enables substantially wider and deeper networks under a limited memory budget.This conclusion is framed as a potential direction for memory-constrained deep-learning studies.
A.1 ATTENTION NEURAL BASES
Attention neural bases treat individual attention heads as reusable nonlinear basis functions whose outputs are scaled and combined to form the attention block.
- The complete attention output concatenates the coefficient-scaled head outputs and combines them through the output projection O.
- The query, key, and value representations for head j are q_j = Q_jx, k_j = K_jx, and v_j = V_jx, with d_H denoting head dimension.
- Each attention head A_j is treated as a nonlinear basis function that extracts a particular interaction pattern from the input sequence.
- The parameters Q_j and K_j form the sensor of attention neural basis A_j, while V_j forms its response.
- Scalar coefficients α_j weight the attention heads and determine their relative contributions to the final output.
A.2 PROJECTION NEURAL BASES
Projection neural bases decompose projection blocks into reusable elementary functions with sensor and response parameters, supporting a geometric vector-field interpretation.
- A sensor–response duplex structure uses two sequential linear layers with identity activation to implement a linear projection.
- The composite projection block uses sensor parameters (S, b) and response parameters P, with hidden dimension m.
- Rows s_j of S, entries b_j of b, and columns p_j of P define the parameters for each projection neuron.
- Each function ϕ_j is identified as a projection neuron or projection neural basis, with (s_j, b_j) as its sensor and p_j as its response.
- Each projection neural basis corresponds to a rank-one affine transformation, with p_jb_j serving as its bias.
- A residual feedforward block is interpreted as a superposition of vector fields, and combining bases can produce rotational dynamics beyond a single directional effect.
ADDITIONAL EXPERIMENTAL RESULTS OF FFN-BASED RESNET
On test-set evaluation, deeper LRNBA variants achieve lower error and more accurate function approximation than the compared classical and directly reused residual models.
- The fourfold-depth LRNBA model achieves the lowest test MSE, while the quarter-depth classical model produces the highest loss.
- LRNBA outperforms the baseline that directly reuses the same residual block.
- Depth expansion provides a larger reduction in test error than width expansion in the reported experiments.
- All models except the quarter-depth classical variant closely approximate the ground-truth function in prediction visualizations.
- The results indicate that LRNBA supports deeper networks under a limited parameter budget while preserving strong generalization performance.