StableMask: Refining Causal Masking in Decoder-only Transformer
This work addresses inefficiencies in attention mechanisms for language modeling, offering a method that enhances model performance and supports extrapolation, though it appears incremental as it refines existing causal masking rather than introducing a new paradigm.
The paper tackled limitations in decoder-only Transformers, where causal masking forces non-zero attention and relative position encoding lacks absolute positional information, by proposing StableMask, a parameter-free method that refines the causal mask to balance attention distributions and encode absolute positions, resulting in significant performance improvements in language models from 71M to 1.4B parameters across diverse datasets.
The decoder-only Transformer architecture with causal masking and relative position encoding (RPE) has become the de facto choice in language modeling. Despite its exceptional performance across various tasks, we have identified two limitations: First, it requires all attention scores to be non-zero and sum up to 1, even if the current embedding has sufficient self-contained information. This compels the model to assign disproportional excessive attention to specific tokens. Second, RPE-based Transformers are not universal approximators due to their limited capacity at encoding absolute positional information, which limits their application in position-critical tasks. In this work, we propose StableMask: a parameter-free method to address both limitations by refining the causal mask. It introduces pseudo-attention values to balance attention distributions and encodes absolute positional information via a progressively decreasing mask ratio. StableMask's effectiveness is validated both theoretically and empirically, showing significant enhancements in language models with parameter sizes ranging from 71M to 1.4B across diverse datasets and encoding methods. We further show that it naturally supports (1) efficient extrapolation without special tricks such as StreamingLLM and (2) easy integration with existing attention optimization techniques.