This work represents progress on removing attention head superposition. We are excited by this approach but acknowledge there are currently various limitations. In the short term, we will be working on adjacent problems are excited to collaborate with anyone thinking about similar things!

Produced as part of the ML Alignment & Theory Scholars Program - Summer 2023 Cohort

Summary: In transformer language models, attention head superposition makes it difficult to study the function of individual attention heads in isolation. We study a particular kind of attention head superposition that involves constructive and destructive interference between the outputs of different attention heads. We propose a novel architecture - a ‘gated attention block’ - which resolves this kind of attention head superposition in toy models. In future, we hope this architecture may be useful for studying more natural forms of attention head superposition in large language models.

Our code can [...]

---

Outline:

(01:10) Background

(03:27) Attention head superposition for OV-Incoherent Skip Trigrams

(06:21) Removing Attention Head Superposition from models trained on OV-Incoherent Skip Trigrams

(06:29) Toy Dataset Generation

(07:18) Gated Attention Mechanism

(12:17) Gated Attention Block Training

(13:51) Summary of Architectural and Training Choices

(15:10) Experimental Results on Toy Model Setup

(19:12) Related Work

(20:14) Conclusion

(21:00) Limitations

(25:32) Future Work

(27:42) Acknowledgements

(28:08) Contributions Statement

(28:46) Appendix

(28:49) Appendix A: OV-Incoherent skip trigram dataset generation algorithm

(29:03) Appendix B1: 2 Heads, 3 Skip Trigrams (Trained via Stochastic Gradient Descent)

(29:47) Appendix B2: 3 Heads, 4 Skip Trigrams (Trained via Stochastic Gradient Descent)

(30:32) Appendix C: Other approaches that we investigated

(32:31) Appendix D: Weights for Manually Specified Model encoding 3 OV-Incoherent Skip Trigrams with 2 heads

---

First published:
April 8th, 2024

Source:
https://www.lesswrong.com/posts/kzc3qNMsP2xJcxhGn/gated-attention-blocks-preliminary-progress-toward-removing-1

---

Narrated by TYPE III AUDIO.

“Gated Attention Blocks: Preliminary Progress toward Removing Attention Head Superposition” by cmathw, Dennis Akar, Lee Sharkey

LessWrong (30+ Karma)