Hi everyone. I want to introduce our new research paper titled “More expressive attention with negative weights”. [2411.07176] More Expressive Attention with Negative Weights
We propose a novel attention mechanism, named Cog Attention, that enables attention weights to be negative for enhanced expressiveness, which stems from two key factors:
(1) Cog Attention can shift the token deletion and copying function from a static OV matrix to dynamic QK inner products, with the OV matrix now focusing more on refinement or modification. The attention head can simultaneously delete, copy, or retain tokens by assigning them negative, positive, or minimal attention weights, respectively. As a result, a single attention head becomes more flexible and expressive.
(2) Cog Attention improves the model’s robustness against representational collapse, which can occur when earlier tokens are over-squashed into later positions, leading to homogeneous representations. Negative weights reduce effective information paths from earlier to later tokens, helping to mitigate this issue.
We develop Transformer-like models which use Cog Attention as attention modules, including decoder-only models for language modeling and U-ViT diffusion models for image generation. Experiments show that models using Cog Attention exhibit superior performance compared to those employing traditional softmax attention modules.
We want to challenge the common belief that attention weights should naturally be non-negative, and we addressed many difficulties such as training instability, numerical overflow, and difficulties in attention normalization due to issues like division by zero, etc.
Our approach suggests a promising research direction for rethinking and breaking the entrenched constraints of traditional softmax attention, such as the requirement for non-negative weights.
Here is an attention pattern figure obtained from our pretrained models. Details can be found in the paper, and hope you find it interesting.
