Tokens, who needs them?

Improving transformer memory efficiency by 2.7x

Introduction

Transformer

Transformers are memory and compute-intensive architectures because of attention.

Attention scales as \(O(n^2)\) due to the all-vs-all token operation. This causes both memory and compute to explode. Whilst the compute can often be justified by enhanced performance, the memory and throughput drawback is significant. It would be great if we could harness the power of attention whilst minimising these bottlenecks.

Attention compute grows quadratically with sequence length, so longer contexts rapidly increase cost.

One way might be to take inspiration from the visual processing field. The U-Net architecture originates from there, and shows that you can save on computation and memory by using compression as part of the architecture. This allows the model to compute in a more abstract but reduced dimensionality space.

U-Net

The U-Net was originally designed for biomedical image segmentation but proved to be a very powerful architecture for many vision tasks. The key architectural innovation was the reasoning by compression. Instead of computing in the high-dimensional pixel space, the U-Net compresses the input layer by layer before expanding it back to its original size. Those lower-dimensional layers also learn to capture more abstract aspects of the input space. The second innovation was the use of skip connections. Skip connections allowed information to pass through the model without requiring it to go through the 'U', the compressive pathway. This means that any high-frequency detail that may have been lost during compression can be recovered.

The efficiencies of the U-Net come from the fact that the dimensionality of the problem is reduced during compression. This means that the total number of FLOPs for a forward pass is much reduced with respect to if each layer held the original image's (C, H, W) dimensionality.

The question is, how can we change the transformer architecture to maintain causality during this compressive process?

Simple U-Transformer

Causality and Downsampling

Language modelling transformers are causal. This means that each token can only attend to all previous tokens. Naively, one could achieve layer-wise compression by just averaging groups of tokens. However, this would break causality, because it means that future token information would begin mixing with past token information during the downstep.

Thus, in order to maintain causality, I decided to take the simplest route, which was to just drop tokens.

Every layer would have a compression factor, e.g. 2x would mean that half the number of tokens would be dropped. By dropping them, causality is surely maintained as there is no chance for information mixing.

Concretely, for 2x compression at a given layer, we keep token positions 0, 2, 4, ... and drop positions 1, 3, 5, ... before the next layer.

My concerns at this stage were that dropping tokens would result in catastrophic information loss, but I was not deterred because of two things:

• Causal attention before dropping means group 'leaders' can absorb information from the tokens behind them before they are dropped.
• Skip connections mean that missing information can be recovered later.

Upsampling

Once tokens are dropped, how can they be recovered again during upsampling so that informaiton can be mixed from the skip connection and from the compressed representation? For this I adopted a strategy analogous to pixel shuffle. Each upsampling layer has a linear head that can expand a tokens dimensionality. For example, in a 2x expansion layer, a token of dims D, will be expanded to 2D. Then simply, it is split in half to create two dim D tokens. The skip connection to the previous layers will then add to these tokens, recovering any lost fine detail from the downsampling.

In practice, each kept token is projected from D to 2D, split into two D vectors, and then aligned with skip features at matching positions so the full token count is reconstructed.

Tokens are dropped by index during compression, then re-expanded by D->2D split and skip-merged to reconstruct full sequence length.

Full model

The full model looked like this:

Before and after each downsampling or upsampling block there would be a regular causal attention block.

Experiment

In order to test whether this U-Net transformer would be performant, I decided to train on enwik8, a Wikipedia excerpt dataset. The task would be simple next-token prediction using byte-level tokenisation. I would then also train two baselines:

• Parameter-matched baseline.
• Compute-matched baseline (so fewer parameters).

As standard, the first 90% of enwik8 was used for training, the next 5% for validation and selecting the best checkpoint, and the final 5% for test evaluation.

These experiments would allow me to understand the different trade-offs:

• Performance trade-off with equal parameters.
• Compute cost and speed with equal parameters.
• Performance with equal compute.

Architectures

The parameter-matched and compute-matched baselines use the same core architecture as the U-Net model, but without the upsampling and downsampling pathway. The compute-matched baseline uses fewer parameters overall.

Results

Parameter and Compute vs Test Loss comparison

Relative tradeoffs versus UNet: compared to baseline_small and baseline, we show differences in inference compute, test performance (2^Δbpb−1), single-forward memory allocation, and train compute to best validation (all as % vs UNet).

Despite having ~4.1% less compute than the compute-matched baseline, the baseline is only 2.25% better (\(2^{\Delta \text{BPB}} - 1\)). This shows that for approximately the same compute, the U-Net does not deteriorate in performance.

On the other hand, the U-Net uses 2.7x less memory than the compute-matched model, and 3.54x less memory than the parameter-matched model.

The parameter-matched model shows 2.9% better performance, but uses significantly more compute than both the compute-matched baseline and the U-Net.

The similarity in performance between the compute-matched and parameter-matched baselines suggests diminishing returns from additional compute on this dataset. In the lower-compute regime where the compute-matched baseline and U-Net operate, small compute differences appear to produce larger performance differences. In the higher-compute regime, larger compute differences produce much smaller performance gains. This likely explains why the small compute gap between U-Net and compute-matched baseline is still visible in performance, while the much larger compute gap from compute-matched to parameter-matched yields a comparatively modest improvement.

Bottleneck Analysis

In theory, if the model is using compression to help calculate the next token, the bottleneck dimension should contain easily identifiable abstract information about the input sequence. To see whether this is the case or not, I investigated:

• From the test set I created non-overlapping sequences.
• I pass the sequences through the U-Net model.
• I extract the bottleneck tokens during the forward pass.
• Once all have been collected I visualise the bottleneck tokens with UMAP and cluster them (HDBSCAN).
• For each cluster, I see whether there are any words enriched in that cluster vs the rest of the dataset.

As you can see from playing with the above demo, the bottleneck embeddings capture real semantic information with clear themes per cluster. For example, cluster 25 is clearly on religious topics. Meanwhile, clusters 22 and 21 are mathematical and 9 is computing.

This is strong evidence that the simple unet model is actually utilising the bottleneck as an abstract reasoning space.

Training Compute Curves

Validation bpb against cumulative training compute, stopping at best checkpoint

Baseline, being the largest compute model, converged with the least training compute and to the lowest loss. This is quite consistent with empirical results from the literature about how scale generally leads to faster convergence and better performance. The compute-matched U-Net and baseline_small take about the same amount of compute to converge to their best values, though it is interesting to note that the U-Net's curve is quite different from the traditional curves, exhibiting a kink around the 150 PFLOPs mark where the trajectory changes significantly.

Live Demo

You can try the model live below. Click 'compare' and scroll down in the widget to see the difference in memory allocation curves.

Conclusion

A simple U-Net style token compression strategy can reduce transformer memory allocation substantially without a large quality drop. In these experiments, the U-Net achieved 2.7x lower inference memory than the compute-matched baseline and 3.54x lower than the parameter-matched baseline, while staying close in test BPB. The bottleneck analysis also shows that compressed representations retain meaningful semantic structure, supporting the idea that this bottleneck acts as a useful abstract reasoning space. Overall, this suggests there is practical headroom in token-space compression for more memory-efficient language modelling.

Comments