SKI to go Faster: Accelerating Toeplitz Neural Networks via Asymmetric Kernels

Alexander Moreno, Jonathan Mei, Luke Walters

Keywords: language models, kernel methods, fourier analysis

TL;DR: We use structured kernel interpolation and a frequency domain relative position encoder to speed up Toeplitz Neural Networks.

Abstract: Toeplitz Neural Networks (TNNs) (Qin et al. 2023) are a recent sequence model with impressive results. They require $O(n\log n)$ computational complexity and $O(n)$ relative positional encoder (RPE) multi-layer perceptron (MLP) and decay bias calls. We aim to reduce both. We first note that the RPE is a non-SPD (symmetric positive definite) kernel and the Toeplitz matrices are pseudo-Gram matrices. Further 1) the learned kernels display spiky behavior near the main diagonals with otherwise smooth behavior; 2) the RPE MLP is slow. For bidirectional models, this motivates a sparse plus low-rank Toeplitz matrix decomposition. For the sparse component's action, we do a small 1D convolution. For the low rank component, we replace the RPE MLP with linear interpolation and use asymmetric Structured Kernel Interpolation (SKI) (Wilson et al. 2015) for $O(n)$ complexity: we provide rigorous error analysis. For causal models, ``fast'' causal masking (Katharopoulos et al. 2020) negates SKI's benefits. Working in the frequency domain, we avoid an explicit decay bias. To enforce causality, we represent the kernel via the real part of its frequency response using the RPE and compute the imaginary part via a Hilbert transform. This maintains $O(n \log n)$ complexity but achieves an absolute speedup. Modeling the frequency response directly is also competitive for bidirectional training, using one fewer FFT. We set a speed state of the art on Long Range Arena with minimal score degradation.

Submission Number: 10091