Waveformer: Linear-Time Attention with Forward and Backward Wavelet Transform
Authors: Yufan Zhuang
Zihan Wang
Fangbo Tao
Jingbo Shang
Venue: Submitted to ICLR 2023 (2023)
Keywords: transformer
efficient attention
long range reasoning

Abstract: We propose Waveformer that learns attention mechanism in the wavelet coefficient space, requires only linear time complexity, and enjoys universal approximating power. Specifically, we first apply forward wavelet transform to project the input sequences to multi-resolution orthogonal wavelet bases, then conduct nonlinear transformations (in this case, a random feature kernel) in the wavelet coefficient space, and finally reconstruct the representation in input space via backward wavelet transform. We note that other non-linear transformations may be used, hence we name the learning paradigm Wavelet transformatIon for Sequence lEarning (WISE). We emphasize the importance of backward reconstruction in the WISE paradigm — without it, one would be mixing information from both the input space and coefficient space through skip connections, which shall not be considered as mathematically sound. Compared with Fourier transform in recent works, wavelet transform is more efficient in time complexity and better captures local and positional information; we further support this through our ablation studies. Extensive experiments on seven long-range understanding datasets from the Long Range Arena benchmark and code understanding tasks demonstrate that (1) Waveformer achieves competitive and even better accuracy than a number of state-of-the-art Transformer variants and (2) WISE can boost accuracies of various attention approximation methods without increasing the time complexity. These together showcase the superiority of learning attention in a wavelet coefficient space over the input space. 

Under review as a conference paper at ICLR 2023
WAVEFORMER: LINEAR-TIME ATTENTION WITH FOR- WARD AND BACKWARD WAVELET TRANSFORM
Anonymous authors Paper under double-blind review
ABSTRACT
We propose Waveformer that learns attention mechanism in the wavelet coefficient space, requires only linear time complexity, and enjoys universal approximating power. Specifically, we first apply forward wavelet transform to project the in- put sequences to multi-resolution orthogonal wavelet bases, then conduct non- linear transformations (in this case, random feature kernel) in the wavelet co- efficient space, and finally reconstruct the representation in input space via back- ward wavelet transform. We note that other non-linear transformations may be used, hence we name the learning paradigm Wavelet transformation for Sequence EEarning (WISE). We emphasize the importance of backward reconstruction in the WISE paradigm - without it, one would be mixing information from both the input space and coefficient space through skip-connections, which shall not be considered as mathematically sound. Compared with Fourier transform in re- cent works, wavelet transform more efficient in time complexity and better cap- tures local and positional information; we further support this through our ablation studies. Extensive experiments on seven long-range understanding datasets from the Long Range Arena benchmark and code understanding tasks demonstrate that (1) Waveformer achieves competitive and even better accuracy than a number state-of-the-e art Transformer variants and (2) WISE can boost accuracies of various attention approximation methods without increasing the time complexity. These together showcase the superiority of learning attention in wavelet coefficient space over the input space.
INTRODUCTION
Transformer [Vaswani et al. 2017) has become one of the most influential models in natural lan- guage processing (Devlin et al. 2018 Brown al., (020)), computer vision (Dosovitskiy et al, 2020), speech processing (Baevski et al. (020), code understanding (Chen et 2021a)) and many other applications. It is composed of the attention layer and the feed-forward layer with layer norms and skip-connections added in between. The original design of the attention layer scales quadrat- ically to the sequence length, becoming scalability bottleneck of Transformers as texts, images, speech, and codes can be of vast lengths.
Sfate-of-the-art attention approximation methods have enabled Transformers scale sub-quadratic or even linearly to the input sequence length. Typical approaches to computing a cheaper pseudo-
attention include sparse attention patterns (Parmar et al. 2018 Wang et al.) 2019 Bellagy et al.
2020 Zaheer et al., 2020), low-rank approximation (Wang et al. 2020 Chen et al. [2021b), and
kernel approximation (Katharopoulos et al.) 20200] [hhormaankii et al. 2020 Peng al., 2020),
where most of these methods are linear in time complexity. For comprehensive review, please
refer to Section 4
Recent works on improving the effectiveness and efficiency of long-range capabilities of Trans- formers start to explore attention learning in a transformed space. For example, conducting low- cost token-mixing with forward Fourier transform leads to remarkable accuracy improvement with a quasi-linear time complexity [Lee-Thorp et al. 2021). foken-mixing ideas (You et al. 2020 Lee-Thorp et al. 2021) are simple and effective, however, they lose Transformer's universal approx- imating power by replacing attention with hard averaging (Yun et al.) 2019). Moreover, without backward transform the model will mix information from both the input and transformed spaces,

---

Under review as a conference paper at ICLR 2023
Outpui
Forward Wavelet Projection x via ((t) Fast Wavelet Transform
Non-lincar Transform
Backward Wavelet Reconsiraction (3) via O(K) Fast Wavele Trandfonm
(eg. Cesture kemel)
Add & Nomualzz
F
Backward Wavelet
Kandom Feature Kernel
In the mevelel confficiinn space
Forward Wavelet
Input sequence (ength n)
Mulli-reshttion orihogonal wavelet bases
Multi-rcillutin orthogonal
wavelet Vaj
Ousput sequence (lougth
Teput
(Different colors denote different bases; 1:2-4 8-15 bases total)
(a) Waveformer Block
(b) Sequence I Leaming with Forward and Backward Waveler Transform (WISE)
Figure 1: An overview of our proposed Waveformer and WISE. (a) The only difference between Transformer block and a Waveformer block is the attention computation. (b) The general flow of computation in WISE with forward and backward wavelet transform.
which is not mathematically sound. Since multiplication in the Fourier coefficient space, after pro- jected back to the input space, is equivalent directly calculating convolutions in the input space, people have also utilized the forward and backward Fourier transform to learn large global filters with linear weights (Rao et al. 2021) and non-linearities (Guibas et al. 2021).
We propose Waveformer that facilitates the attention mechanism learning in a wavelet coefficient space, shown in Figure [a). It requires only linear time complexity and enjoys universal approx- imating power. Specifically, we first apply forward wavelet transform to project the input sequence multi-resolution orthogonal wavelet bases, then conduct non-linearity (c.g.. random feature ker- nel (Rahimi & Recht (007)) in the wavelet coefficient space, and finally, reconstruct the representa- tion in imput space via backward wavelet transform. We name this general learning paradigm WISE, as shown in Figure [b), bb) it can be suited with attention approximation methods to boost their long- range understanding capabilities. We implement wavelet transform using Fast Wavelet Transform (FWT) (Mallai) 1989) so both transform steps are linear in time. Intuitively, WISE operates on lo- cal to global, coarse to fine-grained cascading structure. Compared with Fourier transform, wavelet transform is more efficient in time complexity and better captures local and positional information since the wavelet basis localized in space with ranging granularity. For the non- non-linear transfor- mation in the wavelet coefficient space, attention approximation methods. A reason behind this that since wavelet transformation is invert- ible and exact, WISE will be universal approximating when coupled with universal approximators as its non-linearity.
apply any ineaaiitss,,,,, while we suggest using
We conduct extensive experiments the Long Range Arena (LRA) benchmark and common code understanding tasks to empirically ablate and justify this method. Compared with number of widely-used Transformer variants, Waveformer with a linear time complexity can achieve compet- itive and even better performance. When combined with various representative attention approx- imation functions, WISE can boost their performance without incurring extra time complexities. This shows that learning in wavelet coefficient space provides better long-range understanding capability over direct learning in the input space. Our ablation studies also support the use of the forward-backwant schema and the superiority of wavelet transform over Fourier transform.
In summary, our major contributions are as follows.
We propose WISE to facilitate learning in the wavelet coefficient space following a forward- backward paradigm which can be suited with attention approximation methods while boosting their long-range understanding capabilities.
Based on WISE, we develop Waveformer that requires only linear time complexity and enjoys universal approximating power for sequence- sequence-to-sequenee to-sequence functions.
Extensive experiments on the Long-Range Arena benchmark and code understanding tasks have demonstrated the effectiveness and also justified the design of Waveformer.
Reproducibility. We will release our code on GitHub.
WAVEFORMER
As shown in Figure a the only difference between Transformer block and a Waveformer block the attention computation. In this section, we introduce the details that replace the attention
2

---

Under review as a conference paper at ICLR 2023
computation in Waveformer. The general flow of WISE is shown in Figure T b), which constitutes the forward wavelet transform, the non-linearity in the middle, and the backward wavelet transform.
We list notations here denote scalars vectors x, matrices X; denote inner product between functions f, g as (f.g) = (())((t) dt; we denote function f's transformation in the coefficient space as f.
Background about Attention Let X E R denotes the input sequence of length and hidden dimension d. A dense self-attention layer shown below:
QK Vd
Attention( (X) Softmax
I
(1)
XW, with W WR. . Wi E Rdxm stand for the query, key, and value, respectively. The attention head size is denoted by m. The self-attention layer Attention()) computes a weighted average for each column V according to the dot-product similarity QK" that is of Omm)² number of operations.
where X K =
2.1 GENERALIZED FORWARD-BACKWARI PARADIGM
We propose the generalized forward- backward paradigm to conduct learning in the coefficient space between forward and backward transformation. WISE is special case of the generalized paradigm when the transformation is wavelet transform. The forward transformation, also called analysis, de- composes the input sequence into coefficients of set of orthogonal, complete functional basis. We then conduct non-linear transformation in the coefficient space, which directly operates on the func- tion. In the backward transformation, also called synthesis, reconstruct the target representation in the original function space.
We require the hewardd-bacwward transformation pair be invertible and exact, meaning that one can perfectly reconstruct the same input from the coefficients. Note that orthogonality is not neces- sary for this condition but non-orthogonal basis cannot be transformed via simple inner product.
The general framework is shown below. Without loss of generality, we limit ourselves to 1d func- tions. Given input and output function sx(),, y() R R time domain i, functional basis g t) on both time domain and frequency domain w, and non- non-linear transformation F.
Forward Transform: ((w) = (w,)) Σ
x(4:)g"(w, (w,4)
(2)
Nonlinear Transform: h(w) F f((w)
Backward Transform: v(() ()) = Σ (w
(4)
where (w,t) denotes the complex conjugate of g in case g has complex parts.
As concrete example, we use the Fourier transformation pair (formally defined in Appendix A.2
to illustrate this idea. Under Fourier transform, the functions are decomposed into sinusoidal waves, with g(w, t) = 2axtt The coefficients will represent the magnitude of the corresponding sinusoidal function. Hence, learning the mapping between any sequential function x<t) + y(t) becomes learning coefficient mapping between the transformed ((w)) - (ww) in vector space.
In WISE, we utilize the more localized wavelet transform as the barward-backwarr mechanism and in Waveformer we use the random feature kernel as the non-linear transformation. However we want to highlight that numerous alternative candidates fit in generalized paradigm, we briefly mechanism, any transformation based on orthogonal
list a few below. For the polynomials such Chebyshev transform and Hartley transform will suit purposes. For the non-linear transformation, essentially any neural architecture with sufficiently strong learning power will be valid candidate.
We will explain the inner workings of our Waveformer's 's -arkard- transformation and the reason for choosing it in the next part.
3

---

Under review as a conference paper at ICLR 2023
Chirp Signal from 1Hz 4Hz
Fourier Transform
Wavelet Transform
0.2
www
Frequency (Hz)
Magnitde
Amplitude
0.1
3
10
2
Time
Frequency (Hz)
Time
We show a chirp signal from IHz to 4Hz on the left, its Fourier transform the middle and its wavelet transform on the right. As you can see, from the Fourier transform graph one can only infer the existence of signal in the range of 1-4Hz without information on its happening time, while in the wavelet transform graph both time and frequency information are present and one can tell this is chirp signal.
Figure 2:
2.2 WAVELET TRANSFORM
Fourier transform decomposes the entire function into global sinusoidal waves, It tells people what frequencies are there and in what magnitude, but information is given about when that frequency started or ended. See Figure for an illustration on chirp signal. This limits the capability to understand the local structures of the input and conduct learning on top of it, which is crucial to many machine learning tasks.
Wavelet transform is designed to solve this issue. It employs a function (x)) x E R, called mother wavelet, to generate family of translated and dilated wavelets (see Figure Ib)):
(5) where scale i controls the resolution of the wavelet and j controls the position of the wavelet. With a larger the wavelet will be squeezed shorter in space, hence the normalization factor 2 to ensure the same 2 norm for all wavelets. The wavelet family (x)(()) is orthogonal on this dyadic grid.
(x)(()) 24 (2'x j), E
To be valid mother wavelet (x)), the only requirement is admissibility:
J
{0))) dr
(6)
In other words, the sum of function value should be O.
Given any square integrable function f E (.R( (i.e., J [1(()) dr < oo) and wavelet functions
Wiijs the wavelet transform pair is defined as:
Forward Discrete Wavelet: f(i.j) f(x))i = I(an)tt
(7)
400
Backward Discrete Wavelet f(z) Σ
Σ (11,)))1)(z)
(8)
00
where (1) denotes the complex conjugate of (x.)
Intuitively, in wavelet transform, we are scanning f(x with a microscope that has two knobs. One knob is the location j. the other one is the frequency (i.e., 2¹). We will be able oversee the local structure of the input and calibrate it accordingly with parameterized functions in WISE paradigm. See Figure for direct comparison.
To generalize beyond L²R and avoid using an infinite number of wavelets, we must introduce another function ø, called scaling function with a similar admissibility and orthogonality constraint:
[
+000 (xx) da dzg(z)) = 24 22 1): s.t ( =0, <'> Vj.j (
4

---

Under review as a conference paper at ICLR 2023
P is designed to cover the scale up to i, hence the orthogonality requirement. The decomposition
of f (() therefore becomes:
f(x) Σ (00.), f0 (x) +
Σ> i= j=-00
Σ
Ways (x))
(10)
Note that although still goes +00 in 10 usually has an upper limit in practice since impossible to work with infinite frequency. We also highlight that both the forward and backward discrete wavelet transform have an efficient (n) complexity algorithm (Mallat 1989).
In the d-dimensional case, we do not have general orthogonal discrete R wavelet, unlike the con- tinuous case. However, we can still perform discrete wavelet transform over each spatial dimension of the input, and we'd d still be able to perfectly project and reconstruct the original function.
To summarize, wavelet transform enjoys ((n) time complexity, an already desirable property com- pared to Fourier transform's On log n) complexity. It further provides the capability to examine the local structures with different resolutions via altering the scale. Learning carried out in this wavelet space will correspond to gathering and processing information from local to global level, in coarse to fine-grained fashion.
2.3 RANDOM FEATURE KERNEL
We have explained the transformation mechanism for WISE. One question remains unanswered: which architecture should choose use the parameterized non-linear transformation? For our purpose, the architecture should have linear complexity in sequence length; should be able to tackle sequential inputs effectively: should be as expressive as possible. Among the ones that fit our criteria, we choose random feature kernel (Rahimi & Recht 2007) due to its simplicity, efficiency, and theoretical benefits.
We have shown the self-attention in Equation ① It can be rewritten using the softmax kernel (SM): SM(Q,, Kj) SM(z. v) exp (x yv)
Softmax(QK
(11)
Random feature kernels aim to approximate SM(.,) :) with randomized mapping Rd R,.
K(x, v) E[v(x) v(y)] [v)]]
(12)
We use FAVOR+ from Performer (Choromanski et al. 2020) as the randomized mapping p:
cxp cxp (- 1==]1 (exp (w]x),... x). x)...., exp (w WA La) Vi (1....,m) (13)
v(xx Vm
where m denotes the number of random features. Similar FAVOR+, require ......... Wmn be orthogonal to reduce the variance of estimation. However, we impose additional normalization on beforehand to further stabilize the estimation inspired by Peng et al. (2020). We note that this in-place normalization is important when bounding the MSE error of the random feature kernel, as shown in Appendix A.1 Lemma
2.4 UNIVERSAL APPROXIMATION POWER
In this subsection, we show that Waveformer has the same universal approximation power for seq- to-seq functions as Transformer. The goal is to show that for any f in F, Vp E (o0), V > 0, we can find a f in the class of Waveformer, such that:
d,(f.f) - - - 11(X) - F(X) Se
We define the Waveformer class that has positional encoding, h heads, head size s, hidden dimension as Whe with FAVOR+ kernel described in Section 2.3
5

---

Under review as a conference paper at ICLR 2023
Table I: Evaluation results Long-Range Arena benchmark. We show both the average accuracy (Avg) and average accuracy without Retrieval (Avg (w/r) since LUNA 256, Nyströmformer, and our Waveformer all use prolonged 20k training steps on Retrieval task.
Model
ListOps
Text
Retrieval
Image
Pathfinder
Avg
Avg (w/r)
Transformer
36.37
64.27
57.46
42.44
71.40
54.39
53.62
Local Attention
15.82
52.98
53.39
41.46
66.63
46.06
44.22
Sparse Trans.
17.07
63.58
59.59
44.24
71.71
51.24
49.15
Longformer
35.63
62.85
56.89
42.22
69.71
53.46
52.60
Linformer
35.70
53.94
52.27
38.56
76.34
51.36
51.14
Reformer
37.27
56.10
53.40
38.07
68.50
50.67
49,99
Sinkhorn Trans.
33.67
61.20
53.83
41.23
67.45
51.39
50.89
Synthesizer
36.99
61.68
54.67
41.61
69.45
52.88
52.43
BigBird 36.05
64.02
59.29
40.83
74.87
55.01
53,94
Linear Trans.
16.13
65.90
53.09
42.34
75.30
50.55
49.92
Performer
18.01
65.40
53.82
42.77
77.05
51.41
50.81
FNet
35.33
65.11
59.61
38.67
77.80
55.30
54.23
LUNA 256
37.98
65.78
79.56
47.86
78.55
61.95
57.54
Nyströmformer
37.15
65.52
79.56
41.58
70.94
58.95
53.80
Waveformer
38.20
75.60
78.56
42.98
79.17
62.90
58.99
Theorem Vp E [1, +00, > O, and for any f E F, we can find a Waveformer network w Wa..1, such that d,(s, w) S.
The sketch of the proof is simple: since we have required the transformation pair to be invertible and exact, so for any seq-to-seq function, we can universally approximate in the wavelet space and is equivalent to having universal approximation power in the original space. The detailed proof of Theorem T shown appendix A.I
EXPERIMENTS
Besides intuitively and theoretically showing the superiority of the WISE paradigm, we also em- pirically show the accuracy our design when compared with other recent efficient transformers, along with several ablations on each component of our Waveformer.
3.1 EXPERIMENTAL DESIGN
To examine the performance of various efficient transformer variants, we compare them on datasets that require long-range understandings, (1) five datasets from the publicly available benchmark Long Range Arena, and (2) two code understanding datasets.
Long Range Arena (LRA) Tay al. (2020b) recent benchmark that designed to com- pare different efficient transformers. Since its release which already contains ten different efficient transformers, more and more efficient transformers have chosen it as the primary evaluation. The benchmark contains five datasets evaluate"] The datasets require understanding long sequences of mathematical operations, classifying text based on sentiment, matching similar documents, classify- ing images, and recognizing 2D spacial information. The sequence lengths of the dataset within the range of IK-4K.
Code Understanding Tasks Apart from the rather synthetic datasets in LRA (e.g.. flattening CIFAR-10 10 images to simulate long sequence lengths), we also select real-world long sequence task that involves detecting the vulnerabilities in source code. This task test Waveformer's ability to learn on source code and classify piece of source code as buggy or not. Two datasets are included: CodeXGL.UE (Lu al. 202T) introduced multi- protrammngg language benchmark for code
'There are also hard variants that even full attention transformers fail to solve, which is out-of-scope of this paper.

---

Under review as a conference paper at ICLR 2023
Table 2: Sequence length for each dataset's train split. The task of vulnerability detection requires reasoning over the entire piece of code snippet to deduct the label.
Table 3: Accuracy on vulnerability detec-
tion tasks. Evaluation on validation set.
Model
CodeXGLUE
D2A-Func
Dataset
CodeXGLUE
D2A-Func
Codo/C BERT
62.08
60.2
Average Length
1277.49
1038.86
BERT-medium
59.69
59.731
Median Length
561
717
Waveformer
62.81
62.587
model evaluation, one of the tasks is vulnerability detection on C code, and the labels are created by security experts; D2A (Zheng et al. 2021 P created for vulnerability detection with multi-modal trace and function data, we only use the function subset for evaluation. To show the long- range nature of these tasks, we show the average sequence length (in the number of tokens after tokenization) in Table 2
Experiment Environment. GPUs, and later moved to TPU v3-8s. Our code is written in Jax (Bradbury et al.. 2018) with the Flax framework (Heek et all 2020). The wavelet transformation implementation is primarily based on Jax Wavelet Toolbox (Moritz Wolter 2021 and PyWavelets (Lee al. 2019).
Our early-stage experiments are mostly conducted on RTX 3090
3.2 WAVEFORMER
LRA. We first highlight the performance of our method the LRA benchmark in Table No- lably, among all the attention approximation methods, our Waveformer performed the best on three of the five datasets, and with a close-to-top performance on another. The time complexity of Wave- former is also the theoretically lowest possible among all other approximation methods, matching the size of the input.
Code Understanding. To show that our Waveformer's ' long-range reasoning ability can also be applied to real-world tasks, we compare it with a full attention BERT model on solving the two code understanding tasks. For Waveformer and BERT, we use BERT-medium configuration (8 layers, 8 heads, 512 hidden dimensions) as the model size due to resource constraints.
We mostly follow existing works, Code- BERT (Feng et al. 2020) and C-BERT (Buratti et al. 2020), on the training and evaluation settings on the two datasets. Following Code-BERT's practice, we pre-train on CodeSearchNet (Husain et al. (019) corpus that contains 6.4M code snippets across six programming languages (Python, Java, JavaScript, PHP, Ruby, and Go) with the masked lan- guage modeling objective. We then finetune models the downstream vulnerability detection datasets. Both Code-BERT and C-BERT take in sequence up to 512 tokens, but the code snippets' median length already surpassed 512, meaning that more than half of the inputs will be truncated. Therefore we set the max sequence length as 1,024 for our Waveformer and BERT. We also note that there is more extensive training for Code-BERT and C-BERT that we could not afford, for in- stance, Code-BERT's pre-training uses batch size 2,048 over 100K training steps whereas our batch size set to 64, effectively times less training epochs, and both models choose a 12-layer 768- dimension configuration compared to our 8-layer 512- 12-dimension. The detailed hyperparameters and training configurations are included in Appendix A.4 We report their published test accuracy for CodeXGLUEE and D2A-Func respectively
strong baseline.
Table the performance of Waveformer and HERT-medium on the two code-understanding tasks. We can see that on a fairground of training and model size, Waveformer outperforms full- attention transformer model by a large margin. Further, when compared with larger-sized models that have dedicated training methodologies, Waveformer performs par the vulnerability de- tection task of CodeXGLUE, due to the power of abling to handle long-range sequences and under- standing Long-range relationships.
3.3 ATTENTION IN WISE
Our WISE paradigm has general philosophy of applying non-linearity in the wavelet transformed space and is not limited to a certain type of attention method. We comprehensively evaluate rep-
We note that D2A-Fune hasn't 't released the test set, hence report the validation accuracy for our imple- mentations instead.
7

---

Under review as a conference paper at ICLR 2023
Table 4: We use F/W to denote that the attention is performed in the Fourier/wavelet space (which also incurs an (n log (n) complexity cost). We reran Linformer & Linear Trans. for all (N/A, F, W) with the same additional five sets of hyperparameters because of convergence issues. We note that we are unable to reproduce score close to the original Linformer performance on Pathfinder. $ This the normalized version of Performer as described Section 2.3 with the wavelet space, it is our Waveformer.
when combined
ListOps
Retrieval
Image
Pathlinder
Attention
N/A
F
W
N/A
F
W
N/A
F
W
N/A
F
W
N/A
w
Full O(x²)
36.37
17.80
37.15
64.27
56.42
74.82
57.46
51.78
72.43
42.44
31.41
42.29
71.40
50.55
78.25
Linformer
((n)
35.70 36.15
37.65
53.94
57.06
55.22
52.27
55.93
65.85
38.47
34,89'
39.17
66,44'
61.76'
70.21
Linear Att.
O(x)
16.13 37.68
37.55
65.90
71.66
71.93
53.09
72.71
10.71
42.32*
51.07
40.83'
75,91*
70.45*
76.43*
Longformer
O(x)
35.63
18.95
36.65
62.85
55.36
74.99
56.89
52.52
66.21
42.22
29.12
37.10
49.71
50.38
78.15
Performer
0(n) 18.01
37.15
38.20
65 40
65.52
75.60
53.82
60.56
78.56
42.77
9.99
42.98
77.05
50.49
79.17
resentative attention methods on different space transformations (no transformation, Fourier trans- formation, and wavelet transformation). We show that performing full attention, or many other attention approximation operations in a wavelet transformed space as proposed in WISE paradigm almost always brings great accuracy improvements. In Table 4 almost all attention methods have increased accuracy when applied in the wavelet space compared to an untransformed space, except for the Image dataset, where some methods incur slight drop in accuracy. When compared with Fourier transformation, we can see that in most cases, the wavelet transformation much better.
We do note an interesting finding that Linear Trans. when coupled with Fourier transformation space, showed pretty good results on all the tasks, and especially a significant improvement on the Image [ One hypothesis is that this is due to the fact that Linear Trans.'s 's polynomial kernel in Fourier space, which occurs in the forms of K ((x, y) = (x y), can be interpreted as self-convolution in input space. We leave this for future discovery.
4 RELATED WORK
4.1 ATTENTION APPROXIMATION METHODS
There has been plenty of prior work to enable transformers handle long input more efficiently and effective. Since the inefficiency comes from the quadratic dependency on sequence length because of the dense attention operation, large portion of research simulates the attention operation with certain approximations, for example, replacing the dense attention matrix with a sparse version, or assume that it satisfies certain low-rank structures. We briefly review some methods on this topic in this section. For a more detailed survey, we refer the readers to
Tay et al. (2020c).
Sparse Attention. Perhaps the most intuitive solution to alleviate the quadratic cost, Sparse At- tention only calculates a portion of the full n² attention matrix. Farly stage methods include Local Attention (Parmar et al. 2018) and or chunked blocks to speed up computation. Longformer (Beltagy et al. 2020) and BigBird (Za-) heer et al. 2020) further combine global attention, sliding window attention, dilated sliding window attention, and random attention together to form strong sparse attention mechanisms, and BigBird showed that their method is a universal approximator of sequence functions. To make the block truncation a learnable process, Reformer (Kitaev et al. 2019) groups and sorts input segments via locality-sensitive hashing such that similar tokens are placed in the same chunk. Similarly, Sinkhorn Transformer (Tay et al. (2020a) trains a meta sorting network to reorganize input sequences before applying windowed attention.
BERT (Wang et al.) 2019) use sliding windows
The self-attention matrix, at the center of transformer, has been found to display low-rank behaviors after pre-training. Linformer (Wang et al.) 2020) performed spectrum analysis on the trained attention matrix, and the results indicate that the top 128 singu- lar values composite 88%-96% of the entire 512 singular values across attention heads and layers. Based on this observation, Linformer added low-rank projection matrices in attention to approxi- mate the original attention matrix. On a similar notion, Drone (Chen et al. 2021b) extended the low-rank approximation scope to all matrices in transformer via data-driven optimal compression.
Low-rank Approximation.
Note that the hyperparameters were tuned for Linear Trans, on this task, see Appendix A.3
8

---

Under review as a conference paper at ICLR 2023
Kernel Methods. The kernel methods approximate the whole self-attention by replacing the softmax with a kernel function that can be decomposed to avoid the explicit calculation of the O(n²) matrix multiplication. Linear Transformer (Katharopoulos et al.) 2020) proposed a non-negative elu feature mapping as the substitution for the softmax, they further pointed out the connection between their formulation and RNNs, and argued that transformers and RNNs be unified under the same umbrella. Building on top of this, Random Feature Attention (Peng et al., (020) and Performer (Choromanski et al. 2020) utilized random feature approximation of the attention, one highlights the importance of normalization before random projection while the other one emphasizes the benefits of positive orthogonal random features.
Token Mixing. Token Mixing methods are another version of efficient transformer building
blocks. Different from the methods discussed above, they do not approximate attention, but rather conduct new way of enabling communication between tokens. You et al. (2020) showed the pos- sibility that a random token mixing strategy can work well in transformer encoders, as opposed to delicate (pre-)trained attention heads. Token Mixing is new view towards self-attention methods are not approximating self-attention. Lee-Thorp et al. (202T) pushed this idea further by providing an efficient method to mix the tokens with Fourier forward transformation.
Among these methods, our Waveformer utilizes a wavelet transform, thus, is slightly similar to Token Mixing. However, work should be seen as new approach to efficient transformers, which mixes the idea of a orthogonal space transform that communicates between tokens and attention approximations methods that benefit in the new space. In our study, also pick representatives from each of the of attention approximation types and show that the dense attention operation and these sparse attention operations can benefit from learning in wavelet transformed space.
4.2 FOURIER & WAVELET TRANSFORM IN ML
Fast Fourier Transform (Cooley Tukey, 1965) has been widely used many machine learning domains, probably the most common usage is to speed up the computation of convolution. For similar purpose to our work, AFNO (Guibas et al. (022)) learns global filters for images via adding block-wise MLP in Fourier transformation, equivalent to convolutional layer with large filters. However, AFNO is designed for visual inputs, we show in our ablation study that such architecture cannot fully capture long-range text information.
Fast Wavelet Transform (Mallat, 1989) has been the backbone of numerous modern technologics including JPEG-2000 image compression (Skodras et al.) 2001), digital communication (Akansu & Smith 1993) and many others. FWT has been used for computation speed up (Wolter et al. 2020), speech recognition (Tufekci & Gowdy] 20000, and time series analysis (Michau al. 2022). Recently in computer vision, WaveMix (Jeevan & Sethi, (022) proposed to mix the input images with forward wavelet transform. We note that our work differs from theirs by learning the non-linear transformation in the coefficient space amid forward and backward wavelet transform.
CONCLUSIONS
In this paper, we propose to learn sequential mapping in the wavelet coefficient space, Specifically, the inputs are first forward transformed into the wavelet space, then the sequential mapping is learned and finally, we reconstruct the transformed sequence back in the input space. We combine this paradigm with random feature kernels and propose Waveformer, that has universal approximation power and linear time complexity. When coupled with attention approximation methods, WISE can boost their performance on long-range understanding tasks while enjoying no extra cost time complexity. Our experiments support the superiority of Waveformer and obtain strong performance on the LRA benchmark and code understanding datasets.
Through this work we have focused on performing attention in the transformed space, but how will other seq- seq-to-seq models work in Fourier or wavelet space remains unknown. We also note that there exist hundreds of wavelet bases, each with different design purposes and properties. Findings crucial properties for conducting effective learning in the coefficient space and constructing adaptive wavelet basis accordingly both interesting problems we leave for future work.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Ali N Akansu and Mark JT Smith. Subband and wavelet transforms: design and applications, volume 340. Springer Science & Business Media, 1995.
Alexei Baevski, Yuhao Zhou, Abdelrahman Mohamed, and Michael Auli, wav2vec 2.0: A frame- work for self-supervised learning of speech representations. Advances in Neural Information Processing Systems, 33: 13:12449-12400 2020.
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.0555555 2020.
James Bradbury. Roy Frostig, Peter Hawkins. Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Minn -Milee,, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http: //gi wwt. cccom//ooole/jaaa
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.
Luca Buratti, Saurabh Pujar, Mihaela Bornea, Scott McCarley, Yunhui Zheng, Gactano Rossiello, Alessandro Morari, Jim Laredo, Veronika Thost, Yufan Zhuang, et al. Exploring software natu- ralness through neural language models. arXiv preprint arXiv:2006. /2641, 2020.
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. Evaluating large language models trained on code. arXiv preprint arXiv:2107.0333 2021a.
Patrick Chen, Hsiang-Fu Yu, Inderjit Dhillon, and Cho-Jui Hsich. Drone: Data-aware low-rank compression for large nlp models. Advances in neural information processing systems, 34:29321- 29334, 2021b.
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, 2020.
James W Cooley and John W Tukey, An algorithm for the machine calculation of complex fourier scries. Mathematics of computation, 29900)) 1965.
Ingrid Daubechies. Ten lectures on wavelets. SIAM, 1992.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:/810.0488 2018.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x 16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010. 11929, 2020.
Zhangyin Feng, Daya Guo, Duyu Tang, Nan Duan, Xiaocheng Feng, Ming Gong, Linjun Shou, Bing Qin, Ting Liu, Daxin Jiang, al. Codebert: pre-trained model for programming and natural languages. In Findings of the Association for Computational Linguistics: EMNLP 2020,
pp. 1536-1547, 2020.
John Guibas, Morteza Mardani, Zongyi Li, Andrew Tao, Anima Anandkumar, and Bryan Catan- zaro. Adaptive fourier neural operators: Efficient token mixers for transformers. arXiv preprint arXiv:2111.13577, 2021.
Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Mare van Zee. Flax: A neural network library and ecosystem for JAX, 2020. URL
http:://gtte com/ goog le/ flax
10

---

Under review as a conference paper at ICLR 2023
Hamel Husain, Ho-Hsiang Wu, Tiferet Gazit, Miltiadis Allamanis, and Marc Brockschmidt.
Codesearchnet challenge: Evaluating the state of semantic code search. arXiv preprint 19XX:::909.00436. 2019.
Pranav Jeevan and Amit Sethi. Wavemix: Resource-efficient token mixing for images. arXiv preprint arXiv:2203.03689, 2022.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mns: Fast autoregressive transformers with linear attention. In International Conference Ma- chine Learning, pp. 5156-5165. PMLR, 2020.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In International Conference Learning Representations, 2019.
Gregory R. Lee. Ralf Gommers, Filip Waselewski, Kai Wohlfahrt, and Aaron O' O'Leary. Pywavelets:
A python package for wavelet analysis. Journal of Open Source Software, 4(36):1237, 2019. doi: 10.21 011//sss.02233 URL URLtttps::: /ttp://doooo// org/
001377
James Thorp, Joshua Ainslie, Ilya Eckstein, and Santiago Ontanon. Fnet: Mixing tokens with fourier transforms. arXiv preprint S.Ki::2005.03824 2021.
Shuai Lu, Daya Guo, Shuo Ren, Junjie Huang, Alexey Svyatkovskiy. Ambrosio Blanco, Colin Clement, Dawn Drain, Daxin Jiang, Duyu Tang, et al. Codexglue: machine learning benchmark dataset for code understanding and generation. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 1), 2021.
Stephane G Mallat. A theory for multiresolution signal decomposition: the wavelet representation. IEEE transactions on pattern analysis and machine intelligence, (17):674-693, 1989.
Gabriel Michau, Gaetan Frusque, and Olga Fink. Fully learnable deep wavelet transform for un- supervised monitoring of high-frequency time series. Proceedings the National Academy of Sciences, 199():22106598 1199, 2022.
Moritz Wolter. Frequency Domain Methods in Recurrent Neural Networks for Sequential Data Processing. PhD thesis, Rheinische
wwwph--ihhhhhlln-neeeeee Bonn, July 2021. URL https://hdd www.haandee... 500.
Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. In International conference on machine learning, pp. 4055- 4064. PMLR, 2018.
Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah Smith, and Lingpeng Kong. Random feature attention. In International Conference on Learning Representations, 2020.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. OpenAl blog, 2019.
Ali Rahimi and Benjamin Recht. Random features for large-scale kernel machines. Advances in neural information processing systems, 20, 2007.
Yongming Rao, Wenliang Zhao, Zheng Zhu, Jiwen Lu, and Jie Zhou. Global filter networks for image classification. Advances in Neural Information Processing Systems, 34:980-993, 2021.
Athanassios Skodras, Charilaos Christopoulos, and Touradj Ebrahimi. The jpeg 2000 still image compression standard. IEEE Signal processing magazine, 1865):3655, 2001.
Yi Tay, Dara Bahri, Liu Yang, Donald Metzler, and Da-Cheng Juan. Sparse sinkhorn attention. In International Conference on Machine Learning, pp. 9438-9988888 -447. PMLR, 2020a.
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, 2020b.
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. ACM Computing Surveys (CSUR), 2020c.
11

---

Under review as a conference paper at ICLR 2023
Z Tufekci and John N Gowdy. Feature extraction using discrete wavelet transform for speech recog- nition. In Proceedings of the IEEE SoutheastCon 2000. Preparing for The New Millennium (Cat. No. 00CH37105), pp. 116-123. IEEE, 2000.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- sion processing systems, 30, 2017.
Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04766 2020.
Zhiguo Wang, Patrick Ng, Xiaofei Ma. Ramesh Nallapati, and Bing Xiang. Multi-passage bert: A globally normalized bert model for open- question answering. In Proceedings of the 2019 Conference Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-JJCNIP), pp- 5878-5882, 2019.
Moritz Wolter, Shaohui Lin, and Angela Yao. Neural network compression via learnable wavelet transforms. In International Conference on Artificial Neural Networks, pp. 39-51. Springer, 2020.
Weiqiu You, Simeng Sun, and Mohit lyyer. Hard-coded gaussian attention for neural machine translation. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 7689-7700, 2020.
Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are transformers universal approximators of geuuence-to-squemee functions? In International Con- ference on Learning Representations, 2019.
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, 13:17283-72297, 2020.
Yunhui Zheng, Saurabh Pujar, Burn Lewis, Luca Buratti, Edward Epstein, Bo Yang, Jim Laredo, Alessandro Morari, and Zhong Su. D2a: dataset built for ai-based vulnerability detection meth- ods using differential analysis. In 2021 JEEF/ACM 43rd International Conference on Softwre Engineering: Software Engineering in Practice (ICSE-SEIP), pp. 111-121 IEEE, 2021.
A APPENDIX
A.1 PROOF FOR THEOREM 1
We define the function class F to be the set of all countinous functions that map compact domain in x to xx
We start from making the connection between random feature kernel and regular transformer block:
Lemma (Asymptotic Result for FAVOR+) The following is true for independent random wi.
MSE(SM(I, v)) exp (| + v1 "SM (1, (-Ila= +
m
MSE(SM(Z, v)) + SM(2.9)
lim
where SM denotes the softmax kernel, SM denotes the random feature kernel, and MSE stands for mean-squared error.
The proof of this lemma be found at Choromanski al. (200 Lemma 2). It tells us the the MSE error is upper bounded to a constant since x, y is normal ized beforehand, and vanishes to 0 as the original softmax kernel value tends to 0 and the number of random features tends to +0c.
Next we use the main theorem of Yun et al. (211)). We denote the transformer network class that has positional encoding, h heads, head size 8, and hidden dimension Th.
12

---

Under review as a conference paper at ICLR 2023
Hyperparameter
Configa
Configa
Configa
Config4
Configs
Layers
2
2
2
Embedding Dim.
128
128
128
256
256
Attention Dim.
64
64
64
MLPP Dim.
128
128
256
1024
512
Attention Heads
8
8
2
4
4
Dropout
0.2
0.1
0.1
0.1
0.2
Attention Dropout
0.1
Table 5: Additional hyperparameter configurations tried for Linformer and Linear Trans. in Image and Pathfinder
Lemma Vp E (00,, e <> 0, and for any f F, we can find Transformer network g T21,1 such that dp(f.g) S.
The proof of Lemma constitutes of several steps, of which the first step is to approximate any function f E F as piece- wise constant function f. Since f is continuous, the piece-wise constant approximation can be of arbitrary accuracy. Next they find a modified transformer 9 with hardmax operator and special class of activations. Finally they show that the transformer block g is able to approximate The functional distance is then bounded by:
+ dp(S.)) d
dp(ū.)) Se
We show that with slight modification, the proof will work for Waveformer, and can be generalized to the WISE paradigm under certain constraints.
The proof outlined below:
For Vf F, its wavelet transform (we will also use for to denote this, see for details) is still continuous. Hence, the discretization claim remains valid. We can then effectively approximate the self-attention transformer block with the FAVOR+ block up to á difference by controlling the number of random features m. In the end, the backward reconstruction exact, the distance bound becomes when we control the other three terms to be less than ze as well:
dp(f.w) S dplfwi + dp + dr(ā.9) dolg.w) Se
A.2 FOURIER TRANSFORM PAIR
Given function f(x), x E I, the Fourier transform pair is defined as the following:
Forward Fourier : f(w) = / f(x)e i2nwe dx
(14)
Backward Fourier f(x) = 1 f(w)e dw
(15)
where w stands for frequency in Fourier space.
A.3 LRA CONFIGURATION DETAILS
We tried to follow all hyperparameters suggested for each of the attention approximations with exceptions on Linformer and Linear Trans. in Image and Pathfinder. For them, we experimented with five additional configurations shown in
For all wavelet transform conducted in this work, we use Daubechies Daubechies 1992) as the basis and we set level of decomposition 1.
For Waveformer, the number random features in random feature kernel is set as 256 for all text tasks (ListOps, Text, Retrivial), 512 for all image tasks (Image, Pathfinder).
13

---

We conduct ablation study for Waveformer, shown in Table 6 For (Linear), We limit the transformation in wavelet space to be linear. For (Fourier), we use the Fourier transform as the transformation mechanism for Waveformer. For (Forward Fourier), we only the forward Fourier transform without backward transform. It can be observed that performance dropped significantly in wavelet Under review as a conference paper at ICLR 2023
Table Ablation study on Long-Range Arena benchmark.
Model
ListOps
Text
Retrieval
Image
Pathfinder
Avg
Avg (w/r)
Waveformer
38.20
75.60
78.56
42.98
79.17
62.90
58.99
Linear 37.70
55.36
55.27
15.75
50.58
42.93
39.84
Fourier
36.85
65.52
60.56
9.99
50.49
44.68
40.71
Forward Fourier
37.15
64.91
63.98
37.84
53.39
51.85
48.32
A.4 CODE PRE-TRAINING & FINETUNING DETAILS
For pre-training, we use the CodeSearchNet dataset retrieved from Huggingface dataset. We use GPT-2 (Radford al. 2019) tokenizer for Waveformer and BERT since source code does have fixed vocabulary. We believe byte-level tokenization is more reasonable than sub-word tokenization.
We use the following set of pepeepprraettrs in pre-training: batch size is 64, learning rate is 5e- 5, masking probability is 0.15, warm steps is 10k, training steps is 100k with a linear learning schedule; we use AdamW optimizer with 31 0.9, 82 = 0.98, weight decay is 0.1.
We use the following set of hyper- ppper-parameters in finetuning: batch size is 64, learning rate is 2e-5, training length is 10 epochs; we use AdamW optimizer with 31 0.9, A2 0.999.
A.5 ABLATION STUDY
all cases, indicating the necessity of lineariyy in wavelet space and transform.
14


Reviews:
[{"review_id": "jy3sWYmQL_", "reviewer": "Reviewer_jCxJ", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "In the paper, the author proposed a new Transformer variant with linear complexity, which is called Waveformer. In each Waveformer layer, the input will first project to the coefficient space using a forward discrete wavelet transform. Then a linearized attention operation is applied via random features. Finally, the outputs are transformed back to the \"contextual\" space using a backward discrete wavelet transformer step. The authors studied the expressiveness of the architecture and conducted experimental analysis of the model on standard benchmarks.", "summary_of_the_review": "I believe proposing efficient Transformer is essential but the current quality of the work (motivation, empirical results) are not ready to publish in the venue, and I recommend rejection.", "strength_and_weaknesses": "Strength:\n\n1. The architecture is new as far as I know and studying the efficiency of the model is an important direction.\n2. It is useful to show the model is a universal approximator, which is the same as the original Transformer model.\n\nWeakness:\n\n[Novelty]\n1. Where is the linear complexity coming from?\n\nIf I understand correctly, the operation that reduces the quadratic complexity in attention is to use the random feature kernel developed in Choromanski et at., 2020, but not the proposed wavelet transform. \n\nIn the first round of reading, I got confused as in several parts of the paper, the authors emphasize that the linear-complexity model is a major contribution. I believe the authors have overclaimed that.\n\n2. Why you use wavelet transform?\n\nIn the work, the input will first project to the coefficient space using a forward discrete wavelet transform. Then a linearized attention operation is applied. Note that even if you don;t use the wavelet transform, the model is already a linear transformer with random features. So is the wavelet transform redundant and what is the benefit?\n\n[Universal approximation]\n\nOnce you use the random features, your model is, in fact, a \"randomized\" function. How did you deal with approximating a deterministic function by a \"randomized function\"? I didn't find any place in the proof coping with that. That being said, the proof is incomplete and lacks sufficient details. I don't think you can simply let the random feature dimension approach infinity and say, \"MSE error vanishes to 0\"， since we are concern about a finite-size Transformer's approximation ability\n\n[Experiments]\n\nThe author missed important effcient Transformer S4 (\"Efficiently Modeling Long Sequences with Structured State Spaces\"). In all of the LRA tasks, S4 is roughly better than the proposed model for 20 points. I have little reason to vote a paper for acceptance given the significant performance gap.\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "atwQ5aBi-Uo", "reviewer": "Reviewer_KfzV", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors propose a wavelet transformation based liner attention mechanism named Waveformer. Experiments on long-range arena benchmarks have been conducted to demonstrate the effectiveness of the proposed method.", "summary_of_the_review": "My major concern is the novelty of the proposed method. Applying spectral-based transformation (Fourier or wavelet) for efficient computation has been widely adopted in ML (also indicated in Sec. 4.2 by the authors). The claim of non-linear transformation as a major contribution seems to be incremental. I suggest the authors to better justify the novelty of the proposed approach and the motivation of applying wavelet transformation.\n\nFor experiments, I suggest the authors explicitly indicate how many repetitions each model has been trained and also report the variance of accuracy. The authors choose random feature kernel in the paper and I wonder if the authors have considered other kernels or how different kernels may affect the model performance.", "strength_and_weaknesses": "Strengths:\n* The authors focussed on an important issues in sequence modeling\n* The paper is organized well and easy to follow\n* Technical details are presented clearly\n\nWeakness:\n* The novelty needs to be further justified\n* The evaluation needs to be strengthen", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "EYHjwSq1c1", "reviewer": "Reviewer_DaL7", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper \"WAVEFORMER: LINEAR-TIME ATTENTION WITH FORWARD AND BACKWARD WAVELET TRANSFORM\" proposes integrating analysis- and synthesis-wavelet transforms into transformers.\n\nThe main idea is to move attention computations into the wavelet domain, by framing the attention computation with a forward and backward wavelet transform. The attention module uses the randomized mapping known from the performer architecture.\nThe proposed network is evaluated on the Long Range Arena and the Code Understanding benchmarks.\n\nThe authors observe competitive or improved performance.", "summary_of_the_review": "This paper presents an interesting application of fast wavelet transforms in attention modules.\nThe results are competitive. The experimental part of the paper leaves many questions open.\nWavelet choice and the number of scales for example are not investigated.\nMean, and variance are unknown for all experiments. Experimentally the paper presents the bare minimum.\nI am recommending a weak reject for now, but am willing to adjust my rating depending on the rebuttal.\n\nEdit: In response to the rebuttal, I have raised my rating from five to six.", "strength_and_weaknesses": "### Strengths:\n- To the best of my knowledge the fast wavelet transform has not yet been studied in attention modules.\n- The theory is adequately discussed.\n- The reported results are competitive.\n\n### Weaknesses:\n- The novelty of the proposed approach is limited. Wavelet transforms have previously been studied in other contexts.\n- For a paper which studies the introduction of wavelet transforms, it would have been interesting to study the effect of the wavelet itself.\n- It seems different wavelets have not been explored.\n- Similarly, the depth of the wavelet transforms has not been studied. It would have been interesting to see if and how the level choice impacts the results. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Dc6xvHiJSXy", "reviewer": "Reviewer_ciFm", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper innovates attention block by sandwiching it into wavelet forward and backward transformations. Empirical results suggest that this approach can improve upon vanilla softmax attention, as well as other approximated and more efficient attention mechanisms.", "summary_of_the_review": "Although an in-depth analysis and theoretical justification is lacking, the paper presents a simple method and promising results, which is novel and interesting.", "strength_and_weaknesses": "The method is simple and interesting. Empirical results (especially Tab. 4) are encouraging. It encourages reader to think \"\nwhy attention in (wavelet) transformation domain can outperform the conventional one in time domain?\".\n\nHowever, the answer is still unknown after reading the paper. An in-depth analysis is lacking. Section 2.2 discusses the advantage of wavelet coefficients against Fourier coefficients. The authors suggests maintaining both time and frequency resolution is helpful. But that is a signal processing perspective, whose connection to representations learning is not clear. In that sense, is short-time Fourier transform also able to outperform Fourier transform? Theorem 1 is a combination of two existing theoretical results. It still cannot provide an answer to the question above.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject