Paper: Transformer-VQ: Linear-Time Transformers via Vector Quantization (ICLR-2024)

Reviewer comment (Experiments):
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.

Author response:
We have updated the paper to include a benchmark of the latency and throughput on a TPU v3 with 8 cores. Our model trains 3x faster at length 8192 than the baseline quadratic-time transformer, and 12x faster at length 32768. This is roughly in line with the ratio T^{2}/TS, since a 4x increase in sequence length led to a 4x increase in compute time. 

In addition, our model scales to a sequence length of 131072 without running out of memory and while maintaining comparable throughput. 

We also sought to include an end-to-end FLOP/s estimate for our model's train_op, or better yet an estimate of the MFU [3]. Unfortunately, Jax's analysis tool for FLOP count [4] is still a work in progress, and it does not return any result on TPUs. On CPUs, the backend does succeed in running an analysis, and the FLOP count is about 2x lower for our model than the baseline at length $T = 8192$. However, the outputted FLOP count estimates are also far below the product of the batch size in tokens and the model's non-embedding/non-codebook parameter counts. We unit tested using the same input batch, and verified all the non-embedding, non-codebook parameters changed after the training step. Since the gradients of the per-token losses must be summed, this shows that that the FLOP count estimate is wrong. As a result, we have not yet included the FLOP/s or MFU estimate in our paper yet. However, we will endeavor to raise this issue with the Jax developers and will update the paper if it is solved in time.

Reviewer comment (Experiments):
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.

Author response:
Actually, that's not true. In our Enwik8 experiment, we noted that our model size is 190M parameters. 

The sequence length always exceeds the hidden dimension. We assume you're referring to the heuristics like $6D_{m}$ or $12D_{m}$ minimum until attention compute "matters" at all. For short sequences of length $T=2048$, the slowdown is negligible, and for  $T=8192$ there is already a 3x speedup, as shown Table 6 of our updated paper. Note that for the $6D_{m}$ heuristic, we have $6D_{m} = 6 * 768 = 4608$ for our 190M parameter model.

Reviewer comment (Experiments):
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?

Author response:
Right: if you apply attention dropout, it will elide all the cached timesteps whose keys land in a particular voronoi cell, and in practice that didn't work well. We found that residual dropout and LayerDrop [5] as a pair with drop rates of about (r, r/2) seem to be an effective alternative to attention dropout and MLP fan-out dropout as a pair with drop rates (r, r).

Searching for optimal regularization is always a nuisance, but it is not particularly more so with our model. One helpful caveat however is that the weight decay setting should be set lower than usual, since the model already happens to have a tiny bit of regularization from the commitment loss.

Reviewer comment (Experiments):
Linear-time Encoder is also explained in Section 3.2 but there are no corresponding experiments.

Author response:
As we clarified above, it was introduced to build intuition and simplify the proofs later. Our paper proposes a new decoder-only model. 
If you want to apply it to encoder-only tasks, sort of like S4 does for their autoregressive model [8], you can, but we have not tried this.

Reviewer comment (Writing):
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?

Author response:
It refers to the fact that the cache variables $U(n-2)$ and $L(n-2)$ are of fixed size, yet together with key block $n-1$ and key block $n$, they contain all the information required to compute quantized-key full attention for query block $n$. 

Truncation-free is a reference to prior works, such as Transformer-XL, which implement a sliding window cache that effectively "truncates" the attention range in any given layer, and ultimately, for the entire model.

Reviewer comment (Presentation):
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.

Author response:
As per our promise, we have updated the paper with a schematic illustration of the VQ-attention approximation, and captioned it in a way that explains VQ-Attention intuitively.

Reviewer comment (Writing):
(minor) In Section 3.1, maybe the “gated activation unit” is a typo of “gated attention unit”?

Author response:
Fixed both of these, thanks.

Reviewer comment (Writing):
(minor) Large “K” seems to be used in two different contexts (key and query block size); consider changing the latter to another letter.

Author response:
Fixed both of these, thanks.

Reviewer comment (Experiments):
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.

Author response:
Thank you for your comments. We have recently uploaded a new version of our paper, which we believe is further improved; see our comment at the top for a summary of changes. We will now address your individual comments.

> 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? 

In general, we tried copying hyperparameters from whatever paper was SoTA/near SoTA, then made adjustments specific to the design of GAU [1]. If the settings were slow or nonperformant, we tried other papers. Two examples that didn't work for us were the Focus Attention paper [2], which only used 39M parameters, and was slow to optimize even when the x-axis was wall clock time, and the 60-layer model for ImageNet64 used by Perceiver-AR [3], which was slower per step than a wide but shallower model, and was also slower to optimize when the x-axis was wall clock time. 

For our small model, we ultimately adapted from Sukhbaatar et al., 2021 [4], but copied it as if the MLP fan-out factor was 4x, whereas their setting was larger. For our large model, we ultimately used a model adapted from Hutchins et al., 2022 [5], who train a PG-19 model with 24 layers and 1.3B parameters. From this, we inferred the model width $D_{m}$ was 2048. When applying GAU, one replaces the vanilla transformer's attention and MLP sublayers with one GAU each, both using value and gate widths $D_{v} = 2D_{m}$, so we simply set $D_{v} = 4096$ in our large model. Moreover, GAU always uses query/key width $D_{k} = 128$, so we used this setting too. This strategy yields a model with roughly the same number of parameters, and in our experiments it worked quite well. The only difference between our 1.3B and 1.2B architectures are the token embeddings, since the PG-19 model has a vocabulary size of 32,000. 

For the codebook hyperparameters, we found EMA rate $\gamma = 0.99$ worked best, which is the same as the original VQ-VAE paper recommended. In addition, we found that using a small commitment loss coefficient $\beta = 0.0001$ worked reasonably well for all of our models, but we did not sweep over this choice when training our larger models.

Reviewer comment (Experiments):
Given that the authors are emphasizing a decoder model which is particularly useful for generative tasks, 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 results would be useful to anyone who plans to use this architecture for custom tasks.

Author response:
> ... 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.

We agree it would be beneficial to be able to have a scaling law to predict the validation loss from the codebook size, non-codebook/non-embedding parameter count, and training token count. Ideally, this could be combined with a prediction of the time per training step, and then we could determine the optimal codebook size yielding the lowest loss under a given wall-time budget for training. 

At present, we feel do not have enough data to claim we have a formal scaling law. However, based on our ablations so far, it does appear that the logarithm of validation loss decreases linearly with the logarithm of the codebook size. In other words, the validation loss w.r.t. codebook size appears to follow a power law, $L(S) = (\frac{S_{0}}{S})^{\alpha}$. Similar to Kaplan et al., 2020 [6] we know this model is misspecified, since it implies a zero loss for an infinite codebook size. Adding a constant term similar to Henighan et al., 2020 [7] and Hoffman et al., 2022 [8] solves this problem, and it implies improvements in the logarithm of the validation loss will eventually saturate. 

That being said, we don't want to fit a log-linear model on the basis of 3 datapoints, especially when they're derived from a small dataset like Enwik8 where overfitting could lead to confounded results. To perform a proper study of this matter, it will be necessary to use larger datasets and many additional training runs. For this reason, we did not add a formal scaling law in our paper.

Reviewer comment (Theory):
The paper could benefit if the authors provide intuition for the caching mechanisms and the implementation for efficient attention (Theorems 3.4-3.7).

Author response:
>  the paper could benefit if the authors provide intuition 

The original idea was that the exact attention weights may not matter a lot, and they can be quantized as scalars. However, compressing the bit width of a float32 attention weight scalar can save at most a 32x constant factor on the attention compute using e.g., stochastic rounding when computing QK^T [9]. In contrast, using vector quantization and compressing the "duplicate" keys allows a linear savings, increasing with sequence length, and has zero variance.

Reviewer comment (Writing):
The proofs are written just as a sequence of equations without any conceptual descriptions, making it quite hard to understand what steps are being used in the proofs.

Author response:
Since you mentioned the original proofs were terse, we went through all of them, added commentary, and generally made them easier to follow. Apologies for the delay, and hope you find the updated proofs informative. 

We have also made many more improvements to the latest manuscript, which are catalogued in our comment at the top of the page.

If there's anything else we can do to help, please let us know!

Reviewer comment (Writing):
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?

Author response:
> 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?

No it's not; at the beginning of Section 5 (which was on the previous page) we wrote that the hyperparameters are given by Appendix B "unless specifically noted". Since we noted the choice of $S = 256$, basically we didn't feel the need to make any additional remark about this difference from our main hyperparameters. 

In Sec 5.1.1 and 5.1.2 we are trying to see the cumulative impact of the choices that weaken the attention mechanism. We start with $S = 1024$, then decrease it by 2x twice, then we remove the compressive cache. If you want, we can also run an ablation removing the compressive cache from the $S = 512$ model. If you do want us to add this ablation, please let us know as soon as you can.

Reviewer comment (Experiments):
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?

Author response:
> 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?

We were trying to see if the model can be used as a generative model on domains other than text. In text, the absolute position of distant tokens tends not to matter a lot, since for example localized re-phrasings can shift the absolute index of subsequent tokens without changing the high-level meaning. In contrast, in images things like the row index and column index matter a lot, and it was not clear to us how well Transformer-VQ would work, since its keys in any given layer are compressed to a smaller set of options than the length of the sequence. 

Our main source of optimism for training on ImageNet64 was Perceiver AR, which also compresses the tokens into a relatively small representation space. Of course, Perceiver AR still takes quadratic time to generate a sequence due to the repeated applications of cross attention from each chunk. 

As an aside, since the parameter count of our ImageNet64 model is not known in comparison to Perceiver AR, we are also training a small model on this dataset right now, with a similar parameter count as Efficient VDVAE. Since this will not impact our main results, we assume it is okay to add this result to the paper.

Reviewer comment (Presentation):
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.

Author response:
We have added some pseudocode to Appendix E, please let us know if this helps clarify things. In addition, we will attempt to create a satisfactory diagram illustrating the main idea.

Reviewer comment (Experiments):
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.

Author response:
The advantage of Transformer-VQ is that it is a potential drop-in replacement for full attention, but is efficient to train on long sequences. 

We have added these experiments in our updated paper. Transformer-VQ runs 3x faster at sequence length 8192 with 190M parameters, runs 12x faster at length 32768, and can scale to length 131072 while maintaining similar throughput and without running out of memory.

Reviewer comment (Experiments):
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.

Author response:
The PG-19 and ImageNet64 datasets we evaluate on are standard long-context modeling benchmarks and are in fact the largest publicly-available autoregressive modeling datasets used in well-cited and spotlighted long-context papers from 2022 [3][4][7][8]. 

As for the Pile, it is a great dataset for training LLMs, yet there are several reasons we didn't use it:
1. At 800GB, it requires significant engineering overhead to download and monetary cost to store. 
2. The download link has been broken for some time [9]. 
3. The average document length is only 5.9 KiB, or 6000 bytes [10], which is at most a few thousand BPE tokens. 
4. Only one of the efficient sequence modeling papers we cited used The Pile, and did not report any perplexity improvement beyond 2^12 (4096) tokens, despite training with an 8k context length [11]. 

As for C4, it also has a size of 800GB, it is mainly used for masked language modeling, and its document lengths again tend to be relatively short. We are only aware of one paper credibly using C4 for long-context autoregressive modeling [8], and they used a filtered version of it with 4k tokens as the minimum document length. Since 4k is still quite short, this dataset is not very interesting for long-context modeling, and as far as we know it has not been used since. 

As we experienced minimal issues with overfitting on PG-19 and ImageNet64, which are more popular benchmarks for long-context autoregressive modeling, we do not believe that evaluating on either of these near-terabyte-scale datasets would yield any particular insights at our current scale, especially given their short average document lengths.

Lastly, please note that PG-19 is a BPE subword-token language modeling task and involves 11GB of text. It is only the perplexity that is converted to word-level for benchmarking purposes, per Rae et al., 2020 [12].