Transformer-VQ: Linear-Time Transformers via Vector Quantization
Authors: Lucas Dax Lingle
Venue: ICLR 2024 poster (2024)
Keywords: Transformer
Transformer Decoder
Decoder-Only Transformer
Natural Language Processing
NLP
Vector Quantization
VQ
K-Means
Clustering
Causal Attention
Autoregressive Attention
Efficient Attention
Linear-Time Attention
Autoregressive Modeling
Generative Modeling
Gated Attention
Compressive Attention
Kernelized Attention
Kernelizable Attention
Hierarchical Attention
Segment-Level Recurrent Attention
Long-Context Modeling
Long-Range Modeling
Long-Range Dependencies
Long-Term Dependencies
Cached Attention
Shift-Equivariant Attention
Primary area: generative models
TL;DR: Simple, efficient, and stable decoder-only attention via vector quantization

Abstract: We introduce Transformer-VQ, a decoder-only transformer computing softmax-based dense self-attention in linear time.  Transformer-VQ's efficient attention is enabled by vector-quantized keys and a novel caching mechanism. 
In our large-scale experiments, Transformer-VQ is shown highly competitive in quality, obtaining 0.99 bpb on Enwik8, 26.6 ppl on PG-19, and 3.16 bpb on ImageNet64. In addition, the optimized implementation of Transformer-VQ is over 3x faster than a comparable quadratic-time transformer at sequence length 8k, is over 12x faster at 32k, and can scale to 131k with similar throughput. Code available: \url{https://github.com/transformer-vq/transformer_vq}

Published as a conference paper at ICLR 2024
TRANSFORMER-VQ, LINEAR-TIME TRANSFORMERS VIA VECTOR QUANTIZATION
Lucas D. Lingle Independent Researcher
com
ABSTRACT
We introduce Transformer- VQ, a decoder-only transformer computing softmax- based dense self-attente in linear time. Transformer-Vo VQ'ss efficient attention is enabled by vector-quantized keys and a novel caching mechanism. In our large-scale experiments, Transformer-V VQ is shown highly competitive in quality, obtaining 0.99 bpb on Enwik8, 26.6 ppl on PG-19, and 3.16 bpb on ImageNet64. In addition, the optimized implementation of Transformer-VO than a comparable quadratic-ting transformer at sequence length 8k, is over 12x faster 32k, and can scale 131k with similar throughput. Code available: https: github..hh....o com rans fo
over 3x faster
vQ
k
kz
kg
ki
h
ha
lis
Figure 1: Schematic of the VQ-Altention approximation. The colorful and blank boxes depict the keys and attention weights, respectively. The keys on the right have been rector-quantized Since the green keys k2, ky map to the same code, they have the same attention weights in this attention head.
INTRODUCTION
Transformer (Vaswani et al. 2017) language models would ideally scale to long sequences, since their predictive abilities often improve as context length increases (Dai et al. 2019 Kaplan et al.) 2020). Unfortunately, the standard transformer uses self-attention mechanism with quadratic time complexity with respect to sequence length. This limits the practicality of applying transformers to very long sequences, since increasing the sequence length by factor of 10" increases the attention computations by a factor of 100". Transformer variants that overcome this efficiency bottleneck have the potential facilitate
long-context applications and enable breakthroughs.
Up to this point, a variety of efficient transformers (Tay et al. 2020b) have been proposed to scale
long sequences. Techniques include sparsity (Child et 2019; Ye et al. 2019 Beltagy al.
2020
Kitaev et al. [2020) Qiu et al.) (2020) Roy et al. 2021 Tay et al. 2020al Sukhbaatar et al.
2021
Wu al. [2022] Liu
[2033] Zhang et (202), compression (Liu
2018 Rae et al.
2020
Ainslie et al. 2020 Zhu et al. [221]] Ren et al. [202]] Nawrot et al. 202
2023), low-rank
approximations (Wang al. 2020 Vas al. 20200] Katharopoulos et al. 2020 Xiong al. 202T
Tay et al. 2021 Choromanski et al. 2021), and cross-attention operations (Dai et al. 2019: Ma et al.
2021 2021:Hutchins Hutchins et al. 2022 Hawthorne al. (022). Other efficient sequence models have also been
proposed (Gu al. 2022 Lee- Lee-Thorp al. 2022 Mchta et al. 2022 Smith et al. 2022 Hasami et al.)
2022
Poli et al. 2023 Peng et 2023).
In this paper, present Transformer- VQ, transformer decoder with dense self-attention computible in linear time with respect to sequence length. This is made possible through a combination of vector-
quantized keys, localized positional biases, and compressive cache that can be attended to efficiently, while yielding the same results as an uncompressed variable-length cache. Transformer VQ is also simple to implement sampling for.

---

Published as a conference paper at ICLR 2024
2
PRELIMINARIES
2.1 NOTATION
The real numbers are denoted by R and the extended real numbers R U {- oo indices are used for all tensors. When indexing a matrix M along the first axis, we use M, to denote column vector and M,; : to denote row vector. The functions LN(-), Softmax(-)) Concat(-) denote LayerNorm (Ba et al. 2016), softmax, and concatenation, each applied row-wise. The symbols A, x, O exp(.), 6. SG(.) denote equality by definition, proportionality, element-wise product, element-wise ex exponentiation, Kronecker delta function, and the stop-gradiess operator. We slightly abuse notation to write inner products of vectors u, V as u v. and outer products as uv
by R. Zero-based
We assume familiarity with transformers (Vaswani et al. 2017). and use the notation Dm denote the model width, D& to denote the query/key vector width, and D, to denote the value vector width.
2.2 VECTOR QUANTIZATION
Vector quantization (VQ) a technique used extensively throughout this work. In this subsection we briefly review vector quantization, motivate its use in self-attention, and discuss the backpropagation- compatible VQ scheme introduced by van den Oord al. (2017).
2.3 VECTOR QUANTIZERS AND CODEBOOKS
Definition 2.1. A vector quantizer is a function VQ(:; C) with domain R and codomain RR,,, For an input x, its output x is given by
arg arg min |x
C2||2
8 C
(2)
where C E RSDD is known as the codebook. The row indices ( 13 of C are called shortcodes, and the rows themselves are called codewords.
Theorem 2.2 (Based on Guo et al. (2119). Letq E IR. be a random variable with Eg [qq T = I'ID
for some a > O, and let E R2d be a random variable independent of q. Let ip : R + IR be
deterministic funtiioo Then
Eq,k|1q7 -
(k)]l (kkk)"((k]]
x Exilk - s(k))l??
(3)
Corollary 2.3. Let the conditions of Theorem
2.2 hold. Given the constraint that (p(R. > {C,)
1-11
the choice (p-) VQ(:; C) minimizes kg q q
(k)) 1
Corollary 2.4. Let the conditions of Theorem
2.2 hold. With k vQ(k; C) we have
arg min Eq k q k||² = arg min Exl k - k|l2. C
(4)
Remark 2.5. Fnding the global minimizer C" arg ming EA. k k 2 is expensive, in practice
we approximate it using the method from van den Oord et al. 2017): Razavi et al. (2019).
2.4 SECTOR-QUANTERN REPRESENTATION LEARNING
Definition 2.6 (Based on van den Oord et al. (017 ). A veclor-quantizer with straight-through
estimator is a function STVQ(:; C) with domain R and codomain RD_ For an input x, its output & given by
A arg min |x C2
(5)
SG(C=
Remark 2.7. For any x E R.D, STVQ(x; C) evaluates to VQ(x; C). However, for purposes of backpropagation, the Jacobian of the quantizer w.r.t. its input will now be identity matrix everywhere, instead of zero matrix almost everywhere. Intuitively, when using STVQ, gradients w.r.t. the quantizer outputs copied 'straight through' the inputs. Remark 2.8. We overload the notation STVQ(:; C) to operate row-wise matrix-valued inputs.
2

---

Published as a conference paper at ICLR 2024
3 TRANSFORMER-VO
We now propose Transformer- VQ. a decoder-only transformer that can compute dense self-attention in linear time. Proofs for all theoretical results given in Appendix A
3.1 QUADRATIC- TIME FORMULATION
Definition 3.1. Vector-Quantized Self-Artention is a function VQAut(;; C, ((Q.k.va..o))) with domain R Txxxx Div and codomain R Tx Dm, For an input X, its output Y is defined via
X e LN(X) Tx P...
(7)
_ r - ERTxD
(8)
K A 7 ()(
(9)
V 2 (XWV)) TTD.
(10)
G A b (XWc) TR.N.
(11)
K A STVQ(K; C) ERTxD,
(12)
W A di (QK' + +B) ERTX
(13)
O (WV) O G ETTD.
(14)
Y Y X X RTxD...
(15)
where each W. denotes a trainable projection, B denotes positional biases and/or mask, 7 is fixed constant, and the 6e, &y. &w are element-wise or row-wise nonlinearities. The query/key LayerNorms unit gain and zero bias, and STVQ(:: C) denotes row- wise application of vector-quantization with a thhgghh gradient estimator (van den Oord et al.) 2017).
gated attention unit (GAU) design inspired by Hua et al. (2022). GAU is a single-headed gated attention mechanism and generally uses a small key width D3 128, and large value width D., 2Dm with two GAUs replacing a single transformer layer. This yields a similar parameter count and compute requirement as the usual transformer layer. Remark 3.3. Prior work has also applied LayerNorm or similar to the queries and keys in attention (Henry et 2020 Roy 2021 Zhu et al. 2021 Wu et al. 2022 Hutchins et al. [202 Dehghani
Remark 3.2. Our attention mechanism is applied
et al. 2023 Elsen et al. 2023 generally finding it to improve numerical stability and convergence.
3.2
WARMUP: LINEAR-TINA ENCODER ATTENTION
To simplify the theorems for decoder-only attention and build intuition, we first discuss a setting where there is causal mask.
Theorem 3.4. Suppose B, = 0 for all i. and ow is an element-wise nonlinearity. Then the can can be factored:
attention weights in Definition 3.1
W A Sw &
(QK) +B)
(16)
Si. (QKT
(17)
dw (QC
(18)
where ó (QCT) TXXS A RSXT and A.,t it function and 2, is the VQ shortcode for timestep t.
A 69.2. Here, 8.,. denotes the Kronecker delta
Theorem 3.5. Suppose Bt, for all i, and dw is the row-wise softmax nonlinearity Then the attention weights in Definition 3.1 3.7 can be factored:
W &m (QKT + +B)
(19)
= P.. (QKT
(20)
Diag(exp( (QC A1) (CP(QQT)) )
(21)
where 1 e RT. Diag(exp( QCT) e QC E TXS. A SXT and A & 6,,.. Here, 5.,. denotes the Kronecker delta function and 2, is the VQ shortcode for timestep t.
3

---

Published as a conference paper at ICLR 2024
X
ow ( (QK T) E RTXT
V RTXD.
(QC RRS AV Rgxx.
Figure 2: Schematic of the VQ-Attention factorization with element-wise ow. The column set of W Φ (QKT) RTxT has size < S due to VQ, so the attention output O WV can be obtained by computing the unique attention scores @w (QC and using them to further aggregate to the grouped-sum AV. Transformer- VQ uses softmax-bast extension of this idea for its cache.
3.3 LINEAR-TIME DECODER ATTENTION
Theorem 3.6. Let L be a divisor of Suppose Bij -00 for j i (causal masking), and = = for < L (no bias outside a sliding window). Define A E RSXT with A. - 6g.z1- Let @w be element-wise nonlinearity with (- = 0. For tensor M. let (......... denote the slice M... ....aII((nn11)) I)........ where unsliced dimensions will be denoted by ':' Then the product WV in Definition 3.1 be computed using the following block-level recurrence:
U - 0
((a) V (m.)
i
U(n)
&
(22)
otherwise
(WV)
ci (Qn,,)
)n 2)
(23)
+ Gi Q (m.)) KT
(n-1,)) + B BIn.n-1) V
(24)
+ + (Q (1,5) (n.:) + (B. ((n.)
(25)
where any tensor slice (........ is defined as a zero tensor of width L in the sliced dimension if any block slice index less than zero (zero-padding).
Theorem 3.7.
Let the assumptions of Theorem 3.6 hold, but suppose @w is now the row-wise softmax
nonlinearity, Let 1 E B² Let A & exp(QK + B). Then the product WV in Definition 3.1 can be computed using the following block-level recurrence:
U(n A (s)) (m.)) 0
in otherwise
Unn)
(26)
[L(n- 1) + (in) 1()
if 2 0 otherwise
L(n)
(27)
(AV) exp(Q (m.i) C Un
(28)
+ exxp (n.:) K (n-1.) + B (oan--))((nn-))
(m.))
(29)
+exp (n.:) K (n.:) + (B(nn)) (n.)
(30)
(A1) = exp(Q ) 1-2-2 -2) T
(31)
+exp (n.:) (n- 1,)) + +B B)n.n-)))((-11)
(32)
+ +exp Q (m.:)
(n.)) + +B (B(mn))
(33)
(WV) (n.) Diag((A1) (w)) ('AV)
(34)
Remark 3.8. Theorem 3.7 provides algorithm compute VQ-Attention from the queries, keys, values, gates, and codebook in OLL(S + 2D) + +D))) time per query block, and therefore O(T(S + (D) D.)) time per sequence.
Remark 3.9. For numerical stability, we use an equivalent implementation Theorem 3.7 that stores the running mean of the value vectors assigned given shortcode, instead of the sum as done by U(n - 2). The result is made equivalent by moving the logarithm of the counts L(n 2) inside the exponentials (x.)) CT) appearing in (A.) and (A (A)) See pseudocode in Appendix
4

---

Published as a conference paper at ICLR 2024
Remark 3.10. The general strategy of computing un-normalized softmax and its denominator is also used by many prior methods, including Memory-Efficient Efficient Attention (Rabe & Staats [002)) FlashAttention (Dao et al. 2022), and RWKV (Peng et al. 2023); however, the first two techniques do not run in linear time, and the last one couples a recurrent state size to the model width, which is contrary to the principle of transformers.
3.4 LEARNING ALGORITHM
3.4.1 TRAINING LOsS
Let 0 denote the set of non-codebook parameters of a transformer with N VQ-Attention layers, and let C = {C(() N-1 denote the set of the layers' codebooks. For autoregressive modeling of
1==0
sequence X = X 1 we define the Transformer-VO VQ training loss as
L(X; (xe(X; C) 3Cvo
(35)
where B > is hyperparameter known as the commit loss coefficient, and
T
LcE(X;0,C) &
(36)
T
N-1
1 Lvo 0x(00 &
Σ (=0
K SGG
(37)
T
Thus, the training loss the average next-token cross-entropy loss, plus the average token's com- mitment losses (van den Oord et al., (017), summed over layer codebooks. The non-codebook parameters receive a gradient from both loss terms. Following van den Oord et al. 2017
Razavi
al. 2019), codebooks are parameterized using EMA-smoothed k-means.
3.4.2 TRAINING UPDATES
Instead of updating on the full sequence loss given above, we generally update every W/1 query blocks, where W < T, which resembles a strategy used in prior works (Dai et al. 2019 Wu et al.. 2022
Hutchins et al. 2022).
Each update is obtained by backpropagating through a window of W timesteps, with gradients computed on the corresponding terms in the per-token average losses above. Codebooks are also updated every W/L query blocks.
When W/L 1, using Theorem 3.7 is an efficient equivalent to variable-length key-value cache. When W/L > 1, a learning signal is sent through any value vectors added to the compressed cache within the backpropagation window.
4 RELATED WORK
4.1 HIERARCHICAL ATTENTION
Combiner (Ren al. 2021) proposes an approximation softmax using simple graphical model, and parameterizes its internal probabilities using max-pooling over query/key features, enabling decoder-only self-attention in subquadratic time. H-Transformer-ID (Zhu & Soricut 202T uses average- geerage-poooing operations over queries/keys to reduce the complexity of encoder-only self-attention. Transformer-LS (Zhu al. 202T) uses dynamic projections downsample long-range features in transformers by user-specified factor. Hourglass Transformer (Nawrot et al. (02T) and MegaByte (Yu et al. 2023) eschew pooling favor of convolutions or reshaping for temporal downsampling, and apply these techniques to reduce computation in the interior layers of decoder-only transformers.
Transformer-VO differs from these works that it uses vector quantization (VQ), a well-understood method for compression, instead of newly-designed heuristic methods. In addition, it does not rely on token contiguity to guide the compression process. Instead, it utilizes an equivalence to dense attention. Notably, Transformer- VQ is easier to sample from compared previous hierarchical attention models; since the cache update logic can be equivalently applied every token instead of every I. tokens, there sporadic 'feature consolidation' operations required during sampling.
5

---

Published as a conference paper at ICLR 2024
4.2 KERNELIZABLE ATTENTION
Kernelizable attention (Katharopoulos et al. 2020; Choromanski et al. 2021 Peng et al. 2021
Qin al., 2022b) computes query and key features and applies the same nonlinearity both of them separately, omitting additional nonlinearities when computing attention weights. By using the associativity of matrix multiplication, kernelized attention reduces attention to linear complexity. Transformer- VQ is distinguished from kernelizable attention through an asymmetric treatment of queries and keys, deterministic equivalence to softmax-basse attention, training stability, and strong quantitative results on long-context autoregressive modeling benchmarks.
Clustering attention (Vyas et al. 2020) uses vector-quantized queries and is also kernelizable. However, it requires learning par-layer codebooks for each sequence and uses modified form of Lloyd's iterations based on Hamming distance and locality-sensitive hashing. This yields a complex non-causal algorithm which is only suitable for non-causal attention and is slow TPUs. Transformer- VQ is strongly differentiated from clustering attention by its simplicity, applicability to decoder-only tasks, efficiency on TPUs, and large-scale experimental validation.
4.3 COMPRESSIVE ATTENTION
Compressive Transformers (Rae et al. 2020) directly learn a compression function for long-range features. LUNA (Ma et al. 2021) and Recurrent Transformers (Bulatov et al.] 2022 Hutchins/ ct al., (022) use cross-altention to compress long-range features into recurrent state. Notably, our model implements a kind of block-recurrent mechanism for its cache, but is significantly more parameter-efficient than the mechanisms proposed by Ma al. (202T); Hutchins al. (2022). More generally, Transformer- VQ differs from compressiverrecureett transformers in that it has an equivalence to quadratic-time attention over vector-quantized keys. In other words, if the keys are already vector-quantized, the Transformer-VO cache losslessly reduces the cost to linear time.
Perceivers (Jaegle et al., 2021] Hawthorne et al., 2022) use cross- antesnatttntion to attend to long sequences, and compute self-attention over only a narrow stack of 'latents'. Transformer- VQ differs from Perceivers that it computes dense self-attention in linear time, instead of just cross-attention. Thus, while Perceivers' long-range layers incur quadratic time complexity during sampling, Transformer- VQ generates sequences in linear time.
4.4 GATED SEQUENCE MODELS
Gated attention was introduced in FLASH (Hua et al. 2022 as a fusion of attention sublayers (Vaswani et al. 2017) and GLU-based MLP sublayers (Shazeer 2020). Various gating mechanisms
have previously been used to stabilize training of transformers (Parisotto et al., 2019) and other
sequence models including $4 (Gu al., 2022), GSS (Mehta et al. 2022). MEGA (Ma al. 2023)
and RWKV (Peng et al.) 2023). Transtormer- VQ uses the original gating formulation from Hua et al. (2022), and develops new attention mechanism.
4.5
VQ. K-MEANS, AND BEYOND
Ideas relating k-means, vector quantization, and/or codebooks have also been applied in transform- ers for sparse attention (Roy et al. 2021 Wang et al. [2021] (022), feature learning (Mao et al. 2022 Roy et al. 2022), sparsely-activated MLPs (Lample al. 2019 and expert selection (Roller et al.
2021). These works generally feature codebooks or similar within transformer architecture. Several works also have proposed models that feature codebook somewhere outside a transformer, e.g.. when transformers are priors for VQ-VAEs (Kaiser et al.) 2018 Dhariwal et al., (2000] Ramesh et al. 202T Lee et 2022 Zhou et al., 2022). Transformer- VQ uses one codebook within each layer and, contrast to all of the aforementioned works, computes dense self-attention in linear time.
Transformer- VQ is not directly related to methods which quantize the weights of a transformer e.g.. Dettmers et al. (2022); Dettmers & Zettlemoyer (2023); Frantar et al. (2023). Such methods are typically applied after training to reduce the memory overhead of the model weights, while still computing in higher precision. As such, they do not affect the bitwidth of the queries, keys, or values, nor the complexity of self-attention. However, if applying our method to large models, these approaches may be complementary during inference.

---

Published as a conference paper at ICLR 2024
5 EXPERIMENTS
Transformer- VQ is implemented in Jax (Bradbury et al. 2018) and Flax (Heek et al. 2023
). For
training, we TPU v3 pod slices (Jouppi et al. (2017). Hyperparameters follow Appendix unless specifically mentioned. Generated samples for all models are provided in Appendix
5.1
PRELIMINARY STUDIES
5.1.1 CODEBOOK SIZE ABLATIONS
Larger codebook sizes may allow more flexible attention patterns and could improve the fidelity of the gradients, both of which are likely to benefit model quality at the expense of additional wall time. To investigate, we ablate the codebook size S using the Enwik8 dataset (described in & $5.2.11. 1) and report the lowest validation bits- bits-per-byte (BPB, lower is better) obtained by each model in Table T
Table 1: Codebook size ablations.
Table 2: Compressive cache ablation.
Setting
Val. BPB
Latency (Rel.)
Compressive cache
Val. BPB
Lateney (Rel.)
$ 256
1.010
0.927
No Yes
1.026 1.010
0.836
S 512
1.005
1.0
0.927
S 1024
1.000
1.109
Table confirms the intuition that larger codebooks improve the prediction quality (lower BPB) in return for additional wall time per training step. In particular, for this dataset and model size, increasing the codebook size by a factor of two appears to decrease the validation BPB by about factor of 0.995. This result may suggest that the validation loss follows power-law scaling (Kaplan jet al., (2000) w.r.t. codebook size, though more experiments are needed to verify this phenomenon, and it is subject to the caveat that the validation loss must eventually level off (Henighan et al. 2020 Hoffmann et al. 2022), as the model cannot be expected to obtain zero loss at infinite codebook size.
5.1.2 COMPRESSIVE CACHE ABLATION
Since our model has several architectural differences from most prior works, the benefit of the compressive cache must be shown directly. To investigate, we train a model with the compressive cache omitted, using codebook size 5 256. We report the validation BPB for Enwik8 in Table
As shown in Table 2 removing the compressive cache reduces the wall time per step by factor of about 1.1 at the evaluated model size, but leads to significant drop in quality (higher bits-per-byte). This confirms the importance of our compressive cache mechanism.
5.1.3 LATENCY AND THROUGHPUT
We now measure the training latency (seconds per step) and compute the training throughput (tokens per second). The latter computed as tokens per batch divided by latency, and allows a direct efficiency comparison across different sequence lengths. We benchmark on TPU v3 with cores, using a global batch size sequences. For these experiments, we scale the sequence length T by multiples of 4x, and backpropagate through the entire sequence length.
We compare an unquantized quadratic- cinttaatii-ieeeeeee full attention baseline ('Full') to our proposed linear-tinn time VQ-Attention ('VQ') using Theorem 3.7] Since this theorem does not require access to the output gates, VQ-Attention can be extended to multi-head attention variants as well. For each attention type, therefore benchmark three head types: multi-head attention ("MHA'; Vaswani al. (2017)) multi-query attention ('MQA"; Shazeer (2019)), and single-head gated attention ( ('SHGA' aka GAU; use codebook size S 512 and block length L 512, which is the same as our later experiments. All models use roughly 190M parameters total.
Hua et al. (2022) ) For VQ-
As shown in Table 6 our model has 3x lower latency/3x higher throughput than the quadratic attention baseline at T = 8192 when using SHGA for both. Moreover, Transformer VQ is over 6x faster than the quadratic-time baselines when using MQA/MHA for both models. As the sequence
7

---

Published as a conference paper at ICLR 2024
length increases to T = 32768. Transformer- VQ is over 12x faster than the quadratic time baseline when both use SHGA. For MQA/MHA, the quadratic- attention gives an out-of-memory error at T = 32768, while Transformer- VQ maintains comparable better throughput than with 4x shorter sequences. Table even shows that Transformer- can scale to sequences of length T = 131072 without substantial decrease in throughput and without running out of memory.
5.2 QUANTITATIVE RESULTS
To assess the ability of Transformer- VQ to learn long-range dependencies, we now conduct a series of large-scale experiments, benchmarking on several long-range autoregressive modeling tasks. For fair comparison, we only benchmark against models (a) trained without using any extra data or augmentation, and (b) evaluated with fixed parameters. all cases, we use codebook size S 512.
5.2. ENWIK8
Table 3: Test bits-per-byte on Enwik8.
Enwik8 is a byte-level language modeling dataset con- sisting of 100 million bytes of unprocessed English- language Wikipedia articles (Mahoney) 2011). with long-term dependencies that may span tens of thou- sands of bytes. Per convention, it is split into train, validation, and test sets of 90 million, 5 million, and 5 million bytes, respectively (Child al. 2019 Rae et al. 2020).
Model
BPB
Dai et al. (2019)
XL
0.99
Child et al. (2019) Sparse
0.99
Bellagy
(2020)
Longform.
0.99
Roy ct al. (202T)
Routing
0.99
Sukhbaatar ct al. (2019) Adapt.
0.98
Nawrot ct al. (222)
Hourglass
0.98
For this dataset, we trained Transformer-V VQ with 190M parameters, smaller than the model by Dai et al. (2019). We report test bits-per-byte (BPB) in Table
Rae el al.
2020
Compress.
0.97
Zhu el al.
202T
Long-Short
0.97
Fan al.
20206b
Feedback
0.96
Lei [221
SRD++
0.95
Transformer-Vo VQ obtains a BPB of 0.99, notably match- ing the result of the large Transformer-XL model from Dai et al. (2019), while using 33% fewer parameters and 75% shorter cache that covers a longer context.
Sukhbaalär et al.
2022
Expire.
0.95
Lutati al. (2023
Focus Attn.
0.94
Transformer- VVQ
0.99
For this dataset, we found overfitting was a significant issue, and due to the compressive cache mechanism, using i.i.d. attention dropout was not possible. Sweeping over the residual dropout rate, weight decay coefficient, and layerdrop (Fan et al., (020a) rate, we found a setting yielding good generalization. Nonetheless Transformer- VQ does fall short of state-of-the-art here, with several works using complex recurrence forgetting mechanisms and obtaining better Enwik8 results.
5.2.2 PG-19
PG-19 is an open-vocabulary language modeling dataset consisting of 11 gigabytes of text from over 28,000 freely-available Project Gutenberg books published prior to 1919 (Rae al., 2020). The average number of words per book is nearly 70,000, enabling learning long-term dependencies, especially novels (Sun al. 202 Hutchins et al. 2022).
For this dataset, we trained Transformer-VQ with 1.3B parameters, similar to the largest model by Hutchins et al. (2022). Since PG-19 an open-vocabulary dataset, we first learned SentencePiece vocabulary (Kudo & Richardson 2018) of size 32,000 using the BPE method. Following the calculations of Rae et al. (2020), we report the test set word-level perplexity (WLP) in Table
Transformer- VQ obtains a WLP of 26.6. very close to the state-of-the-art by Transformers (Hutchins al. 2022 Interestingly, since our Transformer design is equivalent to using dense self-attention with vector-quantized keys, our strong result shows that models using self- attention only (no recurrence) can also be highly competitive on PG-19. This affirms the efficacy of standalone self-attention method for sequence processing scale. Furthermore, compared to the Block- Recurrent Transformer, our model can be implemented via intra- intra-block sums and cross-blocd reductions, strategy also used by FLASH (Hua et al. (022) and shown to be faster in Appendix B Lastly, we avoid the instabilities of FLASH QQn et al. 2022a Ma et al. 2023) thanks to softmax normalization and our cache normalization (8 3.9
8

---

Published as a conference paper at ICLR 2024
Table Validation bits-per-byte on ImageNet64.
Table 4: Test word-level perplexity on PG-19.
Model
BPB
Model
WLP
Kingma et al
2021
VDM
3.40
Yu et al.
(2023
MegaByte
36.4
Hawthorne al.
2022
Perceiver AR
3.40
Rae el al
2020
XL
36.3
at.
2022
McgaByte
3.40
Rae et al.
2020
Compressive
33.6
Grete ct al. al.
DenseFlow
3.35
Koy al.
2021
Routing
33.2
Lipman el al.
2023
Flow Matching
3.31
Hawthorne
2022
Pereeiver AR
28.9
Hazami al.
2022
Efficient VDVAE
3.30
Hutchins et al.
2022
Block-Recur. Recur.
26.5
Transformer- VQ (190M)
3.22
Transformer- VQ
26.6
Transformer- VQ (1.2B)
3.16
5.2.3 IMAGENET64
ImageNet64 is an image dataset consisting of over 1.2 million images downsampled to 64x64 resolution (Chrabaszez et al. 2017 Deng et al. 2009). Flattening the images yields an autoregressive density estimation task on sequences of over 12,000 bytes each. Note since the official test set is not public for this dataset, we report results on the official validation set. For validation purposes we used held-out set of about 80,000 examples from the training split.
For this dataset, we trained Transformer- VQ models with 190M and 1.2B parameters, similar to the Enwik8 and PG-19 models, respectively. We report the bits-per-byle on the official validation set in T Several of the earlier baselines used an earlier variant of downsampled ImageNet prepared by van den Oord et al. (2016) with different downsampling algorithm. Since that variant has been unavailable through official channels for about a year, we used the newer variant following Lipman al. (2023). We emphasize that our results using the newer variant cannot be directly compared with basclines using the earlier variant; however, due to several reporting ambiguities, Table does not symbolically distinguish variants used.
Figure 3: Generated samples from our large ImageNet64 model; nucleus 1.0.
Transformer- VQ with 190M parameters is roughly the same size as the Efficient VDVAE Hazami et al. 2022), but obtains better result of 3.22 BPB, setting new state-of-the-art for small models. Transtormer- with 1.2B parameters obtains a 3.16 BPB, setting a new absolute state-of-the-ar on this dataset, and generates high-fidelity samples on par with Perceiver AR while using 33% fewer steps, omitting its image-specific architectural adjustments, and generating samples in linear time,
CONCLUSION
Transformer-Ve VQ a transformer decoder computing softmax-bart self-attention linear time. Its efficient attention is enabled by vector-quantized keys, which allow cache to be attended to in compressed form, while yielding the same result as uncompressed attention over the same keys. Large-scul scale experiments show Transformer- VQ is efficient and flexible autoregressive model, with state-of-fe of-the-art results or near on PG-19 and ImageNet64. Future work directions include formal scaling laws, larger models, and porting to lower-level frameworks like Pallas, Triton, CUDA.

---

Published as a conference paper at ICLR 2024
REPRODUCIBILITY STATEMENT
To facilitate reproducibility, our attention mechanism is described mathematically in Section 3| and pseudocode provided in Appendix F addition, our hyperparameters and other implementation details are given in Appendix C and our implementation is open-sourced at the link in the abstract.
ACKNOWLEDGMENTS
We are grateful to the anonymous reviewers for their helpful feedback. In addition, we would like to express our gratitude to the Python community, especially the Jax ecosystem contributors, for the effective libraries used in this project. This project was generously supported with Cloud TPUs from Google's ' TPU Research Cloud (TRC).
REFERENCES
Joshua Ainslie, Santiago Ontanon, Chris Alberti, Vaclav Cvicek, Zachary Fisher, Philip Pham, Anirudh Ravula, Sumit Sanghai, Qifan Wang, and Li Yang. ETC: Encoding long and struc-
tured inputs in transformers. In Proceedings of the 2020 Conference on Empirical Methods
in Natural Language Processing (EMNLP), pp. 268-284, Online, November 2020. Associ-
ation for Computational Linguistics. doi: 10. 10.18653/w1/2020ccmm
URL https
//aclanthology wttploggyyrrr/////
emnlp mranpp-main...
Jimmy Lei Ba, Jamie Kiros, and Geoffrey Hinton. Layer normalization, 2016. URL https
1/arxiv. org. atth/bba//16600 160.//64500000
Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer. CoRR, abs/2004.05150, 2020. /R../ttps::/azxiv.rr/////
05150
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milh and Qiao Zhang. JAX: composable transformations of Python+NumPP programs, 2018. URL
fitt:://githbe wwwpgogg.///cccocooooooooco jax
Aydar Bulatov, Yuri Kuratov, and Mikhail Burtsev. Recurrent memory transformer. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL https::
Rewon Child, Scott Gray, Alec Radford, and llya Sutskever. Generating long sequences with sparse transformers. CoRR, abs/ 105099 2019. ULLL/EE1:///axxx..0///////////
Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreca Gane, Tamás Sarlós, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy J Colwell, and Adrian Weller. Rethinking attention with performers. In International Conference Learning Representations, 2021. URL https: //openreview. net // 1d
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Tsvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Lev- skaya, Sanjay Ghemawat, Sunipa Dev, Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito, David Luan, Hyeontaek Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thanumalayan Sankaranarayang Pillai, Marie Pellat, Aitor Lewkowycz, Erica Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang, Brennan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy
Douglas Eck, Jeff Dean, Slav Petrov, and Noah Fiedel. Palm: Scaling language modeling with pathways, 2022. URL https::/errs www..orgg/eee abs /204 02311
10

---

Published as a conference paper at ICLR 2024
Patryk Chrabaszez, Ilya Loshchilov, and Frank Hutter. A downsampled variant of imagenet as alternative to the CIFAR datasets. CoRR, abs/1707.088// 19, 2017. httt::tttatxxiiii
abs/1700 08819
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc Le, and Ruslan Salakhutdinov. Transformer-Xi Attentive language models beyond a fixed-length context. In Proceedings of the 57th Anmual Meeting of the Association for Computational Linguistics, pp. 1978-2988, Florence, Italy, jul 2019. Association for Computational Linguistics. doi: 10. 10.18653/v1/P19-11 https:
URL
1285
Tri Dao, Daniel Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattennion: Fast and memory-efficient exact attention with io-awareness, 2022. URL https::/arxiv..o 2205.14133
Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski, Jonathan Heek, Justin Gilmer, Andreas Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, Rodolphe Jenatton, Lucas Beyer, Michael Tschannen, Anurag Amab, Xiao Wang. Carlos Riquelme, Matthias Minderer, Joan Puigcerver, Utku Evci, Manoj Kumar, Sjoerd van Steenkiste, Gamaleldin F. Elsayed, Aravindh Mahendran, Fisher Yu, Avital Oliver, Fantine Huot, Jasmijn Bastings, Mark Patrick Collier, Alexey Gritsenko, Vighnesh Birodkar, Cristina Vasconcelos, Yi Tay, Thomas Mensink, Alexander Kolesnikov, Filip Pavetić. Dustin Tran, Thomas Kipf, Mario Lučié, Xiaohua Zhai, Daniel Keysers, Jeremiah Harmsen, and Neil Houlsby. Scaling vision transformers to 22 billion parameters, 2023. URL http:: org/
05442
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: large-scale hier- archical image database. In 2009 IEEE Conference on Computer Vision and Pattern Recognition, pp. 248-255, 2009. doi: (0.1100092...2
Tim Dettmers and Luke Zettlemoyer. The case for 4-bit precision: k-bit inference scaling laws, 2023. URLhttp: http::/tttxiii...//// org htt///bb/222112 http::///ax...oorrr//////22222. 09720 09720
Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. GPT3.int8() 8-bit matrix multiplication for transformers at scale. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL
/ttps://openzett ott.://www........oo.oo.oooodoo.
Prafulla Dhariwal, Heewoo Jun, Christine McLeavey Paine, Jong Wook Kim, Alec Radford, and Ilya http///aaxxi...r//////////
Sutskever. Jukebox: A generative model for music, 2020. URL 2005.0030
Stefan Elfwing. Eiji Uchibe, and Kenji Doya. Weggoid- linear units for neural network function approximation in reinforcement learning. CoRR, abs/1 100.033 18, 2017. URL http:
//axxxi.org//bs//1111 03118
Erich Elsen, Augustus Odena, Maxwell Nye, Sagnak Taşırlar, Tri Dao, Curtis Hawthorne, Deepak Moparthi, and Arushi Somani. Releasing persimmon-8b, 2023. URL httplttts::/// http::://wwwwaaa adept:
ht/blog/persimenn
Angela Fan, Edouard Grave, and Armand Joulin. Reducing transformer depth on demand with structured dropout. In International Conference Learning Representations, 2020a. URL
htttss://openrev/
net wttp://oormmo03yyy022200
Angela Fan, Thibaut Lavril, Edouard Grave, Armand Joulin, and Sainbayar Sukhbaatar. Addressing some limitations of transformers with feedback memory, 2020b. URL Utt/attpp:://r/..oooog abs/ / 2002.00000000 09402]
Elias Frantar, Saleh Ashkboos, Torsten Hoefler, and Dan Alistarh. OPTQ: Accurate quantization for generative pre- pre-trained transformers. In The Eleventh International Conference on Learning Representations, 2023. wtt.httppee/ppaarrreeeeeeeeee/ee

---

Published as a conference paper at ICLR 2024
Matej Greie, Ivan Grubisic, and Sinisa Segvic. Densely connected normalizing flows. In M. Ranzato, A. Beygelzimer, Y. Dauphin. P.S. Liang, and J. Wortman Vaughan (eds.), Ad- vances in Neural Information Processing Systems, volume 34, pp. 13996-229966 Curran Asso- ciates, Inc., 2021.
fi hatppperrffll//pppapaaeee les/paper/
Albert Gu, Karan Goel, and Christopher Re. Efficiently modeling long sequences with structured state spaces. In International Conference on Learning Representations, 2022. URL https
7 //openreview, net/ wttpeeenneeoooooeeeeeeeeee ottp://ooommmmdooooooooo/oo/
Ruiqi Guo, Quan Geng, David Simcha, Felix Chern, Sanjiv Kumar, and Xiang Wu. New loss functions for fast maximum inner product search. CoRR, abs/1908. 2019. URL http:
7/arxiv. org/ abs/ 908. 10396
Ramin Hasani, Mathias Lechner, Tsun-Hsuan Wang, Makram Chahine, Alexander Amini, and Daniela Rus. Liquid structural state-space models, 2022. URL cnttps://arxitt org/abs
2209. 12951
Curtis Hawthorne, Andrew Jaegle, Cätälina Cangea, Sebastian Borgeaud, Charlie Nash, Ma- teusz Malinowski, Sander Dieleman, Oriol Vinyals, Matthew Botvinick, Ian Simon, Hannah Sheahan, Neil Zeghidour, Jean-Baptiss Baptiste Alayrac, Joao Carreira, and Jesse Engel. General- purpose, long-contess autoregressive modeling with Perceiver AR. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceed- ings of the 39th International Conference on Machine Learning, volume 162 of Proceed- ings of Machine Learning Research, pp. 8535-8558, PMLR, 17-23 Jul 2022. URL https //proceedi.ng...... press v1 http///////attttteeeee222 html
Louay Hazami, Rayhane Mama, and Ragavan Thurairatnam. Pfficient-DDVAA Less is more, 2022.
URL http::://axlv...//// org abs/ wtt.aa///2222222
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, 2023. URL
http:://
ccm google/flax
Tom Henighan, Jared Kaplan, Mor Katz, Mark Chen, Christopher Hesse, Jacob Jackson, Heewoo Jun, Tom B. Brown, Prafulla Dhariwal, Scott Gray, Chris Hallacy, Benjamin Mann, Alec Radford, Aditya Ramesh, Nick Ryder, Daniel M. Ziegler, John Schulman, Dario Amodei, and Sam McCan- dlish. Scaling laws for autoregressive generative modeling. CoRR, abs/2010. 14701, 2020. URL https://arxil: //tp:v//axivvv///m org 2000 14701
Alex Henry, Prudhvi Raj Dachapally, Shubham Shantaram Pawar, and Yuxuan Chen. Query-key normalization for transformers. In Findings of the Association for Computational Linguistics:
EMNLP 2020, pp. 4246-4253, Online, November 2020. Association for Computational Linguistics. doi: 1118655//12220.ffnnff
findinns--mmml.3799
URL httpp:://ccaattheeee wrg. 2020
firdings-emmip.33
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, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Jack W. Rae, Oriol Vinyals, and Laurent Sifre. Training compute-optinal large language models, 2022. URL https: 125561
Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. The curious case of neural text degeneration. In International Conference on Learning Representations, 2020. URL https //openreview net / http:/oommmmmddooyor--ee
Weizhe Hua, Zihang Dai, Hanxiao Liu, and Quoc Le. Transformer quality in linear time. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference Machine Learning. volume 162 of Proceedings of Machine Learning Research, pp. 0099-91119 PMLR, 17-23 Jul 2022. URL www.oeeediiiggs....o www.sp.esse... 62./ua22a.... html
12

---

Published as a conference paper at ICLR 2024
DeLesley Hutchins, Imanol Schlag, Yuhuai Wu, Ethan Dyer, and Behnam Neyshabur. Block-recurrent transformers. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL https://openeeiiee net/fformm d= oen
Andrew Jaegle, Felix Gimeno, Andy Brock, Oriol Vinyals, Andrew Zisserman, and Joao Car- reira. Perceiver: General perception with iterative attention. In Marina Meila and Tong Zhang (eds.), Pruceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 4651-4666 PMLR, 18-24 Jul 2021. URL https:/pprcceanngs www...t...ppss.eee /V39/ https//sss...aaaess//33////.....
jaegle2la.htmm
Norman P. Jouppi, Cliff Young, Nishant Patil, David Patterson, Gaurav Agrawal, Raminder Bajwa, Sarah Bates, Suresh Bhatia, Nan Boden, Al Borchers, Rick Boyle, Pierre-luc Cantin, Clifford Chao, Chris Clark, Jeremy Coriell, Mike Daley, Matt Dau, Jeffrey Dean, Ben Gelb, Tara Vazir Ghaemmaghami, Rajendra Gottipati, William Gulland, Robert Hagmann, C. Richard Ho, Doug Hogberg, John Hu, Robert Hundt, Dan Hurt, Julian Ibarz, Aaron Jaffey, Alek Jaworski, Alexander Kaplan, Harshit Khaitan, Daniel Killebrew, Andy Koch, Naveen Kumar, Steve Lacy, James Laudon, James Law, Diemthu Le, Chris Leary, Zhuyuan Liu, Kyle Lucke, Alan Lundin, Gordon MacKean, Adriana Maggiore, Maire Mahony, Kieran Miller, Rahul Nagarajan, Ravi Narayanaswami, Ray Ni, Kathy Nix, Thomas Norrie, Mark Omernick, Narayana Penukonda, Andy Phelps, Jonathan Ross, Matt Ross, Amir Salek, Emad Samadiani, Chris Severn, Gregory Sizikov, Matthew Snelham, Jed Souter, Dan Steinberg, Andy Swing, Mercedes Tan, Gregory Thorson, Bo Tian, Horia Toma, Erick Tuttle, Vijay Vasudevan, Richard Walter, Walter Wang, Eric Wilcox, and Doe Hyun Yoon. In-datacenter performance analysis of a tensor processing unit. SIGARCH Comput. Archit. News. 45(2):1-12, jun 2017. ISSN 0163-5964. doi: 10.1145/31400559.300 308004
URL https:
Lukasz Kaiser, Samy Bengio, Aurko Roy, Ashish Vaswani, Niki Parmar, Jakob Uszkoreit, and Noam Shazeer. Fast decoding in sequence models using discrete latent variables. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 2390-2399, PMLR, 10-15 Jul 2018. URL htt.thps:///oooceediiiiiio htt.://sss.m.....eeeess/// http://////aaieeeeeeeeeee .tmll
mlr www.s/resseeee.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford. Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. CoRR, abs/2001.08361, 2020. URL /RLtttts:://// http::///aaavaaaa abs /2001.111 08361
Angelos Katharopoulos, Apoory Vyas, Nikolaos Pappas, and François Fleuret. Transformers are RNNs: Fast autoregressive transformers with linear attention. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 5156-5165. PMLR, 13-18 Jul 2020. URL html
https: http:::/ppoceeeinnee
Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Ad- vances in Neural Information Processing Systems, volume 34, pp. 21696-21707, Curran Asso- ciates, Inc., /tt.p///////////peeeeee.....e.e.
files/paper/ /tt.a//rr.......eppooopoeeeaeeee
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In
International Conference on Learning Representations, 2020. URL
net./ fott/Forum?? 7 derkk
Taku Kudo and John Richardson. SentencePiece: A simple and language independent subword tokenizer and detokenizer for neural text processing. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pp. 66-71, Brussels, Belgium, November 2018 Association for Computational Linguistics. doi: 10. 111.1/6////D1112200
URL URR.htttps:://claattti http:///wee.antoooooooo..o//////
Guillaume Lample, Alexandre Sablayrolles, Mare' Aurelio Ranzato, Ludovic Denoyer, and Herve Jegou. Large memory layers with product keys. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems,
13

---

Published as a conference paper at ICLR 2024
volume 32. Curran Associates, Inc., 2019. URL http://prooceeingg.nrrrrr cc/ /aper/2019/Elle/// http://////.aaaaaaaaaa..oo.ao.a. 11
Doyup Lee, Chiheon Kim, Saehoon Kim, Minsu Cho, and Wook-Shin Han. Autoregressive image generation using residual quantization. In Proceedings of the IEEE/CVF Con-
ference on Computer Vision and Patlern Recognition (CVPR), pp. 11523-11532, June thecvf... com/ wwwttenn//PP2OO000oo Lee Autoregressive Image Generation Using Residualls Quantization.
2022. URL
CVPR 2022 paper.html
James Lee-Thorp. Joshua Ainslie, Ilya Eckstein, and Santiago Ontanon, FNet: Mixing tokens with Fourier transforms. In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 4296-4313, Seattle, United States, July 2022. Association for Computational Linguistics. doi: 10.18653/v1/ 2022. main.319. URL wttph:tt:///// http:///aaaaathoooog.o.o.o httpa//ggg.aaape///aaaaaattoogoo 2ttpa//wwwooooooo.............. maaaccl-main.. 319
org/2022
Tao Lei. When attention meets fast recurrence: Training language models with reduced com-
pute. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pp. 7633-7648, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics. doi: 10. 10.18653/v1/2011 (0065////222....mmpmmaaaaaa wmw.pmmmiin662..2 URL https: Bttts:://aclanthoooogy.ooo
emn main. 602
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling, The Eleventh International Conference on Learning Repre-
sentations, 2023. URL https: //teps::/epe //een review net./ forumfrrummooo
Peter J. Liu, Mohammad Saleh, Etienne Pot, Ben Goodrich, Ryan Sepassi, Lukasz Kaiser, and
Noam Shazeer. Generating wikipedia by summarizing long sequences. International Confer- ence Learning Representations, 2018. URL https:: /tttp::///penreviee.. http:/ewwwweeeteeeeeeeeeeee www.ffooofiiiommmmmo Hyg0vbWC-]
Zichang Liu, Aditya Desai, Fangshuo Liao, Weitao Wang, Victor Xie, Zhaozhuo Xu, Anastasios Kyrillidis, and Anshumali Shrivastava. Scissorhands: Exploiting the persistence of importance hypothesis for LLM KV cache compression at test time, 2023. URL http://aaxiv.orrr
abs/2305. 17118
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Representations, 2019. URL
net/forum?i
Bkg6RiCqY7
Shahar Lutati, Itamar Zimerman, and Lior Wolf. Focus your attention (with adaptive IIR filters), 14952
2023. httttt::/aaxiv.orrttt
Xuezhe Ma, Xiang Kong, Sinong Wang. Chunting Zhou, Jonathan May, Hao Ma, and Luke Zettle- moyer. LUNA: Linear unified nested attention. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL net/ 2tee://fooumm1ooooooooor..
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, U fttt/forummid=ggLeee33
Matt Mahoney, Large text compression benchmark, 2011. URL: http:ttpp///aattmmcceee.eeeeee dc/text htmll
Chengzhi Mao, Lu Jiang, Mostafa Dehghani, Carl Vondrick, Rahul Sukthankar, and Irfan Essa. Discrete representations strengthen vision transformer robustness. In International Confer- on Learning Representations, 2022. URL httpttttts//ppeeeeee//e net. ntt/forummiiii
8hM260AZcWkk
Harsh Mehta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. Long range language modeling via gated state spaces, 2022. http:httt:////ccc..oor//aaa/ 13947
14

---

Published as a conference paper at ICLR 2024
Piotr Nawrot, Szymon Tworkowski, Michal Tyrolski, Lukasz Kaiser, Yuhuai Wu, Christian Szegedy, and Henryk Michalewski. Hierarchical transformers are more efficient language models. CoRR, abs/2110.137111 2021. URL
Piotr Nawrot, Jan Chorowski, Adrian Lańcucki, and Edoardo M. Ponti. Efficient transformers with dynamic token pooling, 2023. httttt/tp///xittt otg/abs.....ooo.oo.............t 2211.11111. 09761
Emilio Parisotto, H. Francis Song, Jack W. Rae, Razvan Pascanu, Çaglar Gülçehre, Siddhant M. Jayakumar, Max Jaderberg, Raphael Lopez Kaufman, Aidan Clark, Seb Noury, Matthew M. Botvinick, Nicolas Heess, and Raia Hadsell. Stabilizing transformers for reinforcement learning. 06764
CoRR, abs/1910.06764, 2019.
http::://xxxvv.orrg/
Bo Peng, Eric Alcaide, Quentin Anthony. Alon Albalak, Samuel Arcadinho, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, Xuzheng He, Haowen Hou, Przemyslaw Kazienko, Jan Kocon, Jiaming Kong, Bartlomiej Koptyra, Hayden Lau, Krishna Sri Ipsit Mantri, Ferdinand Mom, Atsushi Saito, Xiangru Tang, Bolun Wang, Johan S. Wind, Stansilaw Wozniak, Ruichong Zhang, Zhenyuan Zhang, Qihang Zhao, Peng Zhou, Jian Zhu, and Rui- Zhu. RWKV: Reinventing RNNs for the transformer era, 2023. URL /tte:///aaxx..oorr/////////
13048
Hao Peng, Nikolaos Pappas, Dani Yogatama. Roy Schwartz, Noah Smith, and Lingpeng Kong. Random feature attention. In International Conference Learning Representations, 2021. URL
httpps://ppenrret
net/ forum (orrummIFFQETTTTVVFF
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, 2023. URL http:://arxa arxiv. www./iwww.w...ioo...oooooooooort wwtgabs/////ooooooooott..pt...et abs/ 2302 10866
Zhen Qin, Xiaodong Han, Weixuan Sun, Dongxu Li, Lingpeng Kong, Nick Barnes, and Yiran Zhong. The devil linear transformer. Yoav Goldberg, Zornitsa Kozareva, and Yue Zhang (eds.), Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp. 7025-7041, Abu Dhabi, United Arab Emirates, December 2022a. Association for Computational Linguistics. doi: /6655////0222emmmmmmmmoae 22022eemnpp-nn
URL http:: htttppps//accaaaaaooorrr.../
Zhen Qin, Weixuan Sun, Hui Deng, Dongxu Li, Yunshen Wei, Baohong Lv, Junjie Yan, Lingpeng Kong, and Yiran Zhong. CosFormer: Rethinking softmax in attention. International Confer- ence on Learning Representations, 2022b. wwt.hhtts///ppeeeeeee B18CQrx2Up4
net. forumm id=
Jiezhong Qiu, Hao Ma, Omer Levy, Wen-tau Yih, Sinong Wang, and Jie Tang. Blockwise self- attention for long document understanding. In Findings of the Association for Computational Linguistics: EMNLP 2020. pp. 2555-2565, Online, November 2020. Association for Computational Linguistics. doi: 10. 10.18653/v1/2020 orq orq/2020. findings-ennnnn 232
www.:oiiinnngnnninn.n emnlp.232. URL https::/aclanthhl..og
Markus N. Rabe and Charles Staats. Self-attention does not need o(n2) memory. CoRR, abs/2112.05682. 2021. URL
baargg/abb/1111 05682
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners, 2019. https
//cdn.openate htw.://aao..ooeetttteeeaaaagaeee mode mode s h//aannuageeeaaae unsupervised multitask lrtiiaskkkeearrae learners. pdf Last visited 2023/09/07.
_s_are__
Jack W. Rae, Anna Potapenko, Siddhant M. Jayakumar, Chloe Hillier, and Timothy P. Lillicrap. Com- pressive transformers for long-range sequence modelling, In International Conference on Learning Representations, 2020. URL https:://ppeerevieeet wttpr//wwwwweeeeeeooummeeeed-eee
Jack W. Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, H. Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young. Eliza Rutherford, Tom Hennigan, Jacob Menick, Albin Cassirer, Richard Powell, George van den Driessche, Lisa Anne Hendricks, Maribeth Rauh, Po-Sen Huang, Amelia Glaese, Johannes Welbi, Sumanth Dathathri, Saffron
15

---

Published as a conference paper at ICLR 2024
Huang, Jonathan Uesato, John Mellor, Irina Higgins, Antonia Creswell, Nat McAleese, Amy Wu, Erich Elsen, Siddhant M. Jayakumar, Elena Buchatskaya, David Budden, Esme Sutherland, Karen Simonyan, Michela Paganini, Laurent Sifre. Lena Martens, Xiang Lorraine Li, Adhiguma Kuncoro, Aida Nematzadeh, Flena Gribovskaya, Domenic Donato, Angeliki Lazaridou, Arthur Mensch, Jean- Baptiste Lespiau, Maria Tsimpoukelli, Nikolai Grigorev, Doug Fritz,, Thibault Sottiaux, Mantas Pajarskas, Toby Pohlen, Zhitao Gong, Daniel Toyama, Cyprien de Masson d' Autume, Yujia Li, Tayfun Terzi, Vladimir Mikulik, Igor Babuschkin, Aidan Clark, Diego de Las Casas, Aurelia Guy, Chris Jones, James Bradbury, Matthew J. Johnson, Blake A. Hechtman, Laura Weidinger, Iason Gabriel, William S. Isaac, Edward Lockhart, Simon Osindero, Laura Rimell, Chris Dyer, Oriol Vinyals, Kareem Ayoub, Jeff Stanway, Lorrayne Bennett, Demis Hassabis, Koray Kavukcuoglu, and Geoffrey Irving. Scaling language models: Methods, analysis & insights from training gopher. / rxi org. 21 12. 11446
CoRR, abs/2112, 22./2112.112 1446, 2021.
Prajit Ramachandran, Barret Zoph, and Quoe V. Le. Searching for activation functions. CoRR, abs/1710.05941, 2017. URL http:tttaaaii...orrr. hbcgg//bs//// 05941
Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ily Sutskever. Zero-shot text-to-image generation. CoRR, abs/2102.12092, 2021, URL.
https:/ /ataiv///rvvvv../ www.orggaa.... abs//aa///// 2 102. 12092
Ali Razavi, Aaron van den Oord, and Oriol Vinyals. Generating diverse high-fidelity images with 2q- In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché- Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Asso- ciates, Inc., 2019. URL Uttttttss::tt http::e::///cccceeeeeeooo ings rww. wwtpeuuwwsspccpppppp.ooeeeoepp ww/pscer.ppppeeeeee fearrr_fileeppppeee les/paper/ 2019 fi /e/ 5f8e2fal1l1111a /tt.:///e..aaaeaaaaaaaaaaa///aaa
Hongyu Ren, Hanjun Dai, Zihang Dai, Mengjiao Yang, Jure Leskovee, Dale Schuurmans, and Bo Dai. Combiner: Full attention transformer with sparse computation cost. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Ad- vances in Neural Information Processing Systems, volume 34, pp. 22470-22482 Curran Asso- ciates, Inc., 2021. URL
2021/fi URL.htpppe//prrrceeeisse..eeeaaa
files/paper/
Stephen Roller, Sainbayar Sukhbaatar, Arthur Szlam, and Jason E Weston. Hash layers for large sparse models. A. Beygelzimer, Y. Dauphin, Liang, and Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. httpttsss///ppeeeii.lteee wwwp:ewwwwwww.weeeeeeeeeeeeeeeee for
1d= 1MgDDWblU 1UW.
Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguistics, 9:53 68, 2021. doi: 10. 10.1162/tacl a 00353. URL bttpps://aclaaathhooo.....
2021
tacl 1.4
Aurko Roy, Rohan Anil, Guangda Lai, Benjamin Lee, Jeffrey Zhao, Shuyuan Zhang, Shibo Wang, Ye Zhang, Shen Wu, Rigel Swavely, Yu Tao, Phuong Dao, Christopher Fifty, Zhifeng Chen, and Yonghui Wu. N-Grammer: Augmenting transformers with latent n-grams, 2022. URL
https://arxis: wttp:b///////./o///// httpps///axxx..orrr///////.....2 06366
Noam Shazeer. Fast transformer decodin One write-head is all you need. CoRR, abs/1911.02110,, 02150
2019. wtRRl/ttt:///axii.rrr/aaa///////
Noam Shazeer. GLUU variants improve transformer, 2020. URL httpss://arxivvorr//bb
2202205202
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. CoRR, abs/ 104.042355 2018.
http::::///x/vv.aa org/ abs/ /184. 04235
Jimmy T. H. Smith, Andrew Warrington, and Scott W. Linderman. Simplified state space layers for sequence modeling, 2022. httphtttt:////xiiii org/ htspgg/bb/22000 04933
Sainbayar Sukhbaatar, Edouard Grave, Piotr Bojanowski, and Armand Joulin. Adaptive attention span in transformers. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics. pp. 331-335, Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10.18653/v1/Pm 1/P19- 1032.
//tpss//aco / anthol ogy. wgt..og/p/9911100/
16

---

Published as a conference paper at ICLR 2024
Sainbayar Sukhbaatar, Da Ju, Spencer Poff, Stephen Roller, Arthur Szlam, Jason Weston, and Angela Fan. Not all memories are created equal: Learning to forget by expiring. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 9902-9912. PMIR, 18-24 Jul 2021. URL https: ftttps::/proceedii ngs. m www.r.press//ee wwtp:///wwaaaaaaaaaaaaaa html
Simeng Sun, Kalpesh Krishna, Andrew Mattarella-Micke, and Mohit lyyer. Do long-range language models actually long-range context? In Proceedings of the 202/ Conference on Empirical Methods in Natural Language Processing, pp. 807-822, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics. doi: 10.18653/1/2021. emnlp- monlp-main.66 main.62. htt.ttss: httpssss/aclaannhoooooooo. org. emn emnip-nnnnmm. in....n... 62
Yi Tay, Dara Bahri, Liu Yang, Donald Metzler, and Da-Cheng Juan. Sparse sinkhorn attention. CoRR, abs/2002. 1296, 2020a.
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: survey, CoRR, abs/2009.06733 06732,, 2020b. URL [RLLtttss///xxii// org. abs 2009:/2200.0000 06732
Yi Tay, Dara Bahri, Donald Metzler, Da-Cheng Juan, Zhe Zhao, and Che Zheng. Synthesizer: Rethinking self-attention for transformer models. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning. volume 139 of Pro- ceedings of Machine Learning Research, pp. 10183-10192 PMLR, 24 Jul 2021. URL https: hrg.:paee..n..l......prss///////
Aaron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu. Neural discrete representation learn- ing. In Guyon, U. Von Luxburg, Bengio, H. Wallach, R. Fergus, Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran As- sociates, Inc., 2017. URL httpstts:::///tttt htttt:ps:///prooeeednngg.e wtw.:roocceeiinnn...........e.ee neurips cc pccpppppr/2001//11.... 2.45a277ee3300000222222633232223
Aäron van den Oord, Nal Kalchbrenner, and Koray Kavukcuoglu. Pixel recurrent neural net- works. In Maria Florina Balcan and Kilian Q Weinberger (eds.), Proceedings of The 33rd International Conference on Machine Learning, volume 48 of Proceedings of Machine Learn- ing Research, pp- 1747-1756, New York, New York, USA, 20-22 Jun 2016. PMLR. URL https: attps::/prrceeedinn
mlr. prwss/peee//ee.e. /4// 048//ood155 html
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, L ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Gamett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Asso- ciates, Inc., 2017. URL htttps::/prooeeednngg www.idiinnnnnnuiiss..ee.. wwtrrips.cc/paaaee wccppppper//11//11.. htt.:///ee...............a..0000
pdf
Apoorv Vyas, Angelos Katharopoulos, and François Fleuret. Fast transformers with clustered attention. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Ad- vances in Neural Information Processing Systems, volume 33, pp. sociates, Inc., 2020. URL wtt.liitttss:/prrceeeedae c
Curran As-
neurips cc .ddff
Ningning Wang, Guobing Gan, Peng Zhang, Shuai Zhang, Junqiu Wei, Qun Liu, and Xin Jiang. ClusterFormer: Neural clustering attention for efficient and effective transformer. Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics ( (Volume I: Long Papers), pp. 2390-2403 Dublin, Ireland, May 2022. Association for Computational Linguistics. doi: acl- acl-long
10. 10.18653/w1/2022.acll 170. URL 170
Shuohang Wang, Luowei Zhou, Zhe Gan, Yen-Chun Chen, Yuwei Fang, Siqi Sun, Yu Cheng, and Jingjing Liu. Cluster-former: Clustering- based sparse transformer for question answering. In Findings of the Association for Computational Linguistics: ACL-IJCNLP 2021, pp. 3958- 3968, Online, August 2021. Association for Computational Linguistics. doi: 10. 10.18653/v1/2221. Andings-acl.3344 / http:: peattp::/acllathhee /at/aallaatthho lgy. htt.tsss///llaatoooooooooog//...
ndi findingn-nnninnnn 346
17

---

Published as a conference paper at ICLR 2024
Sinong Wang. Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: self- with linear complexity. CoRR, abs/2006.04768, 2020. URL https::/axxwv.xxgg// 04768
Yuhuai Wu, Markus Norman Rabe, DeLesley Hutchins, and Christian Szegedy. Memorizing transformers. In International Conference on Learning Representations, 2022. URL https
Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and
Vikas Singh. Nyströmformer: A nyström-based algorithm for approximating self-attention. Pro- ceedings of the AAAI Conference on Artificial Intelligence, 35(16)14138-14145, May 2021.
doi: 1t..1609/aaai.v3351..0 17664. URL a wrt.//dex...i..oooooooooo pwt.dexx.pp/AAAIII. axticle/view/17666
Zihao Ye, Qipeng Guo, Quan Gan, Xipeng Qiu, and Zheng Zhang. BP-Transformer: Modelling long-range context via binary partitioning. CoRR, 125/1911.04070 2019. URL http://arxiv
org ht5/abs/19911 04070
Lili Yu, Dániel Simig, Colin Flaherty, Armen Aghajanyan, Luke Zettlemoyer, and Mike Lewis. Megabyte: Predicting million-byte sequences with multiscale transformers, 2023. URL http
//arxiv.org abs 220.2/335.000 07185
Biao Zhang and Rico Sennrich. Root mean square layer normalization. In H. Wal- lach, H. Larochelle, A. Beygelzimer, d'Alché-Buc,
Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Asso- ciates, Inc., 2019. URL httpps://prooeeednggs pttprrppp.ccppppee//01110030 12823392222443833333771111122--2
Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song. Yuandong Tian, Christopher Ré, Clark Barrett, Zhangyang Wang, and Beidi Chen. H20: Heavy- hitter oracle for efficient generative inference large language models, 2023. URL https
//arxiv.org/abe 2atobb///2266.44144..... 14048
Shangchen Zhou, Kelvin C.K. Chan, Chongyi Li, and Chen Change Loy. Towards robust blind face restoration with codebook lookup transformer. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL https: www.neveeeeeeeeeeoooooee net forum? 11dddddddDD 36FUNn5
Chen Zhu, Wei Ping, Chaowei Xiao, Mohammad Shoeybi, Tom Goldstein, Anima Anandkumar, and Bryan Catanzaro. Long-shor transformer: Efficient transformers for language and vision. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Ad- vances in Neural Information Processing Systems, volume 34, pp. 17723-17736. Curran As- sociates, Inc., 2021. URL httpt:tttttt/ottt htttppp:://pccccciiccee www.rieeeiippp.ee cC 442
Paperr pdf
Zhenhai Zhu and Radu Soricut. H-transformer-ID Fast one-dimensional hierarchical attention for sequences. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pp. 3801-3815, Online, August 2021. Association for Computational Lin- guistics. doi: 10.18653/v 1/202 294. URL https: http:///aclaatoooooo.o/////////0 acl-longe
18

---

Published as a conference paper at ICLR 2024
A THEOREMS
A.1 PROOF OF THEOREM 2.2
Proof. This proof is based Guo et al. 2019). Invoking the fact that q, k, (p(k) RD, the assumed independence between and k, the law of iterated expectations, and the isotropy assumption on ie., Eq In for o² > O, we have
kq k q (k))
(38)
-kq T (k - (ok))]]
(39)
-Eq (k - (o(k)) [q T (k - ((k)
(40)
E (k- ((k)) qq (k- ((k)
(41)
Ek p(k)) Eqlaq lqq [[k ((k))
(42)
Ex *(k)) 'ID(k ((k))
(43)
Ex|lk - =(k)]l?.
(44)
A.2 PROOF OF COROLLARY 2.3
Proof. By definition, VQ(k; C) arg min (C.))) k cll?. In other words, <o(k) VQ(k: C)
minimizes | yo (k)ll under the constraint that the outputs of yp limited the rows of C, i.c. () [C, s=0 Since this choice is pointwise minimizer under the given constraint, it is also a minimizer of the expectation Ex k =(k)]l under the same constraint.
S-1
Under the assumptions of Theorem 2.2 the aforementioned expectation is equal to l k q (kk]] up positive proportionality constant o². As result, VQ(k; C) is also minimizer of the expectation lq,a/l k q (<k) under the same constraint on the output of sp.
A.3 PROOF OF THEOREM
Proof. When &n an element-wise nonlinearity, ((c) is well-defined, where e any scalar. Then using definitions alone, we have
(
)QCT)Ak
A
2n (QC )r. A..
(45)
$-1
dw (QC )r,e
(46)
5
8-1
Σ =
Dil (wi 0,,10,
(47)
iw Qu.. CT.)
(48)
Siw (Q., K ).))
(49)
= ( (QK) T )kk
(50)
A.4 PROOF OF THROREM 3.5
Proof. By Theorem 3.4 with (-) exp(-), we have exp(QK]) exp(QC) ). Invoking the definition of row-wise softmax and applying substitution, we have
Softmax(QK*
= Diag(exp(QR) - eep((kk )
(51)
= -Diageexp(QC (QC A1)- "
(52)
n
19

---

Published as a conference paper at ICLR 2024
A.5 PROOF OF THEOREM 3.6
Proof. For 71 0, the result follows by inspection.
For n 2, by the definition of U(n 2) we have
1n-1)
Unn - 2)
Σ J=0
A. (10mm (0.-11.))
(53)
Note that in notation, the superscripts' block index range is non- non-inclusive the ending value, (1,0m-1) V (0:x-1,,) is equal to the sum of the matrix products for the matrix blocks from 0 to n
Thus, by substitution, we have
Φ (m.)) C CUN 2) 2 (Q C (1,0m-1) V 10:n-1,1))
(54)
We invoke the same argument as in the proof of Theorem 3.4 to conclude ow (m.:) C
(s.0mmm)1))) =
(s,0m-1)
W (m,,0m-1). Substituting this expression into the right-hand side above gives
@w Q (n.:) C (")U(n 2) = W (n.0:n11) V (0:--.:))-
(55)
Substituting this expression into the formula for (WV) (a.)) claimed in the theorem statement, and
invoking the same argument as in the proof of Theorem 3.4 on the middle term, we see the claimed formula has the form W (n.0:n-1) V 10:n-1:;) + W (n.n-1) (n-1,)) W (n.n) (n.)) The diagonal
block W (n.n) of W is causally masked, so the sum of the three terms indeed equals (WV)
A.6 PROOF OF THBOREM
Proof. Recall that we defined A - exp(QK + B). The proposed expression for (AV) (m..) follows from Theorem 3.6 with (.) = exp(-). The proposed expression for (A1) (n) follows by substitution argument using (n.)). Normalizing (R.) by (A1) (m) and iterating over n thus yields all blocks of the product WV when the nonlincarity &w is row-wise softmax.
20

---

Published as a conference paper at ICLR 2024
B THROUGHPUT
We present throughput results for three methods to compute the cache variables: serial sean, matmul, and associative scan. The first are generalizations of the cross-block reduction methods from FLASH (Hua et al. 2022), which were simple cumulative sum and matrix multiplication by lower-trianuule matrix of ones, respectively. We found our proposed generalizations were necessary for stable training, an issue where FLASH has known weaknesses (Qin et al., 2202aa Ma et al. (023))
Pseudocode each of our stable reduction methods given Appendix E
In addition to the three cross-block reduction methods to compute the cache variables from parallel-
computed per-block summaries, we also benchmark an input scanning implementation of VQ- Attention inspired by Wu et al. (2022); Hutchins et al. (2022). such that all the operations for transformer layer are performed one input block at a time. To ground all comparisons, we benchmark the throughput against transformer using unquantized quadratic-time attention, the same attention head type (SHGA, MQA, or MHA), and an identical non-codebook parameter count.
Table 6: Training throughput comparison (tokens/see) on Google Cloud VM with TPU v3 cores, between Full Attention and VQ-Attennton with serial scan reduction.
Sequence Length
Model
2048
8192
32768
131072
VQ
Speedup
Full
VQ
Speedup
Full
VQ
Speedup
Full
VQ
Speedup
SHGA
65.5k
63.0k
0.962x
23.9k
77.2k
3.230x
6.5k
82.1k
12.631x
OOM
OOM
63.0k
1.070x
10.4k
74.8k
7.192x
0OM
79.7k
0OM
MHA
52.0k
49.6k
0.955
9.5k
57.8k
6.084x
0OM
60.7k
0OM
0OM
Table 7: Training throughput comparison (tokens/sec) on Google Cloud VM with 8 TPU v3 cores, between Full Attention and VQ-Attention with matmul reduction.
Sequence Length
Model
2048
8192
32768
131072
Full
VQ
Speedup
Full
VQ
Speedup
Full
VQ
Speedup
Full
VQ
Speedup
SHGA
65.5k
62.5k
0.954
23.90
75.2k
3.148x
6.5k
80.0k
12.250x
OOM
69.5k
MQA
58.9k
62.5k
1.061x
10.4k
74.9k
7.144x
0OM
80.2k
67.7k
MHA
52.0k
48.2k
0.926x
9.5k
58.2k
6.096x
0OM
61.6k
OOM
OOM
Table 8: Training throughput comparison (tokens/sec) on Google Cloud VM with 8 TPU v3 cores, between Full Attention and VQ-Attention with associative scan reduction.
Sequence Length
Model
2048
8192
32768
131072
Full
VQ
Speedup
Full
VQ
Speedup
Full
VQ
Speedup
Full
VQ
Speedup
SHGA
65.5k
58.9k
0.899
23. 23.9k
62.2k
2.603x
6.5k
63.4k
9.754x
OOM
OOM
MQA
58.9k
62.8k
1,066x
10.4k
74.2k
7.134x
0OM
79.0k
67.0k
MHA
52.0k
49.1k
0,944x
9.5k
55.4k
5.831x
OOM
57.8k
OOM
OOM
21

---

Published as a conference paper at ICLR 2024
Table 9: Training throughput comparison (tokens/sec) on Google Cloud VM with 8 TPU v3 cores, between Full Attention and VQ-Attention, both with input scanning.
Sequence Length
Model
2048
8192
32768
131072
Full
VQ
Speedup
Full
VQ
Specdup
Full
VQ
Speedup
Full
VQ
Speedup
SHGA
32.0k
40.8k
1.275x
12.7k
47.4k
3.732x
OOM
49.4k
OOM
OOM
MQA
36.2k
54.8k
1.514x
14.5k
64.9k
4.476x
0OM
68.3k
0OM
OOM
MHA
43.7k
1.457x
11.4k
50.7k
4.447
0OM
53.1k
22

---

Published as a conference paper at ICLR 2024
C
TRAINING DETAILS
C.I HYPERPARAMETERS
Per-dataset hyperparameters are provided below.
Table 10: Hyperparameters.
Name
Symbol
Enwik8
PG-19
ImageNet64
ImageNet64
parameter count
190M
1.3B
190M
1.2B
global batch size
B
128
128
16
128
sequence length
T
8192
8192
12288
12288
backprop length
2048
2048
12288
2048
block length
512
512
512
512
model dimension
Dm
768
2048
768
2048
key dimension
Dx
128
128
128
128
value dimension
D.
1536
4096
1536
4096
code
S
512
512
512
512
num gau
48
sinusoid dropout
Püropsin
0.2
0.1
0.1
0.1
residual dropout rate
Puropres
0.0
layerdrop rate
Phroply
0.3
0.1
0.0
0.0
weight decay optimizer total steps
0,0002
0.0
0.0
0.0
adamw 125000
adafactor
adamw
adafactor
500000
125000
500000
Note that the 190M parameter ImageNet64 result was added after the other experiments had concluded. To avoid biasing its result, we use the exact same architectural hyperparameters' the Enwik8 model, and the exact same regularization as the larger ImageNet64 model. The smaller ImageNet model was trained in a newer version of our codebase optimized for higher throughput and faster compile times, rather than training on long sequences in constant space via input scans and truncated backprop through time. The attention in the optimized codebase was unit-tested match the original.
C.2 IMPLEMENTATION
Weights and token embeddings were initialized following Chowdhery et al. (2022). For the small model, the classifier layer omits LayerNorm and is independently parameterized. For the large model, the classifier layer uses LayerNorm and its projection is tied with the token embedding table, then scaled down by large constant. For image datasets, we add absolute sinusoidal position embeddings, scaled by trainable scalar, to the token embeddings (Hua et al. 2022 Vaswani al. 2017). We used a maximum angular wavelength of 10" for all sinusoidal embeddings.
We used the pre-norm placement of LayerNorm (Radford et al. 2019). and always used the RMS LayerNorm variant (Zhang & Sennrich, (019). For the activations, we used &.. Softmax and d. = Øy SiLU,, the self-gated activation (Elfwing al. 2017; Ramachandran al., 2017). Several models use LayerDrop for regularization (Fan et al.) 2020a), and following the Transformer-Xi codebase (Dai al. (019) models apply dropout to the flipped sinusoidal embeddings used for (local) relative positional biases.
We used float32 parameters, with bfloat16 precision for most computations (Rae al. 2021). For the AdamW optimizer (Loshchilov & Hutter 2019), we used gradient clip 0.1, max learning rate 0.0004 and hyperparameters 31 0.9, 32 0.98, e 10 For the Adafactor optimizer (Saazeer & Stern 2018), we used relative stepsizes, update clip 1.0, max learning rate = 0.01,
0.8 We used weight decay with constant schedule throughout training and omit decay on any oie-dimensional parameter tensors (Radford et al. 2019). always 3 = 0.0001 and codebook EMA rate was always 0.99. Learning rates were linearly warmed up for 10,000 steps, then decayed by 10x factor using cosine schedule.
and hyperparameters 3, 0.0, 32.
The codebook commit coefficient
23

---

Published as a conference paper at ICLR 2024
D GENERATED SAMPLES
D.1 QUALITATIVE ANALYSIS
D.1.1 PG-19
No affort haa been made to explain alanentary nethada of photography, for the reason that auch explanation has bean found publications of every Leading technical journal. endeavor been present what recessary to the ansteur and the profersional photographer, together with suggestious of how to make apparatus
building. The author fully inporfections in the nethods described, and sould like to euphazizo the necessity of studying these aethods carefully before assential point in photography the oparator knava shat light ia. The uriter balieres that much simple apparatus.
student,
give chapter
then, it dusired nake satiefactory photographs. The study light. It. inpoasible to aucoess in photography unleaa
done advance the art of photography
by the nacessary in order to get good reaulta. A lena ia nacassary to bring the inage on the sanaltiva plate up to vary expensive and only
simple apparatus is
focus of the
had
makers.
Figure 4: Sample excerpt from our PG-19 model, generated with nucleus 0.8.
We generated I 128 sequences using nucleus sampling (Holtzman et al. 2020p. In Figure 4 we observe a sample except in which our PG-19 model synthesizes high-quality text, and maintains a consistent tone, topic, and train of thought. These observations were found to hold for the vast majority of the samples we generated.
D.1.2 IMAGENET64
Figure 5: Generated samples from our large ImageNe164 model; nucleus 0.999.
Figures 3 and 5 show subset of samples with the same indices from two batches with different nucleus settings. We see that our large ImageNet64 model synthesizes sequences over 12,000 bytes and is capable of depicting relatively high-fidelity ocean water, shorelines, leaves, insects, trees, animals, people, mountains, and architecture.
D.2 EXTENSIVE SAMPLES
Samples for Enwik8, PG-19, and ImageNet64 be viewed the anonymized URLs in Table II]
Table 11: Generated samples' URLs by dataset.
URL
https:
D.3 IMAGENET64 - FULL BATCI
24

---

Published as a conference paper at ICLR 2024
u
25

---

Published as a conference paper at ICLR 2024
26

---

Published as a conference paper at ICLR 2024
E PSEUDOCODE
import flax. Linen
inport jax
inport jan.mump as jap inport chex
class Von :odule))
nodo::
d_k: 1nt
d_v:
Van. conpact
def c_.lll... (self, x):
"T'Input shape: Dbatch size, blocke, block length, nodel. width].
shape
8, K, V self-n_code, self.d_k, self.d_y
(orta)erroommrittm-- (axis=-1)
RMST WastayerFamman (axia=-1) (a) (mn Denaa (anlf.a)). (K_tilde))
jax. jar.mn. amx.mm. (0 jax, www..mmmmmmmmmmmmmmmmmmmmmmmm dtt.m.c.dcmm...neeeee.........). (x)
quantizer Vectorfumttierr (etpeeooo..iiaeeeeeellee
viddb-eelf.d.b
k_hat, 2, i_commit, Loodebook quantizer(k) quantized keys, shortcodes, atc
quantizer (es_codebookk)
local, biases
1angthe2<c)
sesert__haap((ii (local [8, R, C, 2+c])
local blasea_prwy lncal faaaas__preeee * jnp.aplit(tacal blases, 2, axia=-1) k.hat)
acores_prasant hmtp.einsunntaaeee ("briitbrjjj.jta-r
scores_preseat Local__iases_ppeeeet
ccure_present 1e30 (1 -jnp. (wt.eriii..eess....a...... htg.:poeee..leeeeerrrrprreeeee
haht_prry jap.pad(k_batlt, :-1], (CO, 0), (1, 0), (0. 0), (0, 03)) 0)))
V_prev inp- pad((: :-1), (to, 0), 0), "nnuunn"rri,bbrjeeeaaa acores_ pravv_peeee__ = lesal_ bissl_bisses_prev
9, k. k_haat_prev)
acores, prav jap-pad((
scores_ prev[: 1:3.
0), constant constant_walugge-eeee values=-1600
(0, 0)),
cache jap. eissun (brik ek->bris", cacke bicce_u.div___by.. block, cache L.che__.by_bllle block get_cacbo_
thes assert, aswert_eehppeeeeeeee wwtpeceppppeceeeeeeeeee Beceabbee_ddv_____e w.tpeuuviivllbbbb..eeeo.b block, [6, R, 8, V)) 8])
(caph black, count, count_biases jup.choro(
jmp-greatar(Cachhe b_by_block, block, jnp.zeros. (ekeee back)),
(og./njcc b_b, block, a,nio=1.0)), felpfuul__ikeee
vall_values=-e20),,
cache Jap. exp.erpandddiie (couttt (count, blases, axise-2)
ecores_ present, naz jup.xax(scoree) pressett, axis-1)
ceoresspreeent
scores nax japp (acores, present, axisu-1)
cacreescacch__. ecores, cache_ jaz jnp. www.emrreecccccaaa axige<1) acores rax jep.narimun(
www.aaiuun(corrssppeeaaaa jap.
scores.prav_nme
scores_cache_nas,
scores_ rax lax.lax.ttop...... ant (scores__
ecoree_presett presect = ecoree_zatt None]
ecores_pres prev
scores_cache scores_max [........ Hone]
A presannt: jnp.asp(acores aepppaccoroo_praaaat))
a_prer wwp.eeppacorerapppppppppppppp
jjjp..e wwwppppppacoreeecccmom cache))
jap. present, axiae-1)
jap.eum(a_prevv, axise-1)
eacheems__ccche, axic=-1)
k_present a.preseat / d(..., None]
k_prev a_prez d(..., None]
s_cabbe =_ccache d(..,, None]
jjnp.einsun""rrr (injunn"rri,..bjj--briis
v_present,
we += jng.einaun ("brij""brij.... ("brrij.rjjeeaaee brje"",rir",, k.prer, v.prov)
inp.uin ("bris, brav=sbriv", K_cache, cache_a,di d.v.ivvvllbbb... block)
residual nn. (D) (o)
retura residual, 1.comnit, Loodebook
Code 1: Jax/Flax pseudocode for VQ-Attention.
27

---

Published as a conference paper at ICLR 2024
def set_cacho__arreeee vars(z, R_code):
throughout function, ve in the places where we do this, it does not alter the actual cacae variable astinates will be vhen the clip applied.
clipping elezentwise
since the corresponding entriee
delta = Jar.en. www..anoone..ttto nun classes=n, code, dtyperv.atyp,, axta=-1) deltai_ by block jop eineun("ariee delta)
deltav by block imp.ainauun"bria.. brivmsbara dalta, v)
deltav_ by plock_pornalszed jap.divide(
deltav_by block, bllockk, hipp.aappdaaaaiitp block(.... Konal, a.nto=1.0),
func(carry, in_dict): coaputes runcing average of the value vectors for each shortcode ('upper div lover"),
and running count (*lower"),
carry( ["losar"]]
lower lower_block in_dict("deitaa_bby bt__eetta_.bbbbocrr""
lover louer block
11 aip.dividellower, jap.clip(Lover new, a,xin=1.0))
12 (eppddivideLooee_beeee jap.elip(Lower_ney a_mim=1.0)) upper_div. sepor_diy.lover_eee jap.add(
21[..., None] carry("upper_ds div, wlr.d_v__ve__ jav_lever"),
Nene] in_ dedttarr_byyy btw.kbbbbbbocccccoolleeaaee
carry dict(
upper_ wi__dddvvveeee piverrupppr_ddv_ioooe louer= pev,
retura cee,y_new, carry _ae state to carry, output to save
betore we scan, wo have to transpose since jax only supports scans along axis 0. this still fast, poseibly because jup tranepoce
choosing to view tav by block, normal ized jnp, transpose (deltav, by block, norcck_orrmlieed (1, 0, 2, 8)) blocckoooomllliiiee
delta! by block jep.tranapoa by. block, {1, 0, 2))
eache_ vare lax. feacan, Func,
init=dict(
seper_div_loverriie jap.rjjp..zoo. zoros secosfdssprreel..dyg dtype, chapondeltav_br. www.byy__bboooooeeeecccee pow.alozeeeeeeoom coa..cooooollieeeeeeee__1 ch:pe(1:]), lover=jup.zaros (dtype=aelf dtypa, dtype, shape-deltal, blopk. shapa (1:]),
xsrdict(
by corciinornallli izededeltar boock_normllizedd
byltai_by byloiibbbblociiddelta.... blocirdeleal.byybllcc.
unrollet,
cache, var, uppar div loser jap.pad(
jap. transpone (cache_vare _arrs ("apper_ciy..ccc. div___oovre___________ [1-2], (1, a)),
(10, 0), (2, 0), (o, 0), (0, 0)),
cache, var vov_r jap- pad(
jap. transpore (cache, vars ("Iover) [:-2], (1, 0, 2)),
((0, 0),
0)),
return cache, div lover, cache_vat vau_lover
Code 2: Jax/Flax pseudocode to get cache variables for all blocks; serial scan version.
28

---

Published as a conference paper at ICLR 2024
def set_cachho_arreeee vars(z, R_code):
throughout
clipping elezentwise
avoid
in the places where we do this, it does not alter the actual cacae variable astinates will be vhen the clip applied.
since the corresponding entriee
delta = jar.en. wwr..anoone..ttte
classessn code, dtyperv.atypa, axta=-1) block jop. eineun(""riee delta)
deltav by block jep.ainauu ("baann""rraaanrrrr brivssbar",,, dalta, v)
deltav_ by block__orrmaizzd jap.divida(
deltav_by, block, blockk,, cippeaaaaaa by, block(.... Konal, a_nio=1.0),
deltal block,
jnp.eineuun(
"bre, bgeeeberg',
jap. jap.ones_ (deltaal blokkk)) deltai by. block,
deltal, fia_frraaa__aa by, block jup.divida( (ddltaa by block_ck_ttii
deltat
jap clip(jap. ""neuu(""erg-"ber" deltal_by, ba_as__b__bllee by__bloc_tteeell) (..., Nonel, s_nin=1.0 0),
by. block comllatvvenorrslllee
(ineum(
"burg, "berg, bbea"", deleal_frace bar_bby_bbblbe block, deltav_by, blockkyybblooo..o....o. normalied
deltal, by cumolatvee Jap. (dnsux bybbockk, axts=1)
cacha, var wtr___pparrdde sirr_div__eee jap-pad(
deltar_by, www.aaaatteeeeoonnia.i.
:-2],
((0, 0), 0),
(0, 0)),
cache_ lover jap-pad(
deltai by bllock_cmmllatveeeee :-),, ((O, 0), (2, 0), (0, 0))
www_uaper__eeeeeeeeepp var_upper_div_lowee loverover, vov_lover
Code 3: Jax/Flax pseudocode to get cache variables for all blocks; matmul version
29

---

Published as a conference paper at ICLR 2024
def set_cachoo_rrreee vars(z, R_code):
throughout
clipping elezentwise
in the places where we do this, it does not alter the actual cacae variable astinates will be vhen the clip applied.
since the corresponding entriee
delta = Jar.en. www..anoonn..ttto
classessn code, dtyperv.atyp,, axta=-1) block jop. siieen("arie delta)
deltav by block (mp.einauun"brraa. brivmsbraa"""" dalta, v)
deltav_ by plock_normaiized jap.divida(
boltav_by.an block, bllockk, jap. cippcaippdaltaa block(.... Konal, a.nto=1.0),
nerge.
d_upper_div_love a[0]
b_upper_dit diver b[0]
a[t]]
b_loeer b[i] b_loser
terni jng.mmitiply(
jap (epeeer jpp.clip(Lowe__nev a_ntn=1.0)) Bone],
inp.multiply(
jap dipiddvidette (b. lower, jap. tnp.clip(lower_nng a_min=1.0)) l..., None], upperrr___mm piypprrdd____ae lover,
upper div lovar new taref tam2 return aipor_div_lope Dae, lscer_not
wea_s_ccanntrrre jax. aww.cittreecciiiiieee. ccan
sernerge_fune, www.://r.bbbbbocc..oooooeeeee reverso=Falee, axiast,
by.bckk),
dal by block normal izad, curslative assoc, scan, cattontttttttaa deltal, wlw.blockkooo comulativ
cache_ var wwt_varrr_peeeeeee upper_div diw_eower jap.pad(
deltav by. wtw.:hllooooollla...... nornalizad, curulativee :-2],
((o, 0), 0), (c, (a, 0)),
jap pad[
deltai_ baltsi_b_bblockcoo www.bblocccccsetieeeeee....... :-2), ((o, 0), (2, 0), (0, 0))
return cacha var uppar, div Tover, cache var Tower
Code Jax/Flax pseudocode to get cache variables for all blocks; associative scan version
30


Reviews:
[{"review_id": "VcJd60ytG4", "reviewer": "Reviewer_5aiB", "summary": "* The paper proposes a linear-time attention mechanism based on vector-quantization.\n* This scheme is evaluated quantitatively and qualitatively on various datasets.", "questions": "See weaknesses.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "* The idea of using vector-quantization to achieve linear attention sounds reasonable.\n* The paper includes a detailed discussion of various related works and how they differ from Transformer-VQ.\n* Code is provided, aiding reproducability.\n* The paper provides qualitative samples of the various trained models.", "weaknesses": "* I think the presentation of the paper could be improved. I think some pseudocode or diagrams illustrating the main ideas would be useful, while some of the theorems could potentially be moved to the Appendix.\n* It seems to me that Transformer-VQ only significantly outperforms prior work on ImageNet64 where it uses a 7x larger model than the second best model.\n* It also not entirely clear to me what the real-world advantage of Transformer-VQ is, I assume that the attention mechanism is significantly faster for longer context due to being linear. If so, it would be useful to demonstrate this with some actual speedup experiments.\n* The datasets Transformer-VQ is evaluated on are also not current mainstream LLM datasets (which is the main application of Transformers and attention at the moment); I think results for e.g. standard token-level language modeling on C4 or The Pile would be more convincing.\n\nOverall, while I find the general idea reasonable, I am not fully convinced of it's effectiveness by the current set of experiments, hence I am leaning more towards rejection.", "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": "", "extra_scores": {}}, {"review_id": "Su6PGR8iw5", "reviewer": "Reviewer_1jb5", "summary": "This paper proposes a new attention mechanism that can be realized to be linear in time. The proposed method uses vector quantization (mapping with learnable smooth codebooks) in a way that has non-zero gradients following [van den Oord *et al.* (2017)). Showing that attention as expressed in Definition 3.1. can be realized in linear time is not trivial, so the authors show in Thm. 3.7 how this speed up can be implemented using \"cache\" variables.", "questions": "* In Sec. 5.1.2 the authors say that the codebook size is 256 but Table 6 in the appendix shows S is 512. Is this a typo?\n* What are the authors trying to show with ImageNet64? Is this part trying to show-case the model's applicability as a \"vision transformer\" type of model?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "* Practicality: Instructions on how to implement the architecture is very detailed, and it looks like the authors experimented quite a lot to design an architecture that works well with various architecture sizes (190M, 1.2B, 1.3B parameters)\n* Large-scale experiments across various challenging tasks.", "weaknesses": "The authors experiment with various architecture sizes; how did the authors choose which parameters to scale up? Is there some \"scaling laws\" that the authors observed to be useful? For example, fixing the codebook size to 256 may restrict the model's expressiveness which can be extremely critical for generative models. Given that the authors are emphasizing a decoder model which is particularly useful for generative tasks, I believe this paper would benefit from analysis on this new parameter's effect in addition to Sec. 5.1.1. One type of experiment that could demonstrate this effect would be to see if the model's performance saturates at some parameter count, and see if performance saturation is caused by the codebook size by training the same model with a larger codebook size S. Sharing these would be useful to anyone who plans to use this architecture for custom tasks. \n\nOverall the paper is well written, but I think the paper could benefit if the authors provide intuition for the caching mechanisms and the implementation for efficient attention (Theorems 3.4-3.7). Further, the proofs are written just as a sequence of equations without any conceptual descriptions, so it's quite hard to understand what steps are being used in the proofs.", "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": "", "extra_scores": {}}, {"review_id": "u9wmHfcqlR", "reviewer": "Reviewer_gEU7", "summary": "The paper proposes a novel attention mechanism with linear time complexity. The core idea is to apply vector quantization (VQ) on Keys, leading to the primary computing being (Query x Codebooks), where the number of codebooks is relatively small (256, 512, 1024) compared to the entire sequence length. The proposed method, termed Transformer-VQ, is applicable for both Transformer Encoder and Decoder models. Experimental results in long-context language modeling and image density estimation show that the proposed attention modification achieves performance comparable to prior works.", "questions": "* As the authors mentioned in Section 5.2.1, it seems that simple attention dropout is not applicable; is searching for the appropriate regularization setting a challenging job? In other words, is the model robust enough regarding regularization hyperparameters?\n* The paper uses the GAU Transformer instead of a vanilla Transformer; Is there any reason why you used this variant? If the reason is that GAU employs single-headed attention, is Transformer-VQ applicable to a base (naïve) Transformer architecture?\n* Linear-time Encoder is also explained in Section 3.2 but there are no corresponding experiments.\n* The “truncation-free fixed-size cache” part is confusing. Does this refer to the limitation of the total stored Key size to “LK” (Section 3.4.2)? If so, then where does the ‘truncation-free’ part come?\n* Suggestion: It would be beneficial to include a figure of the computation process (especially visualizing L, C, K, and V) for both the encoder and decoder. The current version only presents mathematical formulations, which can make it difficult to understand the core algorithm.\n* (minor) In Section 3.1, maybe the “gated activation unit” is a typo of “gated attention unit”?\n* (minor) Large “K” seems to be used in two different contexts (key and query block size); consider changing the latter to another letter.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* While the idea of adopting VQ into Transformers is not entirely new (ex: Clustered Attention), the paper’s innovation is in using VQ for Keys and demonstrating its effectiveness (especially for Decoder). This is a nice contribution.\n* The paper provides a comprehensive review of related works, addressing various aspects. Furthermore, the equations and notations are presented clearly to help clarify the core idea.", "weaknesses": "* Although this paper focuses on efficient computation, there is no direct comparison of FLOPs or actual inference latency on GPUs/TPUs with previous works. It would also be beneficial to include a comparison between models with and without VQ.\n* Only one model size (1.3B for PG-19, 1.2B for ImageNet64) is used for experiments. To demonstrate the effectiveness, it would be valuable to test with multiple (smaller) model sizes where the sequence length exceeds the hidden dimension of Transformers. I assume that the actual inference speedup would arise when the sequence length is sufficiently larger than the hidden dimension of Transformers – I’m not sure if the tested models are such a case.", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)