core/nn/aft_transformer library
Attention-Free Transformer (AFT) block + language-model stack.
Same pre-LN skeleton as the vanilla TransformerBlock, but the
self-attention sublayer is AFTAttention rather than
MultiHeadAttention. No cross-attention (decoder-only).
AFTLanguageModel glues this together with token/positional
embeddings and a linear LM head — mirrors TransformerLM's API so
the two are drop-in comparable in training loops and demos.
CPU-only for now (transitively — AFTAttention is CPU-only).
Classes
- AFTBlock
- AFTLanguageModel
- Decoder-only language model built from AFT blocks. Analogous to TransformerLM but with attention-free self-attention.