Scalable Neural Network Kernels

Abstract: We introduce the concept of scalable neural network kernels (SNNKs), the replacements of regular feedforward layers (FFLs), capable of approximating the latter, but with favorable computational properties. SNNKs effectively disentangle the inputs from the parameters of the neural network in the FFL, only to connect them in the final computation via the dot-product kernel. They are also strictly more expressive, as allowing to model complicated relationships beyond the functions of the dot-products of parameter-input vectors. We also introduce the neural network bundling process that applies SNNKs to compactify deep neural network architectures, resulting in additional compression gains. In its extreme version, it leads to the fully bundled network whose optimal parameters can be expressed via explicit formulae for several loss functions (e.g. mean squared error), opening a possibility to bypass backpropagation. As a by-product of our analysis, we introduce the mechanism of the universal random features (or URFs), applied to instantiate several SNNK variants, and interesting on its own in the context of scalable kernel methods. We provide rigorous theoretical analysis of all these concepts as well as an extensive empirical evaluation, ranging from point-wise kernel estimation to Transformers' fine-tuning with novel adapter layers inspired by SNNKs. Our mechanism provides up to 5x reduction in the number of trainable parameters, while maintaining competitive accuracy.

• The paper presents SNNKs as a novel module that replaces traditional FFLs, reducing trainable parameters by up to 5x while maintaining competitive performance.
• SNNKs leverage Universal Random Features and Fourier transforms to approximate FFL operations, enhancing computational efficiency and lowering storage needs.
• The approach enables bundling multiple FFLs into a single efficient kernel layer, benefiting resource-constrained applications in NLP and image recognition.

Scalable Neural Network Kernels

Introduction

The paper introduces Scalable Neural Network Kernels (SNNKs) as a novel computational module designed to replace traditional Feedforward Layers (FFLs). This approach allows for efficient computation by disentangling inputs from model parameters and utilizing them in the final dot-product computation through a kernel function. SNNKs offer several benefits, including model compression, computational efficiency, and theoretical insights into neural network architecture.

Figure 1: Architecture for (a) SNNK layer.

SNNK Design and Implementation

The SNNK module functions by approximating traditional FFL operations using random feature maps, which linearly transform input data and parameters separately. This transformation results in an overall computation that is less complex than standard FFLs—specifically, if the number of random features $m$ is much smaller than the input dimension $d$, then computational complexity and storage requirements are notably reduced.

These kernels rely on Universal Random Features (URFs) to achieve their approximation of FFLs, an approach that leverages Fourier Transforms of activation functions for constructing mappings $\Phi_{f}$ and $\Psi_{f}$. This mechanism provides scalable kernel methods that are independent of specific FFL design, offering adaptability and efficiency across various architectures.

Empirical Evaluation

SNNKs have been empirically assessed in multiple settings, showing effective compression while maintaining competitive accuracy. The mechanism enables up to a 5x reduction in trainable parameters without significant compromise in performance.

Performance on benchmark datasets, such as GLUE for NLP tasks and CIFAR for image recognition, demonstrates that SNNKs can effectively serve as a drop-in replacement for standard FFLs. Notably, SNNK-driven models maintain high accuracy and show improvement in parameter efficiency across different tasks.

Figure 2: Comparison of trainable parameters between various layers/modules and the drop-in replacement NNK layers.

Neural Network Bundling Process

A significant contribution of the paper is the concept of neural network bundling. Bundling is a method by which multiple FFLs are replaced with a single, efficient SNNK implementation, compressing the depth and parameter load of deep networks. This approach can directly benefit both the inference and training phases of neural networks, offering potential theoretical advantages by providing explicit parameter solutions under certain loss functions.

Practical Use Cases

The practical implications of SNNKs include their integration into Transformers for NLP tasks, where models require fine-tuning through efficient adapter layers, notably reducing parameter count. Furthermore, SNNKs' ability to decrease computational requirements makes them suitable for real-world applications constrained by resources, such as in mobile and embedded systems.

Conclusion

SNNKs present a promising direction for scaling efficient neural networks, combining the power of kernel methods with modern neural architecture design. This approach not only seeks computational efficiency but also provides new avenues for model interpretability and theoretical exploration. The realization of SNNKs in numerous applications demonstrates their viability as a replacement for traditional layers, setting a substantial precedent for further research into efficient neural computation.