Paramixer: Parameterizing Mixing Links in Sparse Factors Works Better than Dot-Product Self-Attention (2204.10670v1)

Abstract: Self-Attention is a widely used building block in neural modeling to mix long-range data elements. Most self-attention neural networks employ pairwise dot-products to specify the attention coefficients. However, these methods require $O(N^2)$ computing cost for sequence length $N$. Even though some approximation methods have been introduced to relieve the quadratic cost, the performance of the dot-product approach is still bottlenecked by the low-rank constraint in the attention matrix factorization. In this paper, we propose a novel scalable and effective mixing building block called Paramixer. Our method factorizes the interaction matrix into several sparse matrices, where we parameterize the non-zero entries by MLPs with the data elements as input. The overall computing cost of the new building block is as low as $O(N \log N)$. Moreover, all factorizing matrices in Paramixer are full-rank, so it does not suffer from the low-rank bottleneck. We have tested the new method on both synthetic and various real-world long sequential data sets and compared it with several state-of-the-art attention networks. The experimental results show that Paramixer has better performance in most learning tasks.

• The paper introduces Paramixer, a novel attention mechanism replacing dot-product self-attention by parameterizing sparse matrix factors to avoid the low-rank bottleneck and reduce complexity.
• Paramixer achieves significantly better computational efficiency ($O(N ext{log} N)$ or $O(N ext{log}^2 N)$ vs $O(N^2)$) and empirically outperforms existing Transformer variants on benchmarks involving long sequences.
• By effectively modeling long-range dependencies with full-rank attention, Paramixer offers potential advantages for lossless text data compression compared to traditional or low-rank attention methods.

This paper introduces Paramixer, a novel neural network building block designed to replace the dot-product self-attention mechanism commonly used in Transformers. The authors identify limitations in existing self-attention mechanisms, specifically the quadratic computational cost ($O(N^2)$ for sequence length $N$) and the low-rank constraint imposed by the dot-product, which limits representational power.

Main Contributions:

1. Paramixer: A new attention mechanism that factorizes the interaction (attention) matrix into a product of sparse matrices. Instead of computing attention weights via dot-products, Paramixer directly parameterizes the non-zero entries of these sparse factor matrices using MLPs. This avoids the low-rank bottleneck of dot-product attention.
2. Sparse Factorization Protocols: Two protocols, CHORD and CDIL, are proposed to define the sparsity patterns of the factor matrices. CHORD is based on a peer-to-peer lookup algorithm, and CDIL is inspired by dilated convolutions in Temporal Convolution Networks (TCNs). Both protocols ensure that the product of the sparse factors results in a full-rank, dense matrix, allowing for rich interactions between sequence elements.
3. Computational Efficiency: The proposed methods achieve a computational complexity of $O(N \log N)$ or $O(N \log^2 N)$, significantly improving upon the quadratic cost of standard self-attention. This allows for processing of much longer sequences.
4. Empirical Validation: Extensive experiments on synthetic tasks, the Long Range Arena benchmark, long document classification, and genome classification demonstrate that Paramixer outperforms various Transformer variants (Linformer, Performer, Reformer, Nyströmformer, etc.) in terms of accuracy, especially on tasks involving very long sequences.

Applicability to Lossless Text Data Compression:

From a text compression perspective, Paramixer's ability to model long-range dependencies efficiently and its full-rank attention matrices are highly relevant. Here's a breakdown:

• Modeling Long-Range Dependencies: Text data often exhibits long-range correlations (e.g., a pronoun referring to a noun many words earlier). Traditional compression algorithms, including those based on Lempel-Ziv (LZ) variants, primarily exploit local redundancy. Paramixer's $O(N \log N)$ complexity and its ability to capture long-range interactions without the low-rank constraint of dot-product attention offer a potential advantage. By learning a more complete representation of these dependencies, a model could potentially identify and remove more redundancy than methods focused solely on local patterns.
• Full-Rank Attention vs. Entropy Limits: The paper emphasizes that Paramixer avoids the low-rank bottleneck of dot-product attention. This is crucial. A low-rank attention matrix limits the expressiveness of the model, hindering its ability to capture the complex statistical structure of text. A full-rank matrix, as achieved by Paramixer, has the potential to learn more intricate dependencies, bringing the learned distribution closer to the true data distribution. This, in turn, could lead to better compression, potentially approaching (or even, in theory, surpassing traditional notions of) entropy limits more closely. Entropy limits, such as those defined by Shannon's source coding theorem, represent the theoretical minimum average code length, given a perfect model of the source distribution. However, it represents a theoretical benchmark.
• Comparison with Arithmetic Coding: Arithmetic coding is a powerful entropy coding technique that achieves compression rates very close to the entropy limit given a probability model. The key is the quality of the probability model. Standard self-attention, with its low-rank limitations, might not be able to model the complexities of text as effectively as Paramixer. Paramixer could provide a better probability model to an arithmetic coder, potentially exceeding the performance of standard self-attention approaches.
• Algorithmic Efficiency and Practical Implementation: The $O(N \log N)$ complexity of Paramixer is a significant improvement over the $O(N^2)$ of standard self-attention. While not as efficient as some linear-time algorithms (like LZ variants), it's a reasonable trade-off for the potential gains in modeling long-range dependencies. The practical implementation would likely involve using Paramixer as part of a neural network that predicts the probability distribution of the next character given the preceding context. This probability distribution would then be fed to an arithmetic coder.

Potential Improvements, Limitations, and Future Research:

• Adaptive Sparsity: While CHORD and CDIL provide fixed sparsity patterns, exploring adaptive sparsity patterns (where the sparsity is learned during training) could be beneficial. This might allow the model to focus on the most relevant dependencies for compression.
• Hybrid Approaches: Combining Paramixer with traditional compression techniques (e.g., using it to enhance the context modeling in an LZ-based compressor) could be a fruitful area of research.
• Theoretical Analysis: A deeper theoretical analysis of the relationship between Paramixer's learned representations and the entropy of the data source would be valuable. This could provide insights into how closely Paramixer can approach theoretical compression limits.
• Beyond Character-Level Models: While the paper focuses on character-level models, exploring byte-level or subword-level models with Paramixer could also be beneficial.
• Training data requirements: One limitation of neural network compression techniques is their reliance on a training dataset. This could be improved by having it be adaptive to each text it compresses.

In conclusion, Paramixer presents a promising approach to modeling long-range dependencies in text, which is crucial for lossless data compression. Its computational efficiency and full-rank attention matrices offer potential advantages over existing methods. Future research should focus on exploring adaptive sparsity, hybrid approaches, and a deeper theoretical understanding of its compression capabilities.