Title: SKI to go Faster: Accelerating Toeplitz Neural Networks via Asymmetric Kernels
Venue: NeurIPS 2023
Accepted: no
Keywords: language models

kernel methods

fourier analysis
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.

Abstract

Toeplitz Neural Networks (TNNs) [1] are a recent impressive sequence model requiring 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 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 Structured Kernel Interpolation (SKI) [2] for O ( n ) complexity. For causal models, 'fast' causal masking [3] negates SKI's benefits. Working in 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 improve on speed and sometimes score on the Long Range Arena (LRA) [4].

(a) Long Range Arena (LRA).

(b) Pre-training speedups using Fourier Domain.

Percent Speed Up with FD TNN at 512 Seq Len

80

Bidirectional

Causal

70

60

50

%40

30

20

10

2

3

4

5

#of RPE Layers

1 Introduction 19

Sequence modeling is important in natural language processing, where sentences are represented 20 as a sequence of tokens. Successful sequence modeling typically involves token and channel 21 mixing. Token mixing combines representations of different sequence parts, while channel mixing 22 combines the information across different dimensions of embedding vectors used to encode tokens. 23 Transformers [5] are arguably the most successful technique for sequence modeling, and variants 24 including [6, 7] have achieved state of the art performance on natural language tasks. They use 25 self-attention for token mixing and feedforward networks for channel mixing. 26

Recently, [1] proposed Toeplitz Neural Networks (TNN) using Toeplitz matrices for token mixing. 27 They use a learned neural similarity function, the Relative Positional Encoder (RPE), to form 28 the Toeplitz matrices. Toeplitz matrix vector multiplication can be performed with sub-quadratic 29 complexity using the Fast Fourier Transform (FFT), giving the TNN token mixing layer a total 30 O ( dn log n ) computational complexity, where d is the embedding dimension and n is the sequence 31 length. This achieved state of the art predictive performance and nearly state of the art speed for 32 the long range arena (LRA) benchmark [4]. They also showed strong performance pre-training 33 wikitext-103 [8] and on the GLUE benchmark[9]. Despite strong empirical speed performance, 34 TNNs have two fundamental efficiency limitations: 1) super-linear computational complexity 2) 35 many calls to the RPE: for each layer, one call per relative position. 36

In this paper, we interpret the RPE as a non-SPD kernel and note 1) the learned kernels are discontin37 uous near the main diagonals but otherwise smooth globally; 2) the ReLU RPE learns 1D piecewise 38 linear functions: an MLP is slower than necessary. For bidirectional models, this motivates a sparse 39 plus low-rank decomposition. We apply the sparse component's action via a small 1D convolution. 40 For the low rank component, we replace the RPE MLP with linear interpolation at a set of inducing 41 points and an asymmetric extension of Structured Kernel Interpolation (SKI) [2] for O ( n ) complexity. 42 Further, using an inverse time warp, we can extrapolate beyond sequence lengths observed during 43 training. For causal models, even 'fast' causal masking [3] negates the speed and memory benefits 44 from SKI. Thus, we instead represent the real part of the kernel's frequency response using the RPE 45 MLP, and evaluate the RPE with finer frequency resolution to extrapolate to longer sequence lengths 46 in the time domain. From the real part, we compute the imaginary part via a Hilbert transform during 47 the forward pass to enforce causality. In the bidirectional setting, we remove the causality constraint 48 and represent the complex frequency response of the kernel with the RPE MLP. Levels of smoothness 49 in frequency response imply decay rates in the time domain: thus we model the decay bias implicitly. 50 This maintains O ( n log n ) complexity but achieves an absolute speedup. Further, it often leads to 51 better predictive performance on LRA tasks. 52

This paper has three primary contributions: 1) a TNN sparse plus low rank decomposition, extending 53 SKI to TNNs for the low rank part. We replace the RPE MLP with linear interpolation and apply 54 inverse time warping to efficiently train bidirectional TNNs. We provide rigorous error analysis 55 for our asymmetric SKI application; 2) alternatively, for both causal and bidirectional models, we 56 work directly in the frequency domain and use the Hilbert transform to enforce causality in the 57 autoregressive setting. We prove that different activation choices for an MLP modeling the discrete 58 time Fourier transform (DTFT) lead to different decay rates in the original kernel. 3) Empirical 59 results: we demonstrate that our approaches show dramatically improved computational efficiency, 60 setting a new speed state of the art on LRA [10] on the 1d tasks, with strong LRA score. In section 61 2 we describe related work. In section 3 we propose our new modeling approaches. In 4 we state 62 several theoretical results regarding our modeling approaches. In 5 we extend the empirical results of 63 [1], showing our speed gains with minimal prediction deterioration. We conclude in section 6. 64

2 Related 65

and introduce several techniques to improve on their speed results. [11] took a similar approach, but applied Toeplitz matrices to self-attention rather than departing from it. [12] is also similar, using

Also related are kernel based xFormers, particularly those using the Nyström method [13, 14]. The 71 most related work is [15], which adapts a matrix Nyström method for asymmetric matrices [16] to 72 self-attention. We instead adapt this along with SKI [2] to Toeplitz matrices. [17] extends [15] by 73 embedding the self-attention matrix into a larger PSD kernel matrix and approximating the larger 74 matrix instead. Their final approximate matrix has lower spectral error compared to [15] and higher 75 average validation accuracy on LRA [4]. However, their method is slightly slower. Also somewhat 76 related are random feature self-attention approximations[18, 19]. These extend [20], but use different 77 random features that better approximate self-attention than random Fourier or binning features. 78

Sparse transformers are also relevant. [21] proposed using strided and fixed patterns. [22] alternated 79 between sparse locally banded and dense attention. Finally, [23] proposed combining random 80 attention, window attention and global attention. Our use of a short convolutional filter is most similar 81 to window attention. The space of efficient transformers is huge and there are many models that we 82 haven't covered that may be relevant. [10] provides an excellent survey. 83

Other successful long sequence approaches include state space models [24, 25, 26], long convolution 84 [27, 28], adding moving averages to gated attention [29] and more [30]. 85

3 Modeling Approach 86

We review Toeplitz neural networks (TNNs) in section 3.1. We next speed up the TNN's Toeplitz 87 neural operator (TNO). We discuss using Nyström and SKI approaches to bidirectional training in 88 3.2. We discuss frequency based approaches, particularly for causal training in 3.3. 89

90

3.1 Preliminaries: Toeplitz matrices and Toeplitz Neural Networks

TNNs [1] replace self-attention, which computes the action of self-attention matrices that encode 91 the similarity between both observation values and absolute positions, with the action of Toeplitz 92 matrices that encode similarity only based on relative positions. Toeplitz matrices have, for each 93 diagonal, the same entries from left to right. That is, T ij = t i -j , T ∈ R n × n . Unlike self-attention 94 matrices, which require O ( n 2 ) memory, a Toeplitz matrix has 2 n -1 unique elements and requires 95 O ( n ) memory. Due to close connections with discrete-time convolution, Tx can be computed in 96 O ( n log n ) time by embedding T in a circulant matrix and applying FFT. 97

98

99

100

101

A TNN [1] has multiple sequence modeling blocks, which we show in Figure 3 in Appendix A. Each block has a Gated Toeplitz Unit (GTU), which does both token and channel mixing, followed by a Gated Linear Unit (GLU) [31], which does channel mixing. The core of the GTU is the Toeplitz Neural Operator (TNO), which does token mixing and is the part of the architecture that we modify.

102

103

104

105

We now describe the TNO, shown in Figure 3b of Appendix A. Given a sequence X ∈ R n × d of length n and dimension d in discrete time, there are 2 n -1 unique relative positions/times i -j for i, j = 1 , . . . , n . An RPE : Z → R d neural network maps each relative position to a d -dimensional embedding. These embeddings are used to construct Toeplitz matrices T l for l = 1 , . . . , d using

RPE l ( i -j ) is a learned similarity between positions for dimension l , while λ | i -j | with λ ∈ (0 , 1) 106 is an exponential decay bias penalizing far away tokens to be dissimilar. We can interpret T l ij as 107 evaluating a stationary non-SPD kernel k l ( i -j ) = λ | i -j | RPE l ( i -j ) . Thus T l can be interpreted 108 as a pseudo or generalized Gram matrix. Letting x l be the l th column of X , the TNO outputs 109

where each T l x l is computed via the FFT as described above. 110

The main costs are the RPE's MLP, the FFT, and the decay bias. We aim to eliminate the MLP and 111 decay bias when possible. In the bidirectional setting, we use SKI to apply the FFT using a much 112 smaller Toeplitz matrix. In a separate model we learn the RPE's frequency response directly. In the 113 bidirectional setting, this allows us to both avoid explicitly modeling the decay bias and use one fewer 114 FFT. In the causal setting, it allows us to avoid explicitly modeling the decay bias. 115

SKI TNO (Concept)

Smooth

Output

Input

x d

Interpolate

RPE

2r-1

Toeplitz

Sparse

x d

?m/2?+1

Banded

(a) SKI-TNO.

SKI TNO (Computation)

Approx. smooth

Output

Input

x d

Interpolate

RPE

(

)

,

2r-1

Sparse (W)

Conv.

x d

m

*

(b)

Fast

implementation

of

Output

Input

x d

MLP

RPE

FFT

IFFT

.

2n-1

(c) FD-TNO causal.

TNO FD

Output

Input

x d

MLP

RPE

FFT

IFFT

.

Freq. responses

2n-1

(d) FD-TNO bidirec- tional.

SKI-TNO.

3.2 SKI Based Approaches for Bidirectional Training 116

For a given Toeplitz matrix T , we assume it admits a decomposition that we can approximate with a 117 sparse+low-rank representation, T = T sparse + T smooth ≈ T sparse + T low. Our bidirectional training 118 thus consists of three primary components. The first, the sparse component T sparse is straightforward. 119 Applying the action T sparse x of T sparse ∈ R n × n with m non-zero diagonals is equivalent to applying 120 a 1D convolution layer with filter size m . We then discuss our asymmetric SKI for T low in section 121 3.2.1. Finally, we discuss how we handle sequence lengths not observed in training for T low via an 122 inverse time warp in section 3.2.2. Algorithm 1 summarizes our TNO based on these techniques. 123

Algorithm 1 Sparse Plus Low Rank Bidirectional TNO with Asymmetric SKI

Given sequence X ∈ R n × d with columns x l

Hyperparameters rank r /lessmuch n , sparse filter size m , interpolation degree N , decay parameter λ Compute inducing points p 1 , . . . , p r evenly spaced on [0 , n ]

Compute T l sparse x l with a 1D convolutional filter, size m .

Let x ( t ) = sign ( t ) λ | t | .

Form W l ∈ R n × r degree N polynomial interpolation matrix

Compute

end for

3.2.1 SKI For Asymmetric Nyström 124

Given an asymmetric stationary kernel k : R × R → R , we wish to approximate the (pseudo) Gram 125 matrix T ∈ R n × n using a low-rank approximation based on a smaller Gram matrix A ∈ R r × r , with 126 r /lessmuch n . In context, A is formed using relative positions between a set of inducing points p 1 , . . . , p r 127 instead of the full set 1 , . . . , n that is used for T . That is, 128

In our case, the inducing points are uniformly spaced. Some submatrices of A may be submatrices of 129 T (if inducing points are also observation points). To derive the Nyström approximation, we form an 130

T

l

low

x

with

T

l

low

=

W

A

W

l

l

l

l

/latticetop

augmented Gram matrix K ∈ R ( n + r ) × ( n + r ) in block form as 131

where B ∈ R r × n and F ∈ R n × r are respectively the upper right and lower left partitions of the large 132 Gram matrix K . Explicitly, 133

Extending [16] to allow singular A , 134

where A † is the Moore-Penrose pseudo-inverse satisfying AA † A = A (but not necessarily AA † = I 135 as in [16], which shows up in our different expressions for off-diagonal blocks of ̂ K ). Following 136 structured kernel interpolation (SKI) [2], we approximate F and B using interpolation. Specifically, 137

where W ∈ R n × r is a matrix of sparse interpolation weights with up to two non-zero entries per row 138 for linear interpolation or up to four for cubic. These weights can be computed in closed form from 139 the inducing points p i and the observation points i . Thus we have 140

˜ as desired. We can set T low = ˜ T and compute ˜ Tx by first applying W /latticetop x , which is an O ( n ) 141 operation due to W ∈ R n × r having sparse rows. Next, we apply A ( W /latticetop x ) . Since A is a Toeplitz 142 matrix, this is O ( r log r ) as per Section 3.1. Finally, W ( AW /latticetop x ) , the action of W , is again an O ( n ) 143 operation. Thus computing ˜ Tx is O ( n + r log r ) computation. On a GPU, this factorization achieves 144 a speedup from having small r and being able to leverage efficient parallelized matrix multiplication 145 on specialized hardware. However, in PyTorch [32], we note that for medium sized matrices up to 146 n = 512 , the time required for data movement in order to perform sparse-dense matrix multiplications 147 can be higher than that of simply performing dense matrix multiplication. This means that in practice, 148 we may instead choose to perform batched dense matrix multiplication, which yields an absolute 149 speedup but a worse asymptotic complexity of O ( nr 2 + r log r ) . 150

151

152

153

154

155

156

157

158

159

160

161

162

3.2.2 Inverse Time Warp

TNNs use k l ( i -j ) = λ | i -j | RPE l ( i -j ) , where RPE l ( i -j ) is an MLP. There are two issues: 1) the sequential computations required for an MLP are slow, and we only need to evaluate at 2 r -1 points using SKI instead of 2 n -1 to produce the full matrix; 2) extrapolation is used in extending to longer sequence lengths than the MLP was trained on, which is generally less reliable than interpolation.

In Proposition 1, we note that an MLP f : R → R d with ReLU activations and layer normalization is d piecewise linear functions. As we only need to evaluate at 2 r -1 points, we could let RPE l be a piecewise linear function with r grid points. However, we still need to handle extrapolation. We use an inverse time warp and let RPE l linearly interpolate on [ -1 , 1] with the constraint RPE l (0) = 0 and define x ( t ) = sign ( t ) λ | t | for some 0 < λ < 1 . We then let k l ( i -j ) = RPE l ( x ( i -j )) .

3.3 Frequency Based Approaches

3.3.1 Causal Training

The SKI approach allows training bidirectional TNNs with linear complexity. However, fast causal 163 masking negates SKI's benefits (see Appendix B). Thus we need an alternate causal speedup. We 164 use an MLP in the Fourier domain to avoid an explicit time domain decay bias, and use the Hilbert 165 transform to enforce causality. We now describe how we can learn a causal kernel when working in 166 frequency domain (FD). We first define the discrete Hilbert transform, the key tool for achieving this. 167

Definition 1. The discrete Hilbert transform of the discrete Fourier transform ˆ k is given by 168

where ∗ denotes convolution and

The real and imaginary parts of the Fourier transform of a causal function are related to each other through the Hilbert transform. Thus, in order to represent a causal signal, we can model only the real part and compute the corresponding imaginary part. That is, we first estimate an even real function ˆ k (symmetric about 0 ) using an MLP. We then take ˆ k causal ( ω ) = ˆ k ( ω ) -i H{ ˆ k } ( ω ) .

The inverse Fourier transform k causal of ˆ k causal will thus be causal. For a discussion of why this ensures causality, see [33]. See Algorithm 2 for TNO pseudocode using this approach. Different choices for the smoothness of the frequency domain MLP will lead to different decay rates in time domain, so that smoothness in frequency domain essentially serves the same purpose as the decay bias in [1]. We discuss this theoretically in Section 4.2. Note that we also find that working directly in the frequency domain for bidirectional models (without the Hilbert transform) is often competitive with SKI for speed (despite being O ( n log n ) instead of O ( n + r log r ) ) due to needing one fewer FFT.

Algorithm 2 Causal TNO via Discrete Hilbert Transform

3.3.2 Bidirectional Training with FD TNN

We extend the FD approach to bidirectional training by removing the causality constraint and model the complex frequency response of real valued time domain kernels directly. To do so we simply double the output width of the RPE and allocate each half for the real and imaginary parts of the kernel frequency responses, while explicitly forcing real-valued responses at ω = 0 and π . While increasing the complexity of the RPE slightly, we achieve the speed ups in Figure 1 by eliminating the FFTs for the kernels and causality constraint, in addition to the decay bias.

4 Theory

We show in Proposition 1 that an MLP mapping from scalars with layer norm and ReLU activations is piecewise linear and continuous, suggesting that using an MLP that we only need to evaluate at a small number of points may be overparametrized, justifying the use of interpolated piecewise linear functions. In section 4.1 we analyze the spectral norm of the matrix approximation error for SKI. We assume the sparse component is exactly identifiable and bound the error of approximating the smooth term via a low-rank SKI factorization. We leave the problem of relaxing this assumption to future work. In section 4.2, we analyze how by using different activations with different smoothness when learning the DTFT of the kernel, we obtain corresponding decay rates for the time domain signal.

Proposition 1. A ReLU MLP f : R → R d with layer norm and no activation on its output is d piecewise linear continuous functions.

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

232

233

4.1 Matrix Approximation Spectral Norm Error

We give our main error bound for our SKI based low rank approximation. Note that this requires that our kernel is N +1 times continuously differentiable, while the kernel we use in practice uses a piecewise linear function and is thus non-differentiable. In theory, we would need a smoother kernel, adding additional computation overhead. However, we find that empirical performance is still strong and thus we simply use piecewise linear kernels but include the error bound for completeness. Our results depends on the Nyström error E nyst : its l 2 norm is bounded in [16].

Theorem 1. Assume that A is non-singular and k : [ p 1 , p r ] → R is an N +1 times continuously differentiable function, where p 1 is the smallest inducing point and p r is the largest. Let T r,opt be the optimal rank r approximation to T and let

be the difference between the SKI approximation using linear interpolation and the optimal one, while

is the difference between the Nyström approximation and the optimal one. Then

where ψ N ( i ) = ∏ N j =1 ( i -p n j ) with p n j being the j th closest inducing point to i , L is an upper bound on the N +1 th derivative of k , and σ i ( M ) denotes the i th largest singular value of matrix M .

For linear interpolation | ψ N ( i ) | ( N +1)! ≤ h 2 8 , where h is the spacing between two neighboring inducing points. We have considered the sparse component of the Toeplitz matrix to be identifiable and focused on the error of approximating the smooth component. While there are potential approaches to relaxing this assumption [34, 35, 36, 37, 38, 39, 40], they must be adapted properly to the Toeplitz setting. Thus, this additional analysis is outside the scope of this paper and a fruitful direction for future work.

4.2 Smoothness in Fourier Domain Implies Decay in Time Domain

We now discuss activation function choices when directly learning the discrete time Fourier transform (DTFT) ˆ k as an MLP. In practice, we sample the DTFT to obtain the actually computable discrete Fourier transform (DFT) by evaluating the MLP with uniform spacing. Different levels of smoothness of the MLP ˆ k imply different decay rates of the signal k . One can think of the choice of activation function as a parametric form for the decay bias. For an MLP, using a GeLU activation implies super-exponential time domain decay. Using SiLU implies super-polynomial time domain decay. For ReLU the signal is square summable. While this subsection focuses on the theoretical relationship between smoothness and decay, in Appendix E.3 we show visualizations demonstrating that these relationships are observed in practice. We first define the DTFT and its inverse.

Definition 2. The discrete time Fourier transform [41, 33] ˆ k or F{ k } of k is given by

Definition 3. The inverse discrete time Fourier transform of the DTFT ˆ k is given by 231

We now give three theorems relating smoothness of the DTFT to decay of the signal (its inverse).

Theorem 2. Using a GeLU MLP for the DTFT ˆ k , for all a > 0 , the signal k [ n ] will have decay

Theorem 3. Using a SiLU MLP for the DTFT ˆ k , the signal k [ n ] will have decay 235

for all n = 0 , N ∈ N . 236

237

238

239

240

241

242

243

244

245

246

247

248

249

250

251

252

253

254

255

256

257

258

259

260

261

262

263

264

265

266

/negationslash

Theorem 4. Using a ReLU MLP for the DTFT ˆ k implies ‖ k ‖ 2 < ∞ (the signal is square summable).

Proof. Note that ˆ k ∈ L 2 [ -π , π ] since it is continuous. Then apply Parseval's theorem.

5 Experiments

We perform experiments in two areas: pre-training a causal language model on Wikitext-103 [8] and training bidirectional models on Long-Range Arena. We start with the repositories of the TNN paper 1 and use their training and hyper-parameter settings unless indicated otherwise. We use A100 and V100s for training, and a single A100 for timing experiments.

5.1 Pre-training on Wikitext-103

In the causal case we aim to predict the next token, conditional on a fixed length sequence of previous tokens. Table 1 compares FD-TNN's causal pre-training perplexity [8] to existing models: it almost exactly matches that of TNNs. Our approach is faster for the same capacity: at sequence length 512 with 6 layer RPEs (as in the TNN paper), FD TNN is 15% faster than the baseline TNN on a single A100 GPU. When both use a three layer RPE, FD TNN is 10% faster. We provide some additional details for this experiment as well as for bidirectional pre-training (we see larger speed gains) in Appendix F.

5.2 Long-Range Arena

The Long-Range Arena (LRA) is a benchmark with several long sequence datasets. The goal is to achieve both high LRA score (predictive performance) and training steps per second. Following [1], we take the TNN architecture and their tuned hyperparameter (HP) configurations 2 , simply replacing their TNO module with our SKI-TNO module with r = 64 and m = 32 . We use λ = 0 . 99 where they set λ = 1 , but otherwise perform no additional HP tuning on 1D tasks and use smaller layers r = 32 and m = 16 for the 2D tasks. For FD-TNN, we simply use a same-sized RPE for all tasks except a 3-layer RPE for the CIFAR task. We could potentially achieve even higher accuracy with more comprehensive tuning on the 2D tasks or any tuning for the 1D tasks. We select the checkpoint with the highest validation accuracy and report the corresponding test accuracy. SKI-TNN achieves similar average accuracy than TNN at lower size, while FD-TNN achieves higher accuracy. We suspect that for some of these problems, the square summable signal implied by ReLU in frequency domain is a better parametric form than applying exponential decay bias. We show our results in Table 2.

We additionally perform timing and memory profiling tests on a single 1x A100 instance, keeping 267 the per-GPU batch size constant as in the training runs. In Figure 1a, we plot for each 1D task the 268 percentage of TNN accuracy achieved vs the percentage speedup relative to TNN, with the size of 269 the marker corresponding to the peak memory usage measured. We highlight the 1D tasks because 270 they required no tuning, and they represent the longest sequences at lengths ranging from 1024 to 271 4096 , whereas the 2D tasks are treated as separate 1D sequences in each dimension, so that a 32 × 32 272 image is seen as alternating length 32 sequences. We note that because the effective sequence lengths 273 are shorter, there is less benefit from using our methods over the baseline TNN. 274

275

6 Conclusion

In this paper, we note that [1]'s Toeplitz neural networks essentially apply the action of a generalized 276 Gram matrix (the Toeplitz matrix) for an asymmetric kernel (the RPE times decay bias) as their main 277 computationally expensive operation. The visualized learned Gram matrices motivate a sparse and 278 low rank decomposition. We thus propose two different approaches to improve efficiency. In the 279 bidirectional setting, we extend SKI to the asymmetric setting and use linear interpolation over a 280 small set of inducing points to avoid the MLP entirely, while using an inverse time warp to handle 281 extrapolation to time points not observed during training. This approach reduces the mathematical 282 complexity from O ( n log n ) to O ( n + r log r ) , where r is the number of inducing points. However 283 in practice, we do not actually use O ( n + r log r ) code due to a reshape required for sparse tensors 284 leading to them actually being slower than dense tensors. Thus we actually use O ( nr 2 + r log r ) in 285 code: still much faster than Baseline TNN for small r . For causal training, as causal masking negates 286 SKI's benefits, we instead eliminate the explicit decay bias. We do this by working directly in the 287 frequency domain, enforcing causality via the Hilbert transform and enforcing decay in time domain 288 via smoothness. For the bidirectional case, we eliminate the FFT applied to the kernels. While this 289 maintains O ( n log n ) computational complexity, it leads to a substantial speedup in practice and 290 beats TNNs on LRA score. 291

292

293

294

295

296

297

298

299

300

301

302

303

304

305

306

307

308

309

310

311

312

313

314

315

316

317

318

319

320

321

322

323

324

325

326

327

328

329

330

331

332

333

334

References

Zhen Qin, Xiaodong Han, Weixuan Sun, Bowen He, Dong Li, Dongxu Li, Yuchao Dai, Lingpeng Kong, and Yiran Zhong. Toeplitz neural network for sequence modeling. In The Eleventh International Conference on Learning Representations, 2023.

Andrew Wilson and Hannes Nickisch. Kernel interpolation for scalable structured Gaussian processes (KISS-GP). In International conference on machine learning, pages 1775-1784. PMLR, 2015.

Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are RNNs: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning, pages 5156-5165. PMLR, 2020.

Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long Range Arena: A Benchmark for Efficient Transformers. In International Conference on Learning Representations, 2020.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Ł ukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.

Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland, Katherine Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Oriol Vinyals, Jack William Rae, and Laurent Sifre. An empirical analysis of compute-optimal large language model training. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.

Kevin Clark, Minh-Thang Luong, Quoc V. Le, and Christopher D. Manning. Electra: Pretraining text encoders as discriminators rather than generators. In International Conference on Learning Representations, 2020.

Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer Sentinel Mixture Models. In International Conference on Learning Representations, 2016.

Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. In International Conference on Learning Representations.

Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. ACM Computing Surveys, 55(6):1-28, 2022.

Shengjie Luo, Shanda Li, Tianle Cai, Di He, Dinglan Peng, Shuxin Zheng, Guolin Ke, Liwei Wang, and Tie-Yan Liu. Stable, fast and accurate: Kernelized attention with relative positional encoding. Advances in Neural Information Processing Systems, 34:22795-22807, 2021.

Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena Hierarchy: Towards Larger Convolutional Language Models. arXiv preprint arXiv:2302.10866, 2023.

Evert J Nyström. Über die praktische auflösung von integralgleichungen mit anwendungen auf randwertaufgaben. Acta Mathematica, 54(1):185-204, 1930.

Christopher TH Baker. The numerical treatment of integral equations. Oxford University Press, 1977.

Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin 335 Li, and Vikas Singh. Nyströmformer: A nyström-based algorithm for approximating self336 attention. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 337 14138-14148, 2021. 338

Arik Nemtsov, Amir Averbuch, and Alon Schclar. Matrix compression using the Nyström 339 method. Intelligent Data Analysis, 20(5):997-1019, 2016. 340

Yifan Chen, Qi Zeng, Heng Ji, and Yun Yang. Skyformer: Remodel self-attention with gaussian 341 kernel and nyström method. Advances in Neural Information Processing Systems, 34:2122342 2135, 2021. 343

H Peng, N Pappas, D Yogatama, R Schwartz, N Smith, and L Kong. Random Feature Attention. 344 In International Conference on Learning Representations, 2021. 345

346

347

348

349

Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking Attention with Performers. In International Conference on Learning Representations, 2021.

350

351

352

353

354

355

356

357

358

359

360

361

362

363

364

365

366

367

368

369

370

371

372

373

374

375

376

377

378

379

380

381

382

383

384

385

Ali Rahimi and Benjamin Recht. Random features for large-scale kernel machines. Advances in neural information processing systems, 20:1177-1184, 2007.

Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509, 2019.

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.

Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big Bird: Transformers for Longer Sequences. Advances in neural information processing systems, 33:17283-17297, 2020.

Albert Gu, Karan Goel, and Christopher Re. Efficiently modeling long sequences with structured state spaces. In International Conference on Learning Representations.

Jimmy T.H. Smith, Andrew Warrington, and Scott Linderman. Simplified state space layers for sequence modeling. In The Eleventh International Conference on Learning Representations, 2023.

Tri Dao, Daniel Y Fu, Khaled K Saab, Armin W Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models. arXiv preprint arXiv:2212.14052, 2022.

David W Romero, Anna Kuzina, Erik J Bekkers, Jakub Mikolaj Tomczak, and Mark Hoogendoorn. Ckconv: Continuous kernel convolution for sequential data. In International Conference on Learning Representations.

Daniel Y Fu, Elliot L Epstein, Eric Nguyen, Armin W Thomas, Michael Zhang, Tri Dao, Atri Rudra, and Christopher Ré. Simple hardware-efficient long convolutions for sequence modeling. arXiv preprint arXiv:2302.06646, 2023.

Xuezhe Ma, Chunting Zhou, Xiang Kong, Junxian He, Liangke Gui, Graham Neubig, Jonathan May, and Luke Zettlemoyer. Mega: Moving average equipped gated attention. In The Eleventh International Conference on Learning Representations, 2023.

Ruslan Khalitov, Tong Yu, Lei Cheng, and Zhirong Yang. Chordmixer: A scalable neural attention model for sequences with different length. In The Eleventh International Conference on Learning Representations, 2023.

Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.05202, 2020.

Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.

Alan V Oppenheim and Schafer R W. Discrete Time Signal Processing. Prentice-Hall, 2010. 386

Benjamin Recht, Maryam Fazel, and Pablo A Parrilo. Guaranteed minimum-rank solutions of 387 linear matrix equations via nuclear norm minimization. SIAM review, 52(3):471-501, 2010. 388

Emmanuel J Candes and Yaniv Plan. Matrix completion with noise. Proceedings of the IEEE, 389 98(6):925-936, 2010. 390

Tianyi Zhou and Dacheng Tao. Godec: Randomized low-rank & sparse matrix decomposition in 391 noisy case. In Proceedings of the 28th International Conference on Machine Learning, 2011. 392

Jonathan Mei and José M F Moura. SILVar: Single Index Latent Variable Models. IEEE 393 Transactions on Signal Processing, 66:2790 - 2803, 3 2018. 394

Venkat Chandrasekaran, Sujay Sanghavi, Pablo A Parrilo, and Alan S Willsky. Rank-sparsity 395 incoherence for matrix decomposition. SIAM Journal on Optimization, 21(2):572-596, 2011. 396

Venkat Chandrasekaran, Pablo A. Parrilo, and Alan S. Willsky. Latent variable graphical model 397 selection via convex optimization. Ann. Stat., 40:1935-1967, 8 2012. 398

Teng Zhang and Yi Yang. Robust PCA by manifold optimization. The Journal of Machine 399 Learning Research, 19(1):3101-3139, 2018. 400

John G Proakis and Dimitris G Manolakis. Introduction to digital signal processing. Prentice 401 Hall Professional Technical Reference, 1988. 402

Jeffrey Wong. Math 563 lecture notes, polynomial interpolation: the fundamentals, 2020. 403 URL:https://services.math.duke.edu/ jtwong/math563-2020/lectures/Lec1-polyinterp.pdf. 404

Mhenni Benghorbal (https://math.stackexchange.com/users/35472/mhenni benghorbal). How 405 to prove error function erf is entire (i.e., analytic everywhere)? Mathematics Stack Exchange, 406 2017. URL:https://math.stackexchange.com/q/203920 (version: 2017-04-13). 407

Christopher Heil. Introduction to Real Analysis, volume 280. Springer, 2019. 408

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe present paper offers a Toeplitz matrix architecture which can handle sequence modeling. The architecture comes in two flavors. The first flavor is a fast version that is most useful for bi-directional tasks. It speeds up previous Toeplitz networks by using an interpolation and low rank approximation scheme in its setup. The second is a Fourier based model that appears to offer advantages in causal tasks. Benefits are shown in terms of the speed of training and some marginal benefits in performance.\n\nSTRENGTHS:\nThe main strength of this paper is a noticeable speed up in an alternative architecture to transformers. I am a fan of papers that look at ways to speed up sequence modeling. The Curren paper presents a nice idea. More specifically:\n- The approximations to the TNN appear to be effective in experiments and result in a faster network\n- Approximations do not seem to deteriorate performance and may offer some added performance boosts\n- The changes to the architecture are grounded in some theory\n\nWEAKNESSES:\n**Disclaimer**: I am not an NLP expert and am more focused on the theory side. I have significant concerns about the theory in this paper, but feel the experiments and techniques are solid enough to potentially overcome that issue. Furthermore, none of the theorems are crucial to the crux of the paper, and if any are wrong, they can be removed. For this reason, I placed a borderline accept rating for now, but I believe this will need to be confirmed by people who are closer to the experimental side of the literature and can assess the experiments in a more rigorous fashion. \n\n\\\nBroader comments:\n- Most of my larger concerns revolve around the theory and proofs in this project listed below.\n- Reading through many times, I could not understand what was gained in the “causal training” proposal in section 3.1. First, it seems the parameters are changed to live completely in the Fourier regime. This change was made to obtain “an alternate causal speedup\", but I don’t see where that speedup arises. Following the steps, the main change seems to take the algorithm to do an FFT on the $n$-dimensional space resulting in runtime of $O(n \\log n)$ which is worse than before. Also, changing the algorithm to work in Fourier space introduces a different implicit bias that I’m not sure is desired. For example, the $\\lambda$ parameter controlling the decay is not controlled here. I also have many questions about this approach which I’ve left below.\n- Experiments in table 1 don’t appear to offer much improvement especially considering the added parameters. I would also ask the authors to include citations to the models or results compared to in this table so it is easier to see what is being compared to. \n\n\\\nTheory comments:\n- Theorem about ReLU MLPs being $d$-piecewise linear has assumptions missing or is just wrong. If the authors are implying that any ReLU network from $\\mathbb{R} \\to \\mathbb{R}^d$ has $d$ pieces or contiguous linear regions, this is clearly wrong. MLPs are universal approximators so this is clearly false. If the authors are saying that this only holds for an MLP with a single hidden layer of width at most $d$ then this may be correct. But I don’t see this assumption made anywhere.\n- Theorem 2 has a few confusing elements from my end. First, it is hard to parse. There are many variables and factors like condition number that it is hard to know the scaling of. Second, the bound doesn’t appear to be all that good. The error grows at least linear in $n$ and depends on other factors like singular values or the nystrom error that may also be badly bounded. I suppose the authors would argue that it is exponentially small in the degree of interpolation $N$, but this degree would have to grow at least logarithmically in $n$ to counteract the $n$ factor. This would result in a runtime essentially equivalent to just doing FFTs on the whole space. Third, the error in interpolation is an unusual thing to even bound in my opinion. The weights are updated with this interpolation taken into account. In other words, the algorithm learns a weight matrix with parameters contained in this interpolation. \n- Definition 2 and 3 present the discrete time Fourier transform, but in practice only the DFT of the matrix form is ever used. The resulting statements, regardless of their correctness, do not seem to apply to the setting in practice. Unless I am missing something.\n- Related to the above, I cannot see why Theorem 3 and 4 are correct. Similar to my previous statements, MLPs are universal approximations so they can output any possible function. How can that statement hold true? \n\n\\\nSmall:\n- Simply having a % label on the y axis of Fig 1B is confusing. Percentage relative to what? Also, what does 20% speed-up mean; i.e. that it ran in 20% less time? Simply having this number on the first page can be confusing without the context added.\n- Fig 1A and 1B also appear to be different size fonts.\n- Line 150: I think the dense-case runtime is only a factor of $r$ worse so $O(nr  + r \\log r)$ and not $O(nr^2 + r \\log r)$ unless I’m missing something.\n- For someone outside of the NLP community, section 3.1 needed more motivation and formality. Some details about what “causal masking”, “causal kernel”, and the sequential nature of the data would be helpful.\n- To follow easier, it would be good to define what the role of N is in section 4.1 (i.e., number of interpolating points)\n\n\\\nFormatting:\n- Hyperref links seem to be broken\n- Line 189: sentence is a run-on and hard to follow\n\n\\\nFinally, to add ideas not for criticism, but instead for completing the paper and offering new ideas, there is a wealth of literature on related techniques that could be useful here, or at the very least cited. For example, there are sparse Fourier transforms that can offer speed-ups beyond the $O(n \\log n)$ that the paper aims to improve on, e.g., [1]. Since matrices are low-rank and/or sparse, this could be a more direct way to get the speed-ups desired. Second, there are a number of papers on optimizing structured matrices like Toeplitz matrices, e.g. [2]. In the sequence modeling specifically, there have been a lot of papers on unitary networks for example [3], one paper which actually uses low rank approximations in its implementation [4]. Third, low rank approximations have also been used to speed up other architectures like conv-nets, [5-6]\n\n\n\\\n**References:** \\\n[1] Hassanieh, Haitham, et al. \"Simple and practical algorithm for sparse Fourier transform.\" Proceedings of the twenty-third annual ACM-SIAM symposium on Discrete Algorithms. Society for Industrial and Applied Mathematics, 2012.\\\n[2] Kochurov, Max, Rasul Karimov, and Serge Kozlukov. \"Geoopt: Riemannian optimization in pytorch.\" arXiv preprint arXiv:2005.02819 (2020).\\\n[3] Arjovsky, Martin, Amar Shah, and Yoshua Bengio. \"Unitary evolution recurrent neural networks.\" International conference on machine learning. PMLR, 2016.\\\n[4] Kiani, Bobak, et al. \"projUNN: efficient method for training deep networks with unitary matrices.\" arXiv preprint arXiv:2203.05483 (2022).\\\n[5] Max Jaderberg, Andrea Vedaldi, and Andrew Zisserman. Speeding up convolutional neural networks with low rank expansions. arXiv preprint arXiv:1405.3866, 2014.\\\n[6] Cheng Tai, Tong Xiao, Yi Zhang, Xiaogang Wang, et al. Convolutional neural networks with low-rank regularization. arXiv preprint arXiv:1511.06067, 2015.\n\nQUESTIONS:\nBeyond what's written above, I have the following questions: \n- How do you ensure the complex valued MLP gives real valued outputs in algorithm 2 after the iFFT?\n\n\nQuestions about causal training:\n- Why do we need a Hilbert transform when we can just do a regular FFT on a standard Toeplitz matrix and it results in the same runtime?\n- Is the $\\lambda$ factor which decays farther out correlations just ignored here?\n- Aren’t there much easier ways to enforce a causal filter? E.g., only have the Toeplitz matrix values nonzero along the left of the diagonal of the matrix.\n\nLIMITATIONS:\nThere is a very brief discussion of limitations in the conclusion, though I feel this could be expanded. I would also appreciate some context for this work in relation to other works in NLP and how it fits into the broader landscape of NLP architectures. For someone like me not in the community, this would be useful to better understand its limitations from a practical perspective.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-16"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors of the paper propose two modifications of a recently published alternative to attention mechanism, Toeplitz Neural Operator (TNO), which constitutes the most important part of Toeplitz Neural Networks. The application of TNO is the multiplication of the input sequence by a Toeplitz matrix. Parameters of this Toeplitz matrix are given by a lightweight feed-forward network, called Relative Position Encoder (RPE). The first proposed modification, called SKI-TNN, represents a learned Toeplitz matrix as a sum of a sparse and low-rank matrix. Thus, its multiplication by a vector has the complexity of O(nr^2 + r log r) instead of O(n log n), where r is the rank of a second summand. However, this modification can speed up only the task of bidirectional modeling. In order to speed up the causal modeling (such as autoregressive language modeling), authors view the TNO as an application of kernel to a vector. Then, they train the RPE to model the real part of the Fourier transform of this kernel. The imaginary part is then computed via the Hilbert transform of the real part. This modification does not change the asymptotic complexity, but achieves empirical speed up.\n\nSTRENGTHS:\n1. The article explores an important topic of speeding up the token mixing part of a general sequence modeling pipeline. Nowadays, this topic is highly relevant because of its applications in the field of NLP. The article builds off of a very recent paper [1].\n2. The article creatively combines together a large body of previous work. It uses the ideas of TNN, SKI, FFT, Hilbert transform, Nyström approximation, fast causal masking, etc. \n3. The experimental results show that the proposed modifications do indeed speed up the original TNN.\n4. Overall, the presentation style is mathematically strict and to the point. The formulae in section 3.2.1 and in Appendix are sufficiently well-explained. Both modifications proposed in the paper are succinctly defined in Algorithm 1 and Algorithm 2. This helps the reader significantly to understand the main ideas.\n5. In section 3.2.1, the authors specified not only the theoretical complexity of their modification, but also the practical limitations they meet when implementing it, and specified the practical complexity as well as theoretical.\n6. The theory on the smoothness in Fourier Domain is supported with experimental visualizations\n[1] Zhen Qin, Xiaodong Han, Weixuan Sun, Bowen He, Dong Li, Dongxu Li, Yuchao Dai, Lingpeng Kong, and Yiran Zhong. Toeplitz neural network for sequence modeling. In The Eleventh International Conference on Learning Representations, 2023.\n\nWEAKNESSES:\n1. The main claim of the article is the speedup achieved by the proposed modifications. The only results supporting this claim are Fig. 1 and some percentages in the text (in section 5.1). Fig. 1 shows the performance on specific tasks from the LRA benchmark. Firstly, it is not clear for which task the baseline (TNN)  is evaluated. Secondly, the choice of the tasks shown on the graph is questionable. The hardest task from the LRA benchmark, Pathfinder-X, is not shown. As for the speedups mentioned in the text: it would be better to put them all into a separate table. Moreover, it would be interesting to see a speed comparison in the form of a table similar to Table 5 from the TNN article.\n2. In section 4.2, theoretical results on the choice of activations are presented. Several possible improvement ideas. The ablation study with experimental results for different activation types would be of interest. Moreover, the graphs showing the decay rate for randomly initialized networks might be improved. It would be better to leave only the lowest and highest lines and show the average line in between. Also, if you compare the rate of convergence to some baseline rate (e.g. exponential), plotting it would be appropriate. In addition, it seems to be not quite fair to compare the decay rates of trained and untrained networks.\n3. While the overall presentation style is to the point, as mentioned earlier, it would help the reader if the abstract, introduction and related work were more general. Both abstract and introduction may be hard to read for an unprepared reader, as they contain too much mathematical details and not enough motivation. Moreover, some paragraphs of the introduction repeat the abstract almost word for word, while rephrasing would help the reader to understand the ideas more deeply. The related work section should describe either the ideas of mentioned papers or their connection to your paper more clearly.\n4. The section 3.2.2 is a bit obscure. “Inverse time warp” is not a common term, and it is not described in the section.\n\nQUESTIONS:\n1. I did not understand from the paper how the training process for the sparse summand of the learned Toeplitz matrix is organized. Is it learned without RPE at all?\n\nLIMITATIONS:\nNone", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents several techniques to speed up Toeplitz neural networks (TNNs). In particular, TNNs use convolution of length n (sequence length of the input) and so scales as O(n log n), and TNNs have many calls to the MLP that generate relative positional encoding (RPE) and decay bias. To reduce the time of convolution, for bi-directional modeling the paper proposes to approximate the Toeplitz matrix as a sum of a short convolution and a low-rank matrix, which results in O(n + r log r) complexity where r is the rank of the approximation. For uni-directional model (e.g. auto-regressive modeling), the paper proposes to parameterize the convolution directly in frequency domain and uses the Hilbert transform to obtain the imaginary part from the real part of the filter to ensure causality. The approximation error is then analyzed. Validation on language model (Wikitext-103) and long-range benchmark (LRA) show that the approximation lead to some speedup (10-15%) and the quality stays around the same.\n\nSTRENGTHS:\n1. The idea of using asymmetric Nystrom to approximate the Toeplitz matrix is quite clever. This allows a decomposition into a sparse and a low-rank component, which leads to asymptotically faster algorithm in the case of bi-directional modeling.\n\n2. While uni-directional modeling prevents the Nystrom technique due to causal masking, parameterizing the filters directly in the frequency domain is able to overcome this challenge. While this is not asymptotically faster, it avoids one inverse FFT per layer and leads to some speedup.\n\nWEAKNESSES:\n1. Unclear what problem the paper is trying to address, and how it is motivated.\nThe intro starts out with Toeplitz neural networks, and the paper aims to make it faster. However, it's not clear why we want to make these faster, and what we would enable if we make these faster. Are they being used in very large-scale tasks? Are they being scaled to very long sequences?\nWhile the technical contributions are solid, it's not clear to me why the paper chose to tackle this problem.\n\n2. Unclear what the technical challenges are. \n- The paper mention that they want to avoid O(n log n) computation. But in practice O(n log n) isn't very slow, especially on GPUs. FFTs are pretty much bounded by memory bandwidth, and they take only 2-3 times as long as any pointwise operation. So if the goal is to speed up TNNs, then it makes more sense to have an efficient implementation, rather that using algorithms that faster asymptotically (O(n + r log r)) but is slower than a hardware-friendly algorithm (line 150, where using matmul with O(n r^2 + r log r) is faster). \n- The paper mentioned \"many calls to the RPE\". Why is this a problem? Showing a profile of how much each operation is taking will be much more convincing. That would motivate the approaches in the paper much better.\nWithout knowing how long each operations in TNNs are taking, how do we know that we're solving the right problem?\n\n3. Lack of detailed speed benchmark. Given the goal is to speed up TNNs, I would have expected one of the main results to be speed benchmarks, across different sequence lengths, on different devices, to show the tradeoff. In the main paper, speed is only reported in Figure 1b, which is end-to-end speed for a particular sequence length (512).\nHow do we know that we're close to the maximum speed on these devices (GPU)? Or are we still far from optimal? When we speed up convolution and RPE, what is the remaining bottlenecks.\nHaving these would make the paper stronger.\n\nAt sequence length 512, an optimized implementation of attention (e.g. FlashAttention) is likely faster than FFT and the method in this paper. This is my impression as the Hyena paper [1] reports that TNNs are not faster than FlashAttention until sequence length > 4k.\n\n[1] Hyena Hierarchy: Towards Larger Convolutional Language Models. Poli et al. 2023.\n\nQUESTIONS:\n3.2.2 I don't understand what's an inverse time warp. There doesn't seem to be an explanation of what that is in this section.\n\nLIMITATIONS:\nNot necessary.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes to reduce the computational complexity of Toeplitz neural networks. TNNs are a new form of network for sequence modeling that reduces space complexity of the attention matrix to allow for longer sequences.TNN model consists of a stack of Gated Toeplitz Units (GTU) that includes TNO (Toeplitz Neural Operator) that does token mixing with relative positioning. Then, GTU is a modified GLU layer injected with the proposed Toeplitz Neural Operator (TNO). \nThe paper addresses the TNNs efficiency limitations: 1) super-linear computational complexity 2) many calls to the RPE: for each layer, one call per relative position. Thus, the paper proposes to reduce both the complexity of sequence modeling and of the relative positional encoder. The work proposes solutions by means of both the Structured Kernel Interpolation (SKI) [2] and working with frequency domains.\nIt does so through:\n–\tApproximating Toeplitz matrix using low-rank approximation and replacing the RPE MLP with linear interpolation and using Structured Kernel Interpolation.\n(for O(n) complexity, that is use linear interpolation over a small set of inducing points to avoid the MLP entirely   -using an inverse time warp to handle extrapolation to time points not observed during training)\n–\tCausal training, SKI does not bring benefits, so instead they eliminate explicit decay bias by working in the frequency domain, using Hilbert transform (to force causality) and also use some smoothness\n–\t For the bidirectional case, they eliminate the FFT applied to the kernels.\n\nSTRENGTHS:\nThe work includes a number of solutions to improve TNN speed-up (addressing RPE MLP, the FFT, and the decay bias).\n\nRPE is a neural network to obtain relative position embedding to obtain entries in Toeplitz matrices. These entries could be evaluated with stationary non-SPD kernel which is a good idea.\n\nSo first decomposing Toeplitz matrix and then using interpolation for MLP is a comprehensive pipeline.\n\nThe theory part makes the arguments more sound, and the explanations in supplementary materials are fairly abundant.\n\nThe experiments on LRA show good predictive performance on long range data and on wikitext some speed-ups.\n\nWEAKNESSES:\nThe major paper of the paper talks about the SKI to accelerate the TNNs but in experiments SKI is only shown in the LRA experiment and does worse than both TNN and FD-TNN\n\nAs mentioned in the paper, doing sparse-dense multiplication in practice can be slower than dense-dense matrix multiplication (but that is only part of the potential speed-up)\n\nQUESTIONS:\nCan you elaborate on “We can interpret Tlij as evaluating a stationary non-SPD kernel kl (i j) = |i j| RPEl (i j). Thus Tl can be interpreted as a pseudo or generalized Gram matrix”\n\nWhat is the practical length of a sequence where the benefits are seen over TNN?\n\nA suggestion: some more elaboration in the main paper on causality aspect would be helpful, what is fast causal masking, what [33] says about ensuring causality instead of citing it. \n\nMinor: Provide better notation explanation, e.g. what is PPL, architectures, ListOps, Pathfinder, etc.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-28"}