StableMask: Refining Causal Masking in Decoder-only Transformer
Fuente:
arXiv
Saved in:
| Main Authors: | , , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2024
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866913225814048768 |
|---|---|
| author | Yin, Qingyu He, Xuzheng Zhuang, Xiang Zhao, Yu Yao, Jianhua Shen, Xiaoyu Zhang, Qiang |
| author_facet | Yin, Qingyu He, Xuzheng Zhuang, Xiang Zhao, Yu Yao, Jianhua Shen, Xiaoyu Zhang, Qiang |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2402_04779 |
| institution | arXiv |
| publishDate | 2024 |
| record_format | arxiv |
| spellingShingle | StableMask: Refining Causal Masking in Decoder-only Transformer Yin, Qingyu He, Xuzheng Zhuang, Xiang Zhao, Yu Yao, Jianhua Shen, Xiaoyu Zhang, Qiang Computation and Language Artificial Intelligence 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. |
| title | StableMask: Refining Causal Masking in Decoder-only Transformer |
| topic | Computation and Language Artificial Intelligence |
| url | https://arxiv.org/abs/2402.04779 |