Hi! I just wanted to add a wrap-up about our article in this context.
Sparsifying Transformer Models with Differentiable Representation Pooling 9
By Michał Pietruszka, Łukasz Borchmann
The problem of having quadratic complexity w.r.t. the length of the attention mechanism in transformers is approached using pooling operations for reducing the number of word-vectors in between layers. The paper finds that even a hard selection of word-vectors outperforms Linformer and Reformer-based baselines on a long-documents summarization task both in speed and ROUGE scores. However, this hard selection remains suboptimal, as gradients are not propagated to each element in the sequence. This drawback was eliminated by introducing the novel pooling operation, namely The Successive Halving Differentiable Topk. It allows scoring each element in the sequence and selecting a predetermined number of word-vectors that achieved the highest score.
Findings

Word-vector pooling allows achieving sub-linear complexity. Keeping the lower sequence length after the pooling is beneficial for the complexity of the next layers, FFNs, and even the decoder’s cross attention. More vectors can be eliminated in subsequent layers, further decreasing complexity as in the Pyramidion model.
Massive saving on memory and time (16x and 3.5x respectively) are achieved while outperforming dense baselines at the same time. The time overhead for scoring and pooling is minimal, but the elimination of some information redundancy improves the training significantly.
The best models were reusing a part of the saved computations for deepening the network.

Follow-up questions

The proposed Successive Halving Top-k operator is universally applicable. How do you want to use that in other fields? What are specific examples of tasks and model architectures?
How can other methods (e.g., Linformer) benefit from keeping the lower number of word-vectors?

I hope you will find it interesting!