FiX: Introducing Fine-grained Forget Gate into Softmax Attention

Causal softmax attention is the algorithmic foundation of modern large language models. Inspired by linear attention, recent work has sought to enhance it by incorporating forget gates. However, these efforts, such as FoX, have been limited to coarse, scalar gates. While fine-grained, element-wise gates are shown to be more effective than scalar ones in linear attention, their direct integration into softmax attention is non-trivial due to algebraic constraints. In this work, we introduce Fine-grained Forgetting Transformer (FiX), a novel architecture that successfully enables element-wise forget gates in softmax attention. Our core insight is that the softmax denominator becomes mathematically redundant under a subsequent RMSNorm layer, allowing us to reformulate the forgetting mechanism as a direct element-wise multiplication on the value vectors. This formulation makes FiX the first positional encoding applied to value-output (VO) pairs, designed to be complementary to existing query-key (QK) encodings like RoPE. We systematically address implementation challenges including numerical precision, computational efficiency, and inference memory consumption. Extensive experiments show FiX achieves lower training loss and superior performance on both short-text common sense benchmarks and long-context tasks, opening a new path for building more powerful transformers. Our code is available at https://github.com/AlayaDB-AI/FiX.

Large language models like ChatGPT are powered by a mechanism called "attention", which helps the model decide which words in a sentence are relevant to each other. As these models process longer and longer texts, they need a way to selectively forget outdated information. This is much like how humans don't remember every word of a conversation, only the key points. Current forgetting mechanisms are crude: they treat all types of information the same, like erasing an entire whiteboard instead of just the outdated parts. We discovered a mathematical trick that enables a much more precise form of forgetting. This precise form can independently control what to keep and what to discard for each piece of information. The key insight is that a common normalization step used in modern language models accidentally makes a complicated part of the attention calculation unnecessary, freeing us to apply fine-grained forgetting directly. Our new architecture, FiX, is the first to apply this precise forgetting mechanism to the attention component. Experiments show that FiX trains more efficiently and performs better on a wide range of language tasks, from answering common sense questions to understanding very long documents. This opens a new direction for building more capable and efficient language models.