Introduction

Motivation and Problem Statement

Transformer architectures are the dominant paradigm for sequence modeling in natural language processing , computer vision, and computational biology, with models such as BERT and GPT establishing the standard for representation learning. The field has since diversified rapidly: recurrent and retentive blocks, sparse attention, gated mechanisms, latent projections, memory-augmented layers, fast-weight rules, and continuous-depth formulations each address a different limitation of dense attention, while the optimizer landscape has expanded beyond AdamW with variance-reduction, memory-efficient, schedule-free, second-order, and low-rank methods .

The practical consequence is a fragmented ecosystem. Comparing sequence mixers or optimizers under identical training conditions requires re-implementing each variant, wiring it into a training pipeline, and maintaining the result across framework changes. This effort hinders reproducibility, slows scientific progress, and raises the barrier to entry for new researchers.

Proposal and Contributions

We propose Frankenstein Transformer (https://github.com/erickfmm/frankenstein-transformer), a configuration-driven toolkit that treats architecture, optimizer, and training choices as a single validated document. The project provides:

1. Schema-first design: a strict YAML contract supporting fifty-six sequence mixers across nine families and thirty-eight optimizer families, enforced before training begins.

2. Architecture-agnostic training: one pipeline for dense baselines, recurrent and retentive blocks, sparse patterns, gated mechanisms, latent compression, memory layers, fast-weight attention, and geometric-field attention, with encoder and decoder modes .

3. Optimizer routing: prefixed hyperparameter groups that control embeddings, normalization layers, attention blocks, and remaining parameters independently for each optimizer family .

4. End-to-end workflows: training (including sentence-embedding objectives), deployment and export (quantized, standard, HuggingFace safetensors at 8/4/2 bits, and BitNet GGUF formats), inference, and a schema-driven web builder inspired by SBERT-style workflows .

5. Vision support: a Vision Transformer model class covering masked patch prediction, classification, and segmentation under the same mixer and optimizer infrastructure.

6. Validation tooling: unit tests, YAML preset validation, and continuous integration across supported Python versions.

The command surface is frankenstein-transformer, with four subcommands: train, deploy, infer, and web-server. The last launches a Streamlit configuration builder that renders schema metadata as forms, validates input live, previews the resulting YAML, and generates the equivalent CLI command, serving both as a configuration aid and as documentation of the available design space; a hosted instance runs at https://erickfmm.github.io/frankenstein-transformer/index.html.

Reading Guide

This document is organized as a technical reference:

1. Prerequisites: Appendix 12 introduces transformers and attention for readers new to the field.

2. Configuration contract: Section 3.1 describes the schema and Section 3.2 its validation rules.

3. Architecture selection: Section 4.1 compares sequence-mixer families; Appendices 13, 14, 15, and 16 collect the supporting literature, with normalization details in Appendix 21, BitNet quantization in Appendix 22, and conditional-memory details in Appendix 24, and manifold-constrained residuals in Appendix 27.

4. Optimization: Section 5 summarizes optimizer routing; Appendix 20 provides the full family analysis.

5. Deployment and inference: Sections 6 and 7 cover quantized deployment and sentence embeddings.

Related Work

This work sits at the intersection of deep learning frameworks, training tooling, and configuration-driven experimentation. While Sections 4.1 and 5 survey the architectural and optimization literature in depth, this section focuses on the software ecosystem for building, training, and configuring transformer models.

Base Deep Learning Frameworks

The modern transformer research ecosystem rests on several foundational frameworks that provide automatic differentiation, GPU acceleration, and high-level model-building APIs.

PyTorch has become the dominant framework for transformer research. Its imperative programming model and dynamic computation graphs enable flexible debugging and rapid prototyping of novel architectures. The torch.nn.Module abstraction, automatic mixed-precision training via torch.cuda.amp, and the torch.compile JIT compiler form the backbone of most contemporary transformer implementations. Frankenstein Transformer is built entirely on PyTorch, leveraging its module system for composable architecture definitions.

TensorFlow pioneered production-oriented deep learning with its static computation graph paradigm and the Keras high-level API. While its influence in transformer research has waned relative to PyTorch, TensorFlow’s serving infrastructure (TF Serving, TF Lite) and the tensorflow-text library remain relevant for deployment scenarios.

JAX offers a functional programming model built around composable transformations (jit, vmap, grad) and XLA compilation. Its design enables clean implementations of complex algorithms and has been adopted by several recent transformer libraries (e.g., Flax, Haiku). JAX’s stateless design and automatic vectorization are particularly well-suited for research on novel attention mechanisms and optimization algorithms.

scikit-learn provides classical machine learning algorithms and utilities (preprocessing, metrics, model selection) that serve as baselines and complementary tools in transformer workflows. While not designed for deep learning, its consistent API and extensive documentation make it a standard reference point for comparing transformer-based approaches against traditional methods.

Out-of-the-Box Training Frameworks

Several libraries provide complete training pipelines that abstract away the complexity of distributed training, data loading, and hyperparameter management.

Hugging Face Transformers is the de facto standard for accessing, fine-tuning, and sharing pretrained transformer models. Its Trainer API, model hub with over 500,000 models, and tight integration with the datasets and tokenizers libraries have made it the primary entry point for transformer-based NLP. However, its architecture is optimized for fine-tuning existing pretrained models rather than experimenting with novel sequence mixer designs or training from scratch with custom optimization strategies.

Oumi (https://github.com/oumi-ai/oumi, 2025) is an open-source end-to-end platform for training, fine-tuning, evaluating, and deploying foundation models. It provides a unified interface across multiple cloud providers and supports both supervised fine-tuning and preference optimization (DPO, RLHF). Oumi emphasizes production readiness with built-in monitoring, checkpointing, and deployment tooling.

LLaMA Factory (https://github.com/hiyouga/LLaMA-Factory, 2024) offers a unified framework for efficient fine-tuning of over 100 large language models. It provides both a web-based UI and a CLI, supporting LoRA, QLoRA, full-parameter fine-tuning, and various alignment techniques. Its focus is on adapting existing pretrained LLMs to downstream tasks rather than architectural experimentation.

Axolotl (https://github.com/axolotl-ai-cloud/axolotl, 2024) is a streamlined fine-tuning tool supporting LoRA, QLoRA, DeepSpeed, and FSDP. It emphasizes configuration-driven workflows via YAML files and has become popular in the open-source LLM fine-tuning community. Like LLaMA Factory, it targets fine-tuning of pretrained models rather than training novel architectures from scratch.

Low-Code and Efficiency-Oriented Tools

A growing category of tools prioritizes accessibility and computational efficiency, lowering the barrier to transformer training.

Unsloth (https://github.com/unslothai/unsloth, 2024) provides low-code fine-tuning with 2–5\(\times\) speedup and significantly reduced VRAM requirements through hand-written CUDA kernels and optimized Triton operations. It supports GGUF export for quantized deployment and integrates with the Hugging Face ecosystem. Unsloth’s focus is on making fine-tuning faster and more memory-efficient, not on architectural exploration.

dashAI (https://github.com/DashAISoftware/dashAI, v0.9.7.post2, 2026) is an open-source, local-first desktop workbench for machine learning developed by a Chilean academic consortium (Universidad de Chile/FCFM, CENIA, IMFD, and ANID funding). Its interface is fully schema-driven: configuration forms are generated directly from Pydantic component schemas with no hand-written frontend code, and its catalog spans roughly 76 multitask models—tabular classifiers, NLP models, translation, open-weight LLMs, and text-to-image generation—unified behind a dozen typed base abstractions. Extensibility is central to its design: plugins are ordinary PyPI packages discovered through the dashai.plugins entry-points group and installable from the UI, a mechanism comprehensive enough that even dashAI’s own generative capabilities (e.g., Stable Diffusion, Flux) are distributed as plugins. Where earlier releases targeted classical ML workflows, the current platform spans predictive and generative modeling across domains, making it a natural host for transformer components via its plugin system; Section 8 presents the dashai-frankenstein adapter built on this mechanism.

Configuration-Driven Experimentation

A common thread across the tools surveyed above is their focus on fine-tuning existing pretrained models. Hugging Face Transformers, LLaMA Factory, Axolotl, and Unsloth all assume a pretrained checkpoint as the starting point. Oumi supports training from scratch but targets standard architectures. None of these tools provide a systematic framework for experimenting with novel sequence mixer architectures, comparing diverse optimization algorithms under controlled conditions, or exploring the interaction between architectural choices and training dynamics.

Frankenstein Transformer fills this gap by providing a schema-first configuration system purpose-built for architecture experimentation. Its key differentiators include:

• 50 sequence mixers spanning dense attention, recurrent/retentive blocks, sparse attention patterns, gated mechanisms, latent compression, memory layers, fast-weight attention, and geometric-field attention, all selectable via a single YAML field.

• 23 optimizer families with a prefixed hyperparameter routing system that enables per-parameter-group control across embeddings, normalization layers, attention blocks, and feed-forward networks.

• Dual training modes supporting both encoder-style masked language modeling (MLM) and decoder-style autoregressive (AR) next-token prediction, with specialized fine-tuning workflows for each.

• Strict schema validation via additionalProperties: false constraints that prevent configuration errors before training begins, ensuring reproducibility across experiments.

• End-to-end workflows spanning quantized deployment (ternary weight packing, INT8 activations) and sentence-embedding training inspired by SBERT .

• Web-based configuration interface providing schema-driven form rendering with inline documentation, real-time validation, and CLI command generation.

By decoupling architecture definition from training infrastructure through a validated configuration contract, Frankenstein Transformer enables systematic comparison of sequence mixer families and optimization strategies under identical training conditions—a capability not provided by any existing tool in the ecosystem.

System Design and Architecture

Configuration-Centric Architecture

Experimentation in this project is schema first. Rather than exposing loosely checked flags, model topology, optimizer family, training limits, and telemetry pass through a single validated configuration document, which reduces ambiguity when reproducing results and lets many architectures share one operational interface.

The authoritative contract is , a modular JSON Schema whose modules under cover the model, model class, tokenizer, optimizer, SBERT, training, and dataset blocks. The top-level objects are base_model, tokenizer, model_class, model, training, image, vision_dataset, audio, audio_dataset, text_dataset, inference, and plugins.

Model class.

The model_class field selects among three variants:

• frankenstein_text_encoder: mixed-architecture encoder with the full range of sequence mixers, MoE, and advanced normalization, trained with masked language modeling.

• frankenstein_text_decoder: autoregressive causal decoder for next-token generation; the runtime forces decoder masking.

• frankenstein_vit: Vision Transformer over image patches for masked patch prediction, classification, and segmentation; the runtime forces encoder masking and the configuration requires the image and vision_dataset blocks.

Training mode.

model.dims.mode controls masking: encoder uses bidirectional attention for masked language modeling, and decoder uses causal masking for autoregressive prediction.

Declarative component registry.

The runtime configuration mirrors the validated document one-to-one: the nested model: block is materialised as a typed configuration tree and per-component knobs are bound through the component’s declarative record, with no nested-to-flat rewriting stage. Every pluggable component—sequence mixer, optimizer, normalization, activation, residual strategy, positional encoding and token-embedding variant—is described by one registry entry holding its name, lazily imported implementation, parameter schema fragment, runtime defaults, capability flags (KV-cache support, training-time sparsity, QK-normalization support, gated-FFN status) and bilingual documentation. The schema enumerations, per-mixer parameter objects, optimizer prefix reference and conditional prefix rules are assembled from that registry before validation, so a component is declared once and cannot drift across parallel lists; the closed-world validation guarantee (additionalProperties: false) holds for every registered component. The same mechanism is the extension surface: a plugin module registers new components through a public authoring API, either listed under the configuration’s plugins: key or discovered through the frankenstein.plugins entry-point group. Plugin-declared components therefore validate strictly, appear in the schema-generated web builder, and are dispatched by the same core code as the built-ins, while a plugin that does not declare KV-cache support simply falls back to the full-recompute path (decision D5).

Mixer selection.

model.dims.layer_pattern is an ordered list of mixer identifiers drawn from the eight families described in Section 4; the schema fixes the accepted identifiers, and the complete enumeration appears in Figure 2 and Appendix 38.

Optimizer selection.

training.optimizer.optimizer_class selects one of thirty-eight optimizer families, summarized in Section 5 and detailed in Appendix 20.

Schema Scope and Validation Rules

Top-level and nested objects set additionalProperties: false, so unknown keys fail fast instead of being silently ignored. Optimizer parameters are further constrained by per-family prefix rules. Cross-component validation rejects incompatible combinations before training: hidden size must be divisible by the number of heads, the number of key–value heads must divide the number of heads, vocabulary size must match the tokenizer, a base model requires a task and (for MLM) a tokenizer, each task requires its matching optimizer block, BitNet router flags require BitNet to be enabled, and the decoder and vision classes force their respective masking modes.

Configuration Schema

The complete schema—fields, types, ranges, and validation rules—is documented in Appendix 37. Looped blocks are configured through num_layers and num_loops, giving num_layers\(\times\)num_loops logical layers at fixed parameter count.

Training Tasks

training.task selects the objective: mlm (masked language modeling), causal_lm (autoregressive next-token prediction, decoder class only), sbert (sentence embeddings), text_classification (encoder class only), and the vision tasks patch_prediction, classification, and segmentation (Vision Transformer only). NLP tasks consume text_dataset; vision tasks consume vision_dataset. Sentence-embedding workflows are described in Section 7.

Training Safety and Runtime Semantics

The training step accumulates gradients over micro-batches, clips the global norm, checks for post-clip overflow, retries a bounded number of times on non-finite values, dispatches the configured scheduler, and rotates rolling and best checkpoints. When training.test_percentage is positive, a deterministic fraction of the dataset is held out before training and the final weights are evaluated on it; the loss and task accuracy (character error rate for ASR) are written to test_metrics.json. A single training.seed (default \(42\)) controls the run’s randomness: it seeds the Python, NumPy, PyTorch and CUDA generators, model weight initialization, DataLoader shuffle generators and workers, and the held-out permutation, and it enables deterministic cuDNN, so repeated runs on the same hardware and software stack reproduce initialization, batch order and split. The full pseudocode and the split formulation are deferred to Appendix 24.

Normalization Variants

Normalization is both a modeling choice and a schema compatibility question. The accepted norm.type values are layer_norm, dynamic_tanh, derf, rms_norm, prms_norm (partial RMSNorm), and flash_norm (weightless RMSNorm). Their formulations and comparison are detailed in Appendix 21.

Residual-Stream Architectures

Residual connections are exposed as a macro-design choice through the model.residuals block:

• Standard: the classical identity residual .

• None: no skip connection, retained only as an ablation probe.

• Attention Residuals (full_attn or block_attn): learned attention over depth replaces the fixed residual sum, either over all previous layers or over block-wise summaries .

• Manifold-constrained hyper-connections (mhc.enabled): an \(n\)-stream residual whose mixing matrix is constrained to the Birkhoff polytope .

Attention Residuals compose with mHC and Mixture-of-Depths; the formulations, wiring, and reported results are collected in Appendices 28 and 27.

Architecture Taxonomy and Implementation

Attention and Sequence-Mixer Families

The system implements fifty-six sequence mixers organized into nine functional categories that reflect how each method balances expressivity, computational cost, and memory.

1. Dense baselines (6): standard softmax attention, sigmoid attention, Continuum Attention, Grouped-Query Attention (GQA), Exclusive Self Attention (XSA), and Multi-Token Attention (MTA) provide full global routing at quadratic cost. XSA keeps the standard softmax block and strips the component of each head’s output along the token’s own value vector , a parameter-free exclusion that sharpens the division of labor between attention and the feed-forward path.

2. Recurrent and retentive (3): RetNet, Mamba, and ODE-style blocks maintain state for constant-time inference through recurrence, selective state spaces, or continuous depth.

3. Sparse patterns (11): Sparse Transformer, Longformer, SWA with attention sinks, BigBird, SparseK, NSA, SpargeAttn, FASA, MSA, SparDA, and LSA reduce quadratic cost through structured sparsity, token selection, training-free pruning, or semantic key clustering.

4. Gated mechanisms (9): GLA, DeltaNet, Gated DeltaNet, Gated DeltaNet-2, HGRN2, FoX, Gated Softmax, KDA, and First Passage Attention (FPA) introduce data-dependent control over retention, forgetting, and update strength.

5. Latent compression (13): MLA, GQLA, MLRA, Tucker, KPA, TPA, MFA, IHA, GTA, MTLA, CCA, CCGQA, and GMA compress key–value state into low-rank bottlenecks or route through a Gaussian-mixture latent memory.

6. Memory (3): Titans pairs a local attention window with a surprise-driven associative memory updated at test time, Engram provides test-time read/write memory through learned N-gram hash lookups, and Memory Attention replaces the dedicated value projection with a layer-specific token-ID table.

7. Fast-weight attention (6): Falcon-1/2/3 and Falcon-1A/2A/3A treat the recurrent state transition as an online continual-learning rule .

8. Geometric field (1): SSOG applies a structured geometric prior over the attention matrix .

9. Convolution-augmented (2): the Conformer block fuses multi-head self-attention with a depthwise-convolution module between two Macaron feed-forward layers, giving every layer a local branch alongside the global receptive field; Fast Conformer keeps that design for long-form speech and replaces global self-attention with a limited local window plus a head-pinned global token (Appendix 19).

Dense and Recurrent Families

The dense baselines—standard softmax , sigmoid attention , and Continuum Attention , which replaces the softmax with exact integration of an exponential relevance density over piecewise-linear key/value trajectories—and the recurrent family—RetNet , Mamba , and ODE-style blocks —define the expressiveness-versus-efficiency trade-off. Formulations, parallel and recurrent forms, pseudocode, and a comparison table are collected in Appendix 13.

Memory Family

The memory family groups three mechanisms that add storage beside the attention operator. Titans runs dense attention over a local window while a surprise-driven associative memory is updated at test time; this repository implements the attention window only, and the neural-memory module is not part of the implementation. Engram retrieves N-gram embeddings through per-order hash tables and gates them into the residual stream. Memory Attention keeps the standard softmax attention operator but removes the dedicated value projection: values are the contextual keys plus a layer-specific, token-ID-indexed memory row, \(V = K + \mathrm{Norm}(E[s])\), with RoPE applied only to the scoring pair. It increases parameter capacity without proportional dense computation and retains a conventional KV cache at inference.

Latent and Low-Rank Attention Family

This family compresses per-token key–value state into a low-rank latent: MLA, GQLA, MLRA, Tucker, KPA, TPA, MFA, IHA, GTA, and MTLA unify grouped-query, latent, Kronecker- and CP-factorized, and temporal designs. CCA and CCGQA formulate attention in a compressed latent with decoupled grouping, while GMA replaces pairwise interaction with probabilistic routing through learned Gaussian-mixture components. Formulations, pseudocode, and a comparison table are collected in Appendix 14.

Sparse Attention Extensions

The registry includes twelve sparse blocks . Two inference-time methods, fasa and sparge, apply their sparse heuristic only when the model is in evaluation mode: during training the enclosing hybrid layer substitutes standard dense attention over the same \(\mathbf{q}/\mathbf{k}/\mathbf{v}/\mathbf{o}\) projections, so the trained weights are exactly the ones the sparse path reads at inference. swa is also training-free but is a normal trainable mixer, so it runs its windowed pattern in both training and inference. LSA is the family’s semantic selector: keys are pooled and Ward-clustered per 512-token block, and each query scores the cached bucket centroids instead of every key, selecting top-\(p\) buckets that expand into its \(k\) long-range keys on top of sinks and a one-block window. ProbSparse inverts the axis: it ranks queries with a max–mean approximation of the KL divergence against the uniform distribution over a sampled key subset and scores only the top-\(u = c_2\lceil\ln L\rceil\) dominant ones, which brings the layer to \(\mathcal{O}(L\log L)\) for long bidirectional encoders while leaving decoding dense. Formulations, pseudocode, per-block characteristics, and a comparison table are collected in Appendix 15.

Gated and Fast-Weight Attention

The gated family contains nine blocks . Gating controls what information survives, either in recurrent state updates (GLA, DeltaNet variants, HGRN2, KDA) or in the full attention path (FoX, Gated Softmax, FPA). FPA replaces the path sum of a forget gate with a range minimum of a learned field, so a later token can erase newer keys while preserving older ones: the alive keys form a stack, and FoX is its non-rising-walk special case. Formulations, pseudocode, and a comparison table are collected in Appendix 16.

Fast-weight attention recasts the recurrent update as an online continual-learning rule; the six Falcon variants unify DeltaNet-style regression and linear-attention writes under read-after-write semantics with normalized step sizes and explicit forgetting. Each variant ships a recurrent reference scan and a chunk-parallel kernel. Formulations, configuration knobs, and published results are collected in Appendix 18.

Optimizer Families and Training Dynamics

Optimizing large transformer architectures is complicated by non-convex landscapes, saddle points, and heterogeneous parameter groups. The system supports thirty-nine optimizer families in eight categories: classical baselines (SGD, SGD with momentum, RMSProp, AdamW, RAdam, LaProp); momentum and variance reduction (Adan, ADOPT, AdEMAMix, MARS, Cautious AdamW); memory-efficient and low-rank methods (Adafactor, GaLore, Lion, APOLLO, APOLLO-Mini, Q-APOLLO, AdamPM); schedule-free and parameter-free methods (Schedule-Free AdamW, Prodigy) with large-batch scaling (LAMB); curvature-aware and second-order methods (Sophia, Shampoo, SOAP); geometry-oriented methods (Muon, Turbo-Muon, AF-Muon, AdaMuon, Dion3, ROOT, Muon², NorMuon, Conda, TUON, NAMO, NAMO-D); adaptivity-tunable optimizers (Anon); and evolutionary population-based hybrids (EvE, EA4LLM). Algorithms, pseudocode, and memory and complexity comparisons are collected in Appendix 20.

Quantization and Deployment

The deploy stage converts a trained checkpoint into a deployment artifact through post-training quantization. BitNet-style ternary packing is complemented by two block-affine formats and a ternary safetensors format, all emitted inside a self-contained HuggingFace package whose only weight container is model.safetensors. In the block-affine formats, weights are grouped into blocks of 32: Q8_ stores one FP16 scale and 32 signed 8-bit values per block (8.5 bits per weight), whereas Q4_ stores one FP16 scale and 16 packed nibbles (4.5 bits per weight). The ternary i2_s layout packs four symbols per byte and adds one FP32 scale per tensor (2 bits per weight before the scale). These formats complement, rather than replace, the native standard and quantized artifacts and the BitNet GGUF export used by external runtimes.

The safetensors artifact supports two load-time runtimes. In the dequantize-on-load runtime (v1), packed blocks are expanded to FP32 parameters while the model is loaded, so inference executes exactly like the unquantized backbone. In the packed-linear runtime (v2), each quantized projection keeps its packed buffers and dequantizes inside forward; when a projection replaces a BitLinear layer, the input LayerNorm and INT8 activation quantization are preserved. Only two-dimensional linear weights whose last dimension is a multiple of 32 are quantized; embeddings, normalization parameters, biases and convolutions remain in FP32, which bounds the achievable compression ratio. The experiment harness exports Q8_, Q4_ and i2_s in both runtimes and reports their task accuracy alongside the unquantized export. The formulations, block layouts and size estimates are deferred to Appendix 22.

SBERT Downstream Tasks

Sentence embedding is built on Siamese-style training . For sentence pair \((s_1,s_2)\) with embeddings \((e_1,e_2)\): \[\text{cos}(e_1,e_2)=\frac{e_1^\top e_2}{\|e_1\|\|e_2\|}\] and regression-style cosine loss: \[\mathcal{L}_{\text{cos}}=\left(\text{cos}(e_1,e_2)-y\right)^2\] with \(y\in[-1,1]\) in this pipeline.

Supported downstream modes:

• Similarity: pairwise score between two sentences.

• Search: top-\(k\) nearest neighbors over a corpus.

• Cluster: grouping embeddings (e.g., k-means).

• Encode: persistent embedding export for later retrieval.

The pseudocode for the SBERT inference mode router—dispatching among similarity, search, cluster, and encode modes over a shared encoder—is deferred to Appendix 24.

No-Code Integration: The dashai-frankenstein Plugin

The CLI and web builder described in Section 3 serve researchers comfortable with YAML configuration files. A complementary audience—domain specialists, educators, and analysts who need a trained model but not a training script—is better served by a graphical environment. Frankenstein Transformer therefore ships dashai-frankenstein, an adapter that registers its model classes as components of dashAI (https://github.com/DashAISoftware/dashAI), an open-source, local-first machine-learning workbench developed by a Chilean academic consortium (Universidad de Chile/FCFM, CENIA, IMFD, and ANID funding). Both systems are schema-driven—dashAI renders forms from Pydantic schemas, Frankenstein validates YAML against a strict JSON Schema—so the integration leaves dashAI’s backend untouched; the plugin is discovered through the standard dashai.plugins entry-points group documented in dashAI’s plugin guide (https://docs.dash-ai.com/learn/guides/plugins/).

Registered Components

A single entry-points group registers five components:

• FrankensteinMLMModel (BaseModel, TextClassificationTask): encoder with a classification head over pooled output; predict returns per-class probabilities.

• FrankensteinDecoderModel (BaseGenerativeModel, TextToTextGenerationTask): forces the decoder class and exposes autoregressive top-\(k\) sampling.

• FrankensteinViTClassifier (BaseModel, ImageClassificationTask): uses the ViT classification head (Appendix 32).

• FrankensteinViTSegmenter (BaseModel, plugin SegmentationTask): uses the ViT segmentation head and returns per-pixel class maps.

• SegmentationTask (BaseTask): declares the image-in/mask-out contract that dashAI does not ship.

All share one thin facade over the engine API, so the plugin has a single heavy dependency.

Passthrough Configuration Schema

The Frankenstein JSON Schema remains the single source of truth instead of translating its fields into native dashAI forms. Each component exposes one frankenstein_json field holding a complete training configuration as single-line JSON. Before any training runs, the adapter parses the JSON, validates it against the resolved schema, and applies the config loader’s cross-component checks; failures surface as a readable error tagged with the failing stage. Dataset adapters map dashAI datasets onto Frankenstein’s expectations, inject the vocabulary size from the resolved tokenizer, and stream per-epoch metrics back to dashAI.

Status

The three implementation phases—classification MVP, generative decoder and ViT classifier, and ViT segmenter—are complete, verified inside a live dashAI instance (0.10.x and 0.11.0), and published on PyPI. Known limitations include a direct classification loop rather than MLM pretraining, deferred native form fields for high-leverage knobs, and unconfirmed rendering of segmentation masks. The adapter deliberately leaves dashAI’s core untouched, so future releases cannot break it through coupling.

Empirical Evaluation

The preceding sections describe the toolkit in terms of what it can express: which mixers, optimizers, normalizations, and deployment paths the schema admits. This section turns from design to evidence and reports what happened when that design was exercised at scale. Every mixer pair, every optimizer, and every deployment path in the study was trained, exported, quantized, and evaluated through the same YAML interface, so the measurements that follow are comparable by construction: a difference between two configurations reflects the configurations themselves, not differences in the plumbing around them.

The campaign is organised in rounds. A round is one complete pass of the experiment harness over a generated set of configurations, from training through reporting. Round 1 was a rehearsal at smoke scale—5,129 stage records, with training runs of roughly half a minute each—whose purpose was to debug the pipeline itself before committing serious computation. Round 2 is the measurement campaign analysed here: 868 configurations trained end-to-end, 2,656,372 logged optimisation steps, and 905 CPU-hours of training across those runs, published as an interactive report with downloadable datasets so that every number below can be rederived without re-running anything.¹

Two appendices accompany this section and serve complementary purposes. Appendix 35 documents the harness—the five-stage pipeline, the three independent recorders, the consolidated datasets, and the reporting toolchain—in enough detail that launching a new round reduces to pointing the same scripts at a new folder. Appendix 36 reports Round 2 in full, with the parameter census, the training curves, the complete A\(\times\)B atlas of mixer pairs, the ablation sweeps, and the deployment measurements. The present section states the findings and interprets them; the appendices carry the supporting evidence.

What One Round Measures

Every configuration passes through the same five stages, shown in Figure 12: it is trained for 100 epochs, its checkpoint is exported, the checkpoint is deployed twice (once in a standard format, once in a post-training-quantized format), the two deployed models are evaluated on four held-out masked prompts each, and a Hugging Face package is exported. A single YAML file drives all five stages: it selects the model (one mixer, a pair of mixers declared as layer_pattern: [A, B], a vision transformer, or an ablation of the backbone), the task, and the training budget. Because the file is the only input, any measurement can be traced back to the configuration that produced it, and any configuration can be re-run to reproduce its measurement. The current harness extends the export stage to Q8_, Q4_ and i2_s safetensors variants in both load runtimes and records their task accuracy; the Round 2 tables below were produced before that extension.

Three recorders observe the process independently and are joined only at reporting time: a per-step training log, a parameter census by module group, and a ledger with one row per stage execution. Independence is a deliberate safeguard—no single recorder can silently corrupt the picture, and a disagreement between recorders is itself a diagnostic. The report joins all three by configuration slug, a short identifier derived mechanically from the YAML file.

Round 2 was generated programmatically rather than curated by hand. The hybrid pairs are all combinations of the 42-entry mixer list taken two at a time, materialised as a two-layer model whose layer 0 runs mixer A and layer 1 runs mixer B; single-mixer baselines use one layer. Because the combinations follow the fixed order of the list, no reversed pair \((B, A)\) exists: the campaign measures A__B but never B__A. The “first versus second mixer” asymmetry reported below is therefore confounded with list order—position in the stack coincides with position in the list—and must be read descriptively rather than causally. Table 1 summarises the population by family.

The families differ in what they hold fixed, and the distinction guides every comparison that follows. The hybrid pairs and singles vary the mixer stack itself. The optimizer, structural, normalization, and positional-encoding families share the same 210k-parameter backbone and vary exactly one component—which is precisely what makes them ablations. The task, tokenizer, batch-size, real-data, and audio families change the objective or the data, and therefore the loss scale, which is why they are never compared directly against the mixer sweep.

Findings

Six findings summarise the round. Each is stated as a claim, supported by a measurement, and followed by its interpretation.

The second mixer is the stronger lever. For each candidate mixer \(B\), consider the mean final loss over all pairs in which \(B\) occupies the second position. This marginal mean ranges from 9.37 nats (sparsek) to 10.18 (mamba)—a spread of 0.81 nats. The corresponding range for the first position is only 0.55 nats (9.26–9.81, restricted to mixers with enough pairs for a stable mean). The same asymmetry appears when the atlas is read geometrically: column by column (fixed B, varying A), a good second mixer keeps almost every first mixer inside the 8.9–9.6 nat band, while row by row (fixed A, varying B), no first mixer rescues a bad second one. In other words, the pair inherits its ceiling from its second layer, and at this scale the choice of B is worth roughly 1.5\(\times\) the choice of A. Losses are reported in nats—the natural-logarithm unit of cross-entropy—so a gap of 0.81 nats corresponds to a factor of \(e^{0.81} \approx 2.2\times\) in the geometric-mean likelihood assigned to the correct token.

Pairing is not additive. If mixers contributed independently, the loss of a pair would be predictable from its parts. It is not. Define the marginal mean \(\mu_M\) as the mean loss of all pairs containing mixer \(M\), and the synergy of a pair as the average of its two marginals minus its measured loss, \(S_{A,B} = (\mu_A + \mu_B)/2 - L_{A,B}\): positive synergy means the pair is better than its parts suggest, negative means worse. The best interaction is nsa+engram at \(+1.304\) nats and the worst is gqa+tucker at \(-1.018\). The positive tail consistently pairs a cheap, stable first mixer with a memory-heavy second one; the negative tail is dominated by second mixers that look average in isolation (gta, gqa, tucker) but clash when stacked after a strong first mixer. The practical consequence is that the atlas cannot be summarised by two ranked lists: interaction noise of \(\pm 0.3\) nats is larger than the difference between most optimizers in the ablation.

Scale is a weak lever at this size. Within the hybrid pairs, the parameter count correlates with final loss at \(r = -0.20\) and with held-out accuracy at \(r = -0.07\): at this scale, adding parameters barely buys quality. What parameters do predict reliably is cost (\(r = 0.68\) with step time, \(r = 0.65\) with wall-clock). The largest model of the round (engram__conformer, 415,904 parameters) is not the best. Across all 868 rows the correlation appears much stronger (\(r = -0.43\) for loss), but this is a family confound in the manner of Simpson’s paradox: vision transformers are simultaneously larger and far easier to fit than MLM pairs, so the pooled coefficient measures the gap between families rather than the effect of size. The comparison of interest is the within-pairs one, and there size is close to irrelevant.

Post-training quantization is essentially free. Quantization here means converting trained weights to a lower-precision format after training, with no retraining. Of the 810 configurations deployed in both formats, 800 keep an identical top-5 hit rate; the mean top-1 moves from 5.62% to 5.43% and the mean top-5 from 13.36% to 13.24%—deltas of \(-0.19\) and \(-0.12\) percentage points, within the resolution of a four-sentence evaluation. Quantized deployment is not faster end-to-end—it wins in only 32.3% of runs, with a median ratio of 0.64\(\times\)—because the deployment stage is dominated by serialisation and format conversion rather than by arithmetic. Evaluation itself, which does measure inference, is faster in 51.2% of runs with a median 1.03\(\times\) speed-up. The quality claim is firm (quantization does not move the metrics); the speed claim is about the pipeline, not the format.

The sweep is honest about its failures and its blind spots. The ledger holds 5,271 stage records, of which 275 are failures (171 in training, 104 in deployment); all of them fall outside the 868 analysed configurations, which were selected precisely because they trained. Four mixers (cca, ccgqa, msa, ssog) never train in any configuration, which accounts exactly for the 158 missing cells of the 861-cell atlas. Four others (standard, bigbird, longformer, titan) produce bit-identical losses because their attention windows (at least 64 tokens) cover the whole 32-token sequence—at this context length they collapse to the same dense model, and the atlas measures one model four times. Separately, the ledger records 37 pairs under both retnet and the now-removed retnet_attn alias with identical metrics—one implementation under two names—which the analysis collapses. Reporting these defects explicitly is a design choice: a silent zero or a hidden duplicate would contaminate every downstream statistic.

The vision block works. Segmentation reaches 90.7% mean held-out accuracy, classification 75.8%, and patch prediction 69.9%—an order of magnitude above the MLM pairs. The gap reflects the difficulty of the tasks, not the quality of the mixers: guessing a masked token in a 512-token vocabulary is harder than the vision objectives at this scale. This is why the families are kept separate throughout the report, and why no cross-family ranking is stated.

Threats to Validity

The findings above are conditioned on the scale of the round, and four limitations should temper their reading.

• Scale and convergence. The models are small (127k–416k parameters), the vocabulary is 512 Spanish tokens, the sequences are 32 tokens long, and each configuration trains for 100 epochs with a single seed. Every learning curve is still descending at epoch 100, so the ranking rewards fast learners rather than converged quality.

• Resolution of the evaluation signal. Inference uses four sentences per configuration. One additional correct token moves a top-5 rate by 25 percentage points, so the held-out signal is a smoke detector for catastrophic failures, not a ranking instrument.

• Optimizer ties. Several optimizers converge to the same loss because the run is too short to separate them; the optimizer ranking is indicative at best.

• Known failure modes. One text-classification run collapsed to a constant prediction—a failure mode recorded as such, not a usable model—and the four never-training mixers and the retnet alias are carried forward as explicit caveats.

The rankings in Appendix 36 are therefore directional: they indicate where to look next, not final verdicts. The next round should re-run the atlas at 256–512 tokens, with a 200-sentence evaluation set, at least three seeds, and 300 epochs for the top pairs.

Discussion

The design choices in Frankenstein Transformer reflect several engineering and research tensions in modern deep learning tooling.

The campaign reported in Section 9 turns several of those tensions into measurements rather than opinions. The strict schema paid off operationally: 868 configurations trained, deployed and were evaluated through the same contract, and the failures that remained were concentrated in four mixers and in checkpoint serialisation, not in configuration drift. The prefixed optimizer system proved expressive but, at 100 epochs, undecisive—a dozen optimizers converged to the same loss—which is itself an argument for the schema: only a controlled sweep can show that a design dimension does not matter at a given scale.

Schema-Driven Design Trade-offs

The schema-first approach provides significant reproducibility benefits by enforcing explicit contracts and failing fast on invalid configurations. However, this approach also introduces rigidity: adding new architectures or optimizers requires schema extensions rather than loose command-line arguments. The prefixed hyperparameter system enables fine-grained control but increases configuration complexity for users accustomed to simpler interfaces.

The decision to enforce additionalProperties: false at all schema levels eliminates silent parameter swallowing that has plagued earlier configuration systems, but this strictness requires careful schema maintenance when extending system capabilities. Each new attention mechanism or optimizer variant must be properly integrated into the validation framework, including schema field definitions with appropriate types and constraints, prefixed hyperparameter mapping for optimizer-specific groups, default values aligned with research best practices, and documentation strings for web interface rendering.

Architectural Coverage and Gaps

The forty-nine implemented mixers span the major research directions in sequence modeling, but certain gaps remain. The system lacks recent hybrid architectures such as Griffin and Jamba , which combine gating with state-space models. MoE routing is implemented for FFN layers but not for attention computation, where recent work has shown benefits .

The sparse attention coverage is broad. The two inference-time methods (FASA, SpargeAttn) keep their training-free character in the sense that their sparsification rule is never learned; instead, the implementation trains the layer as standard dense attention over the same projections and activates the sparse selection/pruning path only at evaluation, so both mixers can be trained from scratch and their sparse behaviour is exactly the one described in the original papers.

The gated and fast-weight coverage spans the major categories, including the six Falcon variants . Residual-stream macro-design is covered by Attention Residuals and by the manifold-constrained \(n\)-stream residual (mHC) , whose Sinkhorn–Knopp projection onto the Birkhoff polytope restores the identity-mapping property lost by unconstrained hyper-connections; mHC is validated only at the level of the published 3B–27B MoE experiments and would benefit from independent reproductions at smaller scales (Appendix 27).

Optimizer Landscape Fragmentation

The support for thirty-eight optimizer families across eight categories demonstrates comprehensiveness but also highlights the fragmented state of optimization research. Users face significant decision complexity when choosing among variance-reduction methods (Adan, MARS), memory-efficient variants (GaLore, Adafactor, APOLLO), and curvature-aware approaches (Sophia, Shampoo). The prefixed hyperparameter system, while powerful, requires understanding of which parameters are relevant for each optimizer class.

The implementation quality varies across optimizers: classical methods (AdamW, SGD with momentum) are highly optimized in PyTorch, while newer methods (Muon, Turbo-Muon, AF-Muon, SOAP) may require custom implementations that affect numerical stability and performance characteristics.

Deployment and Production Considerations

The quantization pipeline demonstrates practical deployment concerns but makes specific engineering trade-offs. Ternary weight packing reduces storage to approximately \(1.58\) bits per parameter, but this aggressive compression may degrade performance, especially for smaller models where quantization error is more significant. The safetensors exporter instead offers block-affine Q8_ and Q4_ formats (8.5 and 4.5 bits per weight) and applies all three only to two-dimensional linear weights whose input dimension is divisible by 32; embeddings, normalization parameters and biases remain in FP32, so the effective compression ratio depends on the linear fraction of the model. The packed-linear runtime keeps the artifact compact at the cost of a dequantization step in every forward pass.

The SBERT workflows provide practical utility for semantic similarity and retrieval tasks, but implementation assumes standard pooling strategies (CLS token, mean pooling). Recent advances such as Matryoshka embeddings and contrastive learning refinements are not yet incorporated.

Integration and Extensibility Challenges

Extensibility is handled by a declarative component registry rather than ad-hoc dispatch. Every pluggable component—sequence mixer, optimizer, normalization, activation, residual strategy, positional encoding and token embedding—is described by a frozen ComponentSpec registered through , and the schema enumerations, per-component parameter fragments, optimizer prefix reference and conditional prefix rules are assembled from that registry before validation (Section 3). A third-party plugin registers the same kind of spec through the public authoring API, either from the configuration’s plugins: list or through the frankenstein.plugins entry-point group, so it validates strictly, appears in the schema-generated builder and is dispatched by the same core code as the built-ins. The maintenance pressure is therefore the size of the component catalog itself—forty-nine mixers and thirty-eight optimizers—rather than the dispatch logic, which contains no per-optimizer branches and no per-mixer selection ladder.

The web configuration interface provides significant usability improvements but introduces deployment complexity: running Streamlit alongside training jobs requires additional resources and infrastructure considerations that may not be appropriate for all environments, particularly HPC clusters without web access.

Conclusion

Frankenstein Transformer presents a unified, configuration-driven experimentation platform addressing three challenges in modern deep learning research: architectural fragmentation across dense attention, recurrent models, sparse patterns, gated mechanisms, latent compression, memory layers, fast-weight memories, and geometric priors; optimizer landscape complexity spanning classical baselines, variance-reduction methods, memory-efficient variants, schedule-free approaches, curvature-aware algorithms, and geometry-oriented methods; and end-to-end deployment workflows spanning quantization and sentence embedding.

The system’s primary contributions are:

1. Schema-First Design: A strict YAML-based configuration contract with validation and prefixed hyperparameter routing enabling reproducible experiments across forty-nine sequence mixers and thirty-eight optimizer families.

2. Comprehensive Architecture Support: Implementation spanning dense baselines, recurrent and retentive blocks, sparse attention, gated mechanisms, latent compression, memory layers, fast-weight attention, and geometric-field attention, with encoder, decoder, and Vision Transformer model classes.

3. Unified Optimizer Framework: Prefixed hyperparameter groups enabling fine-grained control over embeddings, normalization layers, attention weights, and the remaining parameters across classical baselines (SGD+Momentum, AdamW, RAdam), variance reduction (Adan, ADOPT, AdEMAMix, MARS, Cautious), memory-efficient and low-rank methods (Adafactor, GaLore, Lion, APOLLO, APOLLO-Mini, Q-APOLLO), large-batch and schedule simplification (LAMB, Schedule-Free AdamW, Prodigy), curvature-aware (Sophia), second-order (Shampoo, SOAP), geometry-oriented (Muon, Turbo-Muon, AF-Muon), and adaptivity-tunable (Anon) optimizers.

4. End-to-End Workflows: Integrated deployment pipeline supporting ternary weight packing, block-affine 8/4-bit safetensors export at both load-time runtimes, and INT8 activation quantization; SBERT-inspired training and inference for semantic similarity, retrieval, and clustering tasks.

5. Interactive Configuration: Streamlit-based web interface providing schema-driven form generation, real-time validation, inline documentation, and CLI command synthesis.

This system enables rapid experimental iteration while maintaining reproducibility through strict configuration contracts. By consolidating diverse research contributions into a unified toolkit, it lowers barriers to exploring novel architectures and optimization strategies, particularly for researchers who may lack resources to implement and validate each variant independently.

Limitations and Future Directions

Several limitations and promising directions for future work emerge from this system’s design and implementation:

1. Architecture Integration: Recent hybrid architectures (Griffin, Jamba, Mamba-X) demonstrate benefits of combining multiple mechanisms into unified blocks. Future versions should integrate these architectures and explore systematic composition patterns.

2. Advanced Quantization: Current safetensors quantization is weight-only and packs two-dimensional linear weights, dequantizing without fused integer kernels; layer-wise, channel-wise and importance-aware policies remain future work.

3. Automated Hyperparameter Optimization: The schema supports extensive hyperparameter spaces, but users must manually explore these spaces. Integration with Bayesian optimization, multi-armed bandit strategies, or gradient-based hyperparameter tuning could automate effective configuration discovery.

4. Production Deployment: The web interface improves usability but may not be appropriate for all deployment environments. Headless configuration modes, API-based configuration management, or improved CLI ergonomics could serve HPC and production workflows.

5. Evaluation Benchmarking: Round 2 (Section 9) provides the first comparative atlas across the mixer list, but it is a 32-token, single-seed, 100-epoch measurement. Future rounds should extend it to longer context, more held-out sentences and multiple seeds before the rankings are read as architectural conclusions.

6. Training Stability Guarantees: Current implementation includes NaN/Inf guards and gradient clipping, but formal analysis of stability conditions for different mixer-optimizer combinations, particularly with looped blocks and aggressive quantization, remains open.

7. Multimodal and Task-Specific Extensions: The current design covers text and single-image tasks through the Vision Transformer class. Extensions to vision-language models, multimodal architectures, and task-specific fine-tuning workflows (e.g., instruction tuning, RLHF) would broaden applicability.

The research trajectory of sequence modeling continues toward hybrid approaches that combine the strengths of multiple paradigms—compression from recurrence, selectivity from attention, gating for memory management, and sparsity for efficiency. A unified experimentation platform like Frankenstein Transformer becomes increasingly valuable as this convergence accelerates, enabling researchers to explore this expanding design space with reproducible, well-engineered infrastructure.

Conceptual Introduction—Transformers and Attention for Beginners

This appendix is the on-ramp to the paper: it explains what a transformer is and how attention works, for readers with no deep-learning background. Every mechanism is introduced through one real-world analogy, and — this is the point of the appendix — the same analogy returns every time the mechanism returns, so a single mental image carries you from the first intuition to the mathematics. Table 2 lists the cast; the running example is always the same sentence, “The cat eats fish”.

What Is a Transformer? Words in Context

A transformer is a neural architecture that processes text (or any sequence) by interpreting the meaning of each word in the light of all the other words around it. The problem it solves is easiest to see with an ambiguous word. Read this:

What does “bank” mean in each case? Taken alone, the word has three unrelated meanings: a financial institution, a bench to sit on, the ground beside a river. Inside the first sentence, the neighbors (“full of people waiting”) vote for the branch office. Inside the second, “Mary went to the…” leaves two readings alive — she could be heading to the counter or to the river. Figure 13 contrasts the two situations.

This ability to re-read every word in the light of the whole sentence is called contextualization, and it is the heart of every modern transformer. The question this appendix answers is how the model decides which neighbors to listen to. The answer is attention.

The Attention Mechanism: One Voice in a Noisy Room

Attention is the “magic trick” behind contextualization. Think of a conversation in a noisy room:

1. Lots of noise: many people talk at once, but one of them matters to you.

2. You ignore the rest: your brain automatically focuses on that one voice and filters the rest.

3. You understand better: by concentrating, you catch every word clearly.

The neural mechanism works the same way, with one twist: every word in the sentence is such a listener. Each word “asks” how important every other word is for understanding it, then focuses its attention on the most relevant ones. The room is the sentence, the voices are the words, and the focusing ability is attention. Keep this image: every formula below is just the arithmetic version of listening selectively.

A Worked Example: The Two Lives of “eats”

Let us watch the mechanism work on one word in two different sentences. The numbers that follow are illustrative, but they are exactly the kind of table a trained transformer computes — Figure 14 shows both cases.

Context 1: “The cat eats fish”

Processing “eats”, the word-as-listener asks about every other voice in the room:

• How relevant is “The”? Little — it is just an article. Low attention: 5%.

• How relevant is “cat”? Very — it is the subject, who eats. High attention: 55%.

• How relevant is “eats”? A little — itself. Low attention: 5%.

• How relevant is “fish”? Very — it is what gets eaten. High attention: 35%.

The new representation of “eats” is dominated by “cat” and “fish”: an action, with its subject and its object.

Context 2: “The restaurant eats into profit margins”

Same word, different room. Now the listeners that matter are different:

• “restaurant” (the subject): 30%.

• “profit” (what is affected): 25%.

• “margins” (what gets consumed): 35%.

• “The”, “into”, and “eats” itself: the remaining 10%.

The word “eats” ends up with a completely different representation in each sentence because it attends to different words. That is the whole trick: one word, many possible meanings, and attention picking the right neighbors to define it.

How Attention Works Mathematically: A Library Called Attention

Behind the listening there is arithmetic. The standard analogy — and the one this paper sticks to — is a library, shown in Figure 15.

Step 1: Queries, Keys, and Values

Each word is transformed into three versions of itself:

• Query (Q): the visitor’s question — “what information do I need?”

• Key (K): the catalog card — “what information do I offer?”

• Value (V): the book’s content — “here is my actual information.”

Crucially, every word plays all three roles at once: it asks (Q), it advertises (K), it delivers (V). One visit to the library, with every book simultaneously a visitor and a shelf.

Step 2: Compatibility

The librarian compares the question with every card. Questions similar to a card get a high compatibility score; the score is computed by multiplying the query by the key (a dot product). For our “eats” query: \(\mathrm{s}(\text{eats} \rightarrow \text{cat}) = 3.40\), \(\mathrm{s}(\text{eats} \rightarrow \text{fish}) = 2.95\), \(\mathrm{s}(\text{eats} \rightarrow \text{The}) = 1.00\). Figure 16 places the four steps in one pipeline.

Step 3: Focus

Raw scores are just numbers; attention is a budget. A function called softmax converts the scores into percentages that always sum to 100%: the “eats” query turns the scores above into 5% / 55% / 5% / 35% over “The / cat / eats / fish”. Figure 17 shows the conversion. High compatibility means “focus attention here”; low compatibility means “ignore this” — back in the noisy room, the voice you decided to filter out.

Step 4: Combine

Finally the visitor leaves the library with a weighted mix of contents. If “eats” places 55% of its budget on “cat” and 35% on “fish”, then 55% of cat’s value and 35% of fish’s value are blended into the new representation of “eats”: \[\text{new ``eats''} \;=\; 5\%\cdot v_{\text{The}} + 55\%\cdot v_{\text{cat}} + 5\%\cdot v_{\text{eats}} + 35\%\cdot v_{\text{fish}}. \label{eq:tut-mix}\] That blended vector is what the rest of the network sees: a representation of “eats” that already contains who eats and what gets eaten.

Multiple Attention Heads: A Team of Analysts

Transformers never run one attention but many heads in parallel — as if the sentence were analyzed by a team of specialists, each asking a different question about it (Figure 18):

• Analyst 1 (grammar): “who does it?” — links cat \(\rightarrow\) eats.

• Analyst 2 (object): “what receives it?” — links eats \(\rightarrow\) fish.

• Analyst 3 (time): “when / how often?” — tracks verb tense and time words.

• Analyst 4 (style): “which articles and modifiers?” — tracks “the”, “a”, adjectives.

Each head learns its own focus pattern from data, and their reports are concatenated and re-mixed by one learned projection. Together they capture far more than any single head could — grammar, roles, time and style at once.

Stacking Layers: Editing the Essay in Passes

One attention step still sees shallow patterns, so transformers stack the same block many times (12 to 48 in practical models). The analogy: editing an essay in successive passes (Figure 19).

1. First pass — fix spelling and basic grammar: recognize words, forms, gender.

2. Second pass — improve clarity of each sentence: who does what to whom.

3. Third pass — ensure consistency and narrative flow: idioms, relations, abstract meaning.

Each layer re-reads what the previous layer produced and refines it: early layers capture basic features, later layers understand complex relations and abstract concepts. No layer is told what to look for — the passes organize themselves during training.

Why Does It Matter?

The attention mechanism was revolutionary because:

1. Parallelism: any word can be compared with any other word at once (no sequential processing, unlike earlier recurrent systems) — fast on modern hardware.

2. Flexibility: the patterns are learned from data; nobody programs the rules.

3. Scalability: the same mechanism works from a sentence to millions of words of context.

4. Generality: translation, summarization, question answering, generation, computer vision, audio — one mechanism, many tasks.

Practical Challenges: The Price of Listening to Everyone

Dense attention buys its power by comparing every word with every other word, and that price grows fast.

Computational Complexity

A 1,000-word document forces \(1{,}000 \times 1{,}000 = 1{,}000{,}000\) word-to-word comparisons, and the cost squares as the document grows. Figure 20 shows the growth against the friendlier alternatives discussed below.

Memory Requirements

During inference (using a trained model), the keys and values of every past word must be stored — the library keeps every book on the table just in case a later query needs it. For long contexts this cache dominates GPU memory.

Device Efficiency

Transformers were designed for powerful TPUs and GPUs; running them on phones or embedded devices is challenging.

Modern Solutions: Smarter Ways to Listen

The rest of this paper is a catalog of ways to keep the benefits of attention while paying less. Each keeps our analogies:

• Sparse attention: stop browsing every card in the library — visit only nearby shelves and a few strategic ones. Complexity drops from \(\mathcal{O}(n^2)\) to \(\mathcal{O}(n \log n)\) or \(\mathcal{O}(n)\).

• Recurrent models (e.g. Mamba): take notes as you read, instead of re-listening to the whole room for every new word — an old idea (recurrence) rebuilt with modern efficiency.

• Gated attention: a gatekeeper decides which of your notes are worth keeping and which to fade out, so less needs to be stored.

• Quantization: write each number with fewer digits (small integers instead of wide decimals), shrinking memory with minimal loss of precision.

These innovations are exactly what this toolkit (Frankenstein) lets you experiment with, and the annexes that follow are guided tours of each family.

Key Takeaways

• A transformer is a network that contextualizes: it re-reads every word in the light of all its neighbors (the ambiguous “bank”).

• Attention is selective listening: each word focuses its limited 100% budget on the neighbors that define it (the noisy room).

• The arithmetic is a library visit: query asks, key advertises, value delivers; dot products score compatibility, softmax turns scores into a 100% budget, and the values are mixed by that budget.

• Multiple heads are a team of analysts with one question each; stacked layers are successive editing passes, each refining the last.

• The catch: dense attention compares everyone with everyone — \(\mathcal{O}(n^2)\) — which is why sparse, recurrent, gated and quantized variants exist.

• Every later annex of this paper is a deeper tour of one of those variants; the analogies of Table 2 will keep coming back.

Dense, Recurrent, and Memory-Augmented Transformers

This annex synthesizes the dense, recurrent, and memory-augmented transformer families. The field has evolved toward several primary paradigms: dense global attention with variants, key-value head compression via grouped-query attention, recurrent state compression with decay, selective state-space models, continuous-depth numerical integration, token-indexed value memory, and memory-augmented architectures. Understanding this taxonomy illuminates fundamental tradeoffs between expressiveness, computational cost, memory footprint, and deployment simplicity.

Dense Attention Baselines: Standard, Sigmoid, Continuum, Exclusive, and Multi-Token

Standard Softmax Attention

Standard multi-head self-attention computes scaled dot-product similarity between token embeddings:

For autoregressive generation, KV caching stores past keys and values to avoid \(\mathcal{O}(n^2)\) recomputation. However, this linearly growing cache occupies \(\sim n \cdot d_{\text{hidden}}\) bytes, which can consume hundreds of gigabytes for billion-parameter models. Standard attention achieves perfect expressiveness within the context window—any token can attend to any other token with learned weights—but pays the price of dense computation.

• Training complexity: \(\mathcal{O}(n^2 \cdot d)\) time, \(\mathcal{O}(n^2)\) space (attention matrix materialization).

• Inference complexity: \(\mathcal{O}(n)\) time per token, \(\mathcal{O}(n)\) space (KV cache).

• Strengths: Unparalleled expressiveness; perfect history recall; highly parallelizable.

• Weaknesses: Quadratic bottleneck prohibits very long contexts; KV cache dominates memory during generation.

Sigmoid Attention

Sigmoid attention replaces row-wise softmax with element-wise sigmoid activation :

Unlike softmax, sigmoid does not enforce a probability distribution (weights need not sum to 1), enabling stronger token independence. Theoretical analysis via mixture-of-experts shows sigmoid achieves superior sample complexity: \(\mathcal{O}(n^{-0.51})\) convergence for ReLU experts versus softmax’s \(\mathcal{O}(n^{-0.24})\). However, empirical training revealed gradient instabilities at scale. The remedy is hybrid-norm—adding normalization after the attention output—which stabilizes gradients without sacrificing the theoretical benefits of element-wise gating.

• Training complexity: \(\mathcal{O}(n^2 \cdot d)\) (identical to standard), but element-wise operations enable 17% inference speedup via FlashSigmoid.

• Inference complexity: \(\mathcal{O}(n)\) per token with KV cache (asymptotically same, but lower constant factors).

• Strengths: Overcomes zero-sum competition; avoids row-wise synchronization; hardware-friendly implementation.

• Weaknesses: Requires careful stabilization (hybrid-norm); training instability at large scales without auxiliary loss.

Grouped-Query Attention (GQA)

Grouped-query attention bridges multi-head attention (MHA) and multi-query attention (MQA) by partitioning query heads into groups, each sharing a single key-value head. If \(h\) is the total number of attention heads, \(h_k\) is the number of key-value heads, and \(G = h / h_k\) is the group size, then every group of \(G\) query heads shares one KV head:

GQA interpolates smoothly between MHA (\(h_k = h\), \(G = 1\)) and MQA (\(h_k = 1\), \(G = h\)). The key advantage is a tunable tradeoff between KV-cache footprint and representational capacity. Each additional KV head consumes \(2 \cdot d_h \cdot n\) extra bytes in the cache, so reducing \(h_k\) yields proportional memory savings.

• Training complexity: \(\mathcal{O}(n^2 \cdot h \cdot d_h)\) (identical to MHA after head expansion).

• Inference complexity: \(\mathcal{O}(n)\) per token, \(\mathcal{O}(h_k \cdot n \cdot d_h)\) KV-cache space.

• Strengths: Flexible head-count tradeoff; proven at scale (Llama 2 & 3 use \(h_k = 8\) KV heads with \(h = 32\) query heads); reduces KV-cache by \(G \times\); supports fast inference with FlashAttention.

• Weaknesses: Fewer KV heads can reduce representational capacity in small models; requiring \(h_k\) to divide \(h\) constrains architectural choices.

Continuum Attention

Continuum Attention replaces discrete softmax retrieval with exact integration over a continuous, piecewise-linear memory trajectory. Adjacent keys and values are interpolated so that \(\mathbf{k}(t)\) and \(\mathbf{v}(t)\) are defined at every position of the visible domain: \[\mathbf{k}_j(u) = (1-u)\,\mathbf{k}_{j-1} + u\,\mathbf{k}_j, \qquad \mathbf{v}_j(u) = (1-u)\,\mathbf{v}_{j-1} + u\,\mathbf{v}_j,\] with \(u = t-(j-1)\in[0,1]\) on the interval \([j-1,j]\); constant half-token extensions cover \([-1/2,0]\) and \([i,i+1/2]\). Each query induces an exponential relevance density over the visible prefix, \[s_i(t) = \frac{\gamma}{\sqrt{d_h}}\,\mathbf{q}_i^\top \mathbf{k}(t), \qquad p_i(t) = \frac{\exp(s_i(t))}{Z_i}, \qquad \mathbf{o}_i = \int_{-1/2}^{i+1/2} p_i(t)\,\mathbf{v}(t)\,dt,\] where \(\gamma>0\) is the score gain (\(\gamma=1/\tau\)). Since the key trajectory is linear within a segment, the score is affine, \(s_{ij}(u) = a_{ij} + u\,\Delta_{ij}\) with \(a_{ij}=s_i(j-1)\), \(b_{ij}=s_i(j)\) and \(\Delta_{ij}=b_{ij}-a_{ij}\), so the integrals have closed forms in the exponential moments \(J_n(\Delta)=\int_0^1 u^n e^{\Delta u}\,du\): \[\begin{aligned} J_0(\Delta) &= \frac{e^{\Delta}-1}{\Delta}, \qquad J_1(\Delta) = \frac{(\Delta-1)e^{\Delta}+1}{\Delta^2}, \\ M_{ij} &= e^{a_{ij}} J_0(\Delta_{ij}), \qquad \mathbf{U}_{ij} = e^{a_{ij}}\big[(J_0-J_1)\,\mathbf{v}_{j-1} + J_1\,\mathbf{v}_j\big], \\ Z_i &= \tfrac12 e^{s_i(0)} + \tfrac12 e^{s_i(i)} + \sum_{j=1}^{i} M_{ij}, \qquad \mathbf{o}_i = \frac{\mathbf{U}_i^{\mathrm{cap}} + \sum_{j=1}^{i}\mathbf{U}_{ij}}{Z_i}, \end{aligned}\] with \(\mu(\Delta)=J_1/J_0\) the normalized retrieval position inside a segment. Because the contributions are nonnegative, CA is equivalently a normalized nonnegative weighted sum of the original values, \(\mathbf{o}_i = \sum_{j\le i} w_{ij}\mathbf{v}_j\) with \(w_{ij}\ge 0\) and \(\sum_j w_{ij}=1\); equal visible scores reduce exactly to uniform causal attention and \(o_0=v_0\). The learned-gain variant adds one scalar per head (\(\gamma=\exp(\rho)\), initialized at \(\gamma=1\) in the 101M model; the fixed CA-P variant uses \(\gamma=1/0.30\) and adds no parameters). The encoder (bidirectional) form integrates over the full trajectory \([-1/2,n-1/2]\) – the natural symmetric extension of the causal-only paper.

• Training complexity: \(\mathcal{O}(n^2 d)\) – identical to dense softmax; the closed-form moments replace the softmax (no quadrature).

• Inference complexity: \(\mathcal{O}(n)\) per token with a KV cache; the effective weights are recomputed row-locally over the cached prefix each step (exact).

• Strengths: Exact continuous retrieval with normalized nonnegative effective weights; zero parameters beyond the projections in the fixed-gain form; strong multilingual and zero-shot results (FLORES macro 1.8915 BPB, a 5.8% reduction versus the strongest baseline and the best score in every language; HellaSwag and ARC-Challenge best among the baselines); 42% fewer training bytes to a matched 1.10 BPB threshold than a matched Transformer.

• Weaknesses: The quadratic attention arithmetic persists; source-code performance is the weakest of the five architectures (Stack-smol 3.3451 BPB); reduced robustness to byte deletion and “antspeak” perturbations.

Exclusive Self Attention (XSA)

Exclusive Self Attention starts from the observation that in trained Transformers the attention output \(\mathbf{y}_i\) keeps a large cosine similarity with the token’s own value vector \(\mathbf{v}_i\)—the attention similarity bias—and that this similarity grows with depth. Since the residual path and the following FFN already carry the point-wise information of position \(i\), that component is redundant and competes with the layer’s actual job, context modeling. XSA therefore subtracts it explicitly. The first line is standard causal softmax attention: \[\mathbf{q}_i = \mathbf{W}_q\mathbf{x}_i, \quad \mathbf{k}_j = \mathbf{W}_k\mathbf{x}_j, \quad \mathbf{v}_j = \mathbf{W}_v\mathbf{x}_j, \quad a_{i,j} = \frac{\exp(\mathbf{q}_i^\top\mathbf{k}_j)}{\sum_{j'=1}^{i}\exp(\mathbf{q}_i^\top\mathbf{k}_{j'})}, \quad \mathbf{y}_i = \sum_{j=1}^{i} a_{i,j}\mathbf{v}_j,\] and the exclusive step projects the output orthogonally out of the self value direction: \[\hat{\mathbf{v}}_i = \frac{\mathbf{v}_i}{\|\mathbf{v}_i\|_2}, \qquad \mathbf{z}_i = \mathbf{y}_i - \big(\mathbf{y}_i^\top \hat{\mathbf{v}}_i\big)\,\hat{\mathbf{v}}_i .\] The step is parameter-free, row-local (it only consumes the current token’s value) and differentiable, so it trains end-to-end with the usual four projections \(\mathbf{W}_q,\mathbf{W}_k,\mathbf{W}_v,\mathbf{W}_o \in \mathbb{R}^{H\times H}\); the parameter count is exactly \(4H^2\), identical to standard attention. Removing the self direction can be read as an implicit attention sink: undesired mass is parked on the diagonal \(a_{i,i}\) and then discarded, which is consistent with the observed robustness to explicit sink tokens.

• Training complexity: \(\mathcal{O}(n^2 d)\) – the softmax attention dominates; the exclusion is \(\mathcal{O}(nd)\) and measured as minimal overhead (speed and memory curves overlap the baseline across sequence lengths and model widths).

• Inference complexity: \(\mathcal{O}(n)\) per token with a plain KV cache; the exclusion reads only the newly computed \(\mathbf{v}_i\), so incremental decoding is exact (class (i)).

• Strengths: Zero additional parameters; two-line change; consistent training/validation gains at 0.7B, 1.3B and 2.7B parameters on FineWeb-100BT (+0.26, +1.03 and +1.36 average downstream points over SA); the margin grows with sequence length (512 to 16k) and is stable across learning rates and attention-sink configurations.

• Weaknesses: Evidence is language-modeling only; the paper offers no theoretical grounding beyond the empirical motivation; per-token L2 normalization is an extra (small) kernel compared with vanilla SDPA; the encoder (bidirectional) variant used by this toolkit is a mask-agnostic extension of the paper’s causal formulation.

Multi-Token Attention (MTA)

Multi-Token Attention targets the single-token bottleneck of soft attention: every weight \(a_{ij}\) is decided by one query and one key vector, so a context region identified by several tokens (a sentence mentioning both “Alice” and “rabbit”) must be compressed into a single vector. MTA keeps the quadratic attention skeleton and makes neighbouring queries, keys and heads interact through small learned convolutions over the score map. The key-query convolution mixes the logits of a window of \(c_q\) queries and \(c_k\) keys, per head: \[a_{ij} = \operatorname{Softmax}\Big(\sum_{i'=0}^{c_q-1}\ \sum_{j'=-\lfloor c_k/2\rfloor}^{\lceil c_k/2\rceil-1} \mathbf{1}_{i\ge j-j'}\ \theta_{i',j'}\ \mathbf{q}_{i-i'}\mathbf{k}_{j-j'}^{\top}/\sqrt{d_h}\Big), \label{eq:mta-exact}\] and, because that indicator would require bespoke CUDA kernels, the default implementation applies the existing causal mask twice: \[\mathbf{A} = \operatorname{Softmax}\Big(\operatorname{Mask}_{-\infty}\big(\operatorname{Conv2d}_{\theta}(\operatorname{Mask}_{0}(\hat{\mathbf{A}}))\big)\Big), \qquad \mathbf{A} = \operatorname{Mask}_{0}\Big(\operatorname{Conv2d}_{\theta}\big(\operatorname{Softmax}(\operatorname{Mask}_{-\infty}(\hat{\mathbf{A}}))\big)\Big), \label{eq:mta-pre}\] where the first expression is the pre-softmax variant (the paper’s default) and the second the post-softmax variant (additive rather than multiplicative interaction); both are applied by default. Head mixing partitions the \(n_h\) heads into \(n_h/c_h\) non-overlapping groups and mixes each group with its own \(c_h\times c_h\) kernel, \[\mathbf{A}^{1}_{\text{new}} = w_{11}\mathbf{A}^{1}+w_{12}\mathbf{A}^{2}, \qquad \mathbf{A}^{2}_{\text{new}} = w_{21}\mathbf{A}^{1}+w_{22}\mathbf{A}^{2}, \label{eq:mta-head}\] so attention maps can be combined across heads (a post-softmax \(2\times2\) mixing is equivalent to doubling the rank of the value/output projections, Appendix B of the paper). Finally a grouped normalization over each set of \(c_h\) heads followed by a per-head sigmoid gate, \(\operatorname{diag}(\sigma(\mathbf{g}))\operatorname{GroupNorm}(\mathbf{A}\mathbf{V})\), pushes back against the residual stream and lets the model switch heads on and off across layers; the paper reports 10.90 validation perplexity with scalar gating versus 11.03 without scaling and 11.41 with layer-norm scaling. Kernels are identity-initialized (the best-converging choice) and the mixer adds only about \(0.001\%\) of the parameters at 880M scale: \(n_h(2c_qc_k + c_h) + 2H + n_h\), i.e. one scalar per head for the gate, an affine group norm, and the kernels.

• Training complexity: \(\mathcal{O}(n^2 d_h n_h + n^2 c_q c_k n_h)\) – still quadratic in the sequence length, with a small constant from the kernels; the score map must be materialised, so fused SDPA kernels are not compatible (the paper’s stated limitation).

• Inference complexity: \(\mathcal{O}(n d_h n_h)\) per step with a KV cache plus an \(\mathcal{O}(c_q c_k n_h)\) convolution correction; decode keeps the classic KV cache and a FIFO tail of the last \(c_q-1\) queries, so the double-masked convolution is reproduced exactly (classes (i)+(v)).

• Strengths: Consistent gains at 880M/105B SlimPajama tokens over Transformer, Diff Transformer and Talking-Heads baselines (average validation perplexity 10.91 versus 11.25, 11.15 and 11.04; 44.9 versus 43.7 average zero-shot points), large long-context retrieval gains (multi-needle, BabiLong, LAMBADA) with \(+0.001\%\) parameters, and near-identity kernels that are learned to be diagonal “sentence-matching”, priming or edge-detecting filters.

• Weaknesses: Incompatible with optimized attention kernels (the score map is materialised); the key-query convolution is applied on every fourth layer in the reference run for efficiency, so the full-per-layer variant costs more arithmetic than reported; the encoder (bidirectional) form used by this toolkit drops both causal masks, an extension of the paper’s causal-only formulation.

Memory Attention

The memory family groups Titans (the attention window implemented here), Engram (Appendix 24) and Memory Attention; this annex covers the attention-side members.

Memory Attention removes the dedicated value projection and constructs values from contextual keys plus layer-specific token memory. With token IDs \(s\), vocabulary size \(N\) and \(d_v = H_{KV} d_h\):

The output decomposes as \(O_{\mathrm{MA}} = A(K{+}M) = AK + AM\): a contextual retrieval term over keys and a memory retrieval term aggregated under the same attention distribution. Across \(L\) replaced layers, MA swaps the standard \(L\,d\,d_v\) value-projection parameters for \(L(N d_v + p_{\mathrm{norm}})\) memory and normalization parameters — added capacity accessed by lookup instead of dense computation. At inference the per-row normalization is folded into the table (\(E[i] \gets \mathrm{Norm}(E[i])\)), reducing online value construction to lookup and addition while retaining a conventional KV cache. The paper reports improved LAMBADA/WikiText perplexity and average downstream accuracy across MHA/GQA/MQA backbones at matched token budgets, with \(1.42\times\)/\(1.16\times\) token efficiency; the comparison includes the additional memory parameters, so it does not isolate structure from capacity. MA-Offload (CPU-resident tables with prefetching) and MA-Recall (reconstructing historical values as \(V_{1:T} = K_{1:T} + E[s_{1:T}]\) for a \({\sim}50\%\) smaller value cache) are analyzed but not implemented here.

• Training complexity: \(\mathcal{O}(n^2 d)\) attention with \(\mathcal{O}(n d_v)\) lookup+norm+add value construction (replacing the \(\mathcal{O}(n\, d\, d_v)\) projection).

• Inference complexity: \(\mathcal{O}(n)\) per token with a conventional KV cache; \(N d_v\) extra parameters per layer.

• Strengths: Extra capacity without dense-compute growth; layer-specific token representations; exact standard-attention aggregation; KV-cache-compatible by construction.

• Weaknesses: Parameter count grows with the vocabulary (\(N \gg d\)); needs token IDs (text-only); the reported gains are not parameter-matched.

Recurrent and Retentive Architectures

Retentive Networks (RetNet)

RetNet unifies three computation paradigms: parallel training, recurrent inference, and chunkwise deployment. Its core innovation is the retention mechanism, which uses a fixed exponential decay matrix to model temporal importance:

The decay scalar \(\gamma \in (0,1)\) controls the temporal window. RetNet uses multi-scale retention with different \(\gamma\) values per head (e.g., \(\gamma = 1 - 2^{-5}, 1 - 2^{-6}, \ldots\)), allowing short-term and long-term dependencies simultaneously. Chunkwise recurrent mode divides sequences into chunks, processes each chunk in parallel, and threads a recurrent state between chunks.

• Training complexity: \(\mathcal{O}(n^2 \cdot d)\) (parallel), or \(\mathcal{O}(n \cdot c \cdot d)\) (chunkwise recurrent with chunk size \(c\)).

• Inference complexity: \(\mathcal{O}(1)\) per token, \(\mathcal{O}(d^2)\) state space (fixed matrix).

• Strengths: Constant-time inference; triple computation paradigm; multi-scale consolidation.

• Weaknesses: Fixed decay imposes rigid inductive bias; may truncate learned long-range patterns.

Mamba: Selective State-Space Models

Mamba frames recurrence as a continuous-time dynamical system with input-dependent parameters, achieving both linear training and constant-time inference:

The critical innovation is that \(\mathbf{A}\) (the system matrix) is not input-dependent, but \(\Delta_t\), \(\mathbf{B}_t\), and \(\mathbf{C}_t\) are, making the system time-varying. This selectivity allows the model to ignore irrelevant information by setting \(\Delta_t\) near zero, effectively creating a gate. A hardware-aware parallel scan algorithm implements the recurrence efficiently on GPUs by fusing computation within SRAM, avoiding expensive HBM bandwidth.

• Training complexity: \(\mathcal{O}(n \cdot d)\) via hardware-aware scan (linear in sequence length).

• Inference complexity: \(\mathcal{O}(1)\) per token, \(\mathcal{O}(d)\) state (hidden vector).

• Strengths: Achieves linear training and constant-inference simultaneously; practical on long sequences (millions of tokens).

• Weaknesses: State vector compression can weaken exact copying and dense associative recall versus full attention.

Continuous-Depth Transformers: ODE Integration

ODE Transformer

The ODE Transformer interprets network depth as numerical integration of a continuous dynamical system, using higher-order Runge-Kutta solvers to reduce truncation error:

Instead of simple Euler residual connections \(h_{t+1} = h_t + f_\theta(h_t)\), the RK4 block computes four intermediate evaluations and combines them with classical RK4 weights. To avoid vanishing gradients, the architecture introduces learned gating that interpolates between intermediate approximations:

This formulation reduces the effective number of parameters through weight sharing—the same \(f_\theta\) is evaluated multiple times—while providing richer trajectory refinement.

• Training complexity: \(\mathcal{O}(k \cdot n^2 \cdot d)\) where \(k\) is the RK order (4 for RK4).

• Inference complexity: \(\mathcal{O}(k \cdot n)\) per token (higher constant overhead).

• Strengths: Significantly higher accuracy on generation tasks (state-of-the-art BLEU); parameter-efficient via weight sharing.

• Weaknesses: High per-step compute; inference latency increases by constant factor \(k\); complex gating required for stability.

Test-Time Memory: Titans

The Titans architecture introduces an orthogonal dimension: instead of static weights, the model maintains a learnable memory that is updated during inference based on a surprise-driven signal:

The memory module \(\mathcal{M}\) is literally updated during the forward pass by computing gradients of an associative loss and applying SGD steps with momentum. The decay rates \(\eta_t, \theta_t, \alpha_t\) are themselves input-dependent, allowing the model to switch memory paradigms when context shifts.

• Training complexity: Roughly \(\mathcal{O}(c^2 + n \cdot f)\) where \(c\) is local window and \(f\) is memory update overhead.

• Inference complexity: \(\mathcal{O}(c^2)\) local attention plus \(\mathcal{O}(1)\) memory retrieval per token.

• Strengths: Handles extreme context lengths; enables true associative recall; memory adapts to input.

• Weaknesses: Significantly more complex; inference includes gradient computation; higher coordination overhead.

Pattern-Driven Mixer Dispatch

The codebase selects a mixer per block from the configured layer_pattern. The dispatch logic routes inference-time sparse blocks (fasa, sparge) to standard dense attention while the model is in training mode, and to their sparse implementation during evaluation; all other blocks route to the corresponding family implementation.

The field exhibits a clear progression along two axes: (i) computational efficiency, moving from \(O(n^2)\) to \(O(n)\) or \(O(1)\), and (ii) memory adaptivity, shifting from static weights to dynamic, test-time updated models. Standard attention remains the expressiveness baseline; Mamba and RetNet represent the practical efficiency frontier; Titans introduces an orthogonal innovation axis (test-time learning). The choice of architecture reflects the fundamental engineering constraint: expressiveness versus deployment cost. The latent-attention family, which compresses the KV cache into a low-rank latent, is covered in Appendix 14.

Latent and Low-Rank Attention Families

This annex covers the latent / low-rank KV compression family of attention mechanisms. These architectures compress the per-token key–value state into a low-rank latent vector that is the only quantity materialised in the KV cache during decoding, then reconstruct full multi-head keys and values on the fly. The family unifies grouped-query, latent, factorised, and temporal-latent designs under a common low-rank bottleneck.

Multi-Head Latent Attention (MLA)

Multi-Head Latent Attention (MLA) compresses the per-token key–value state into a single low-rank latent vector \(c_{KV}\) of rank \(r_{kv}\), which is the only quantity materialised in the KV cache during decoding. A pair of up-projection matrices reconstructs the full multi-head keys and values on the fly, while RoPE is applied to the decompressed query and key tensors so that positional structure is preserved. The authors show that the Pareto-optimal latent width sits near \(r_{kv} = d_{\text{hidden}} / 2\), at which point MLA halves the KV cache yet matches multi-head attention (MHA) quality on small models.

Mathematical Formulation

\[c_{KV} = W_{DKV}\, x \in \mathbb{R}^{r_{kv}}, \qquad k = W_{UK}\, c_{KV}, \qquad v = W_{UV}\, c_{KV}, \qquad q = W_{Q}\, x.\] \[\tilde q = \text{RoPE}(q), \quad \tilde k = \text{RoPE}(k), \qquad \text{attn} = \texttt{softmax}\!\left(\frac{\tilde q\, \tilde k^\top}{\sqrt{d_h}}\right) v.\]

Algorithmic Pseudocode

Key Characteristics

• Training complexity: \(\mathcal{O}(n^2 \cdot d)\) (full attention on decompressed heads).

• Inference complexity: \(\mathcal{O}(n)\) per token, \(\mathcal{O}(r_{kv} \cdot n)\) KV-cache space (latent only).

• Trainable: fully trainable; \(r_{kv}\) is a hyperparameter (Pareto-optimal \(\approx d/2\)).

• Strengths: Halves KV cache versus MHA at matched quality; RoPE-compatible; clean latent bottleneck.

• Limitations: Up-projection adds decode FLOPs; latent rank is a fixed global budget shared across heads.

Group-Query Latent Attention (GQLA)

Group-Query Latent Attention (GQLA) is a minimal modification of MLA that exposes two algebraically equivalent decoding paths over the same trained weights: an MQA-absorb path (fusing the up-projection into the query, optimal on H100-class hardware) and a GQA-expanded path (materialising full keys/values, optimal on H20 / multi-token-prediction hardware). The two paths are mathematically identical because the up-projection is linear, so \(\texttt{softmax}(q_{\text{absorbed}} \cdot c_{KV}^\top) = \texttt{softmax}((q \cdot W_{UK}^\top) \cdot c_{KV}^\top)\), allowing hardware-adaptive dispatch with no retraining.

Mathematical Formulation

Same latent as MLA: \[c_{KV} = W_{DKV}\, x, \qquad k = W_{UK}\, c_{KV}, \qquad v = W_{UV}\, c_{KV}, \qquad q = W_{Q}\, x.\] \[\text{Path A (MQA-absorb): } \text{attn} = \texttt{softmax}\!\left(\frac{(q\, W_{UK}^\top)\, c_{KV}^\top}{\sqrt{d_h}}\right) (W_{UV}\, c_{KV}).\] \[\text{Path B (GQA-expanded): } \text{attn} = \texttt{softmax}\!\left(\frac{q\, k^\top}{\sqrt{d_h}}\right) v, \quad k = W_{UK} c_{KV},\ v = W_{UV} c_{KV}.\]

Algorithmic Pseudocode

Key Characteristics

• Training complexity: \(\mathcal{O}(n^2 \cdot d)\), identical to MLA.

• Inference complexity: \(\mathcal{O}(n)\) per token, \(\mathcal{O}(r_{kv} \cdot n)\) latent state; path chosen by hardware.

• Trainable: trained once; the two decoding paths reuse the same weights.

• Strengths: Hardware-adaptive decode with no quality loss; algebraic equivalence guarantees consistency.

• Limitations: Same latent-rank tradeoffs as MLA; path selection is a deployment-time heuristic.

Multi-Head Low-Rank Attention (MLRA)

Multi-Head Low-Rank Attention (MLRA) refactors the single MLA latent into \(L\) disjoint sub-heads each of rank \(r/L\), so that the KV cache can be partitioned across \(L\) devices for 4-way tensor-parallel decoding (each device loads only \(1/L\) of the cache). Per-sub-head up-projections reconstruct the corresponding slice of the full heads and are concatenated, yielding a 2.8\(\times\) decode speedup over MLA while preserving the global latent budget.

Mathematical Formulation

\[c_{KV} = [c_1, c_2, \ldots, c_L], \quad c_i \in \mathbb{R}^{r/L}, \qquad c_{KV} = W_{DKV}\, x.\] \[k_i = W_{UK}^{(i)}\, c_i, \quad v_i = W_{UV}^{(i)}\, c_i, \qquad k = [k_1; \ldots; k_L], \quad v = [v_1; \ldots; v_L].\] \[\text{attn} = \texttt{softmax}\!\left(\frac{q\, k^\top}{\sqrt{d_h}}\right) v.\]

Algorithmic Pseudocode

Key Characteristics

• Training complexity: \(\mathcal{O}(n^2 \cdot d)\) (full attention on concatenated heads).

• Inference complexity: \(\mathcal{O}(n)\) per token, \(\mathcal{O}(r_{kv} \cdot n)\) state partitioned across \(L\) devices.

• Trainable: fully trainable; \(L\) controls partition granularity.

• Strengths: 2.8\(\times\) decode speedup; clean 4-way tensor-parallel cache sharding; preserves total latent budget.

• Limitations: Requires \(L\) to divide \(r_{kv}\); more projection matrices than MLA; inter-device all-reduce on output.

Tucker Attention

Tucker Attention unifies MHA, GQA and MLA under a single Tucker-style factorisation of the query/key/value weight tensors with independent ranks \(q_{\text{rank}}, k_{\text{rank}}, v_{\text{rank}}\). MHA is recovered when all ranks equal the hidden size, GQA when \(k_{\text{rank}} = v_{\text{rank}} < d\), and MLA when \(k_{\text{rank}} = v_{\text{rank}}\) equals the latent rank. Because the factorisation is a pure linear reparameterisation, the method is fully compatible with FlashAttention and RoPE, and yields an order of magnitude fewer parameters for comparable metrics.

Mathematical Formulation

\[\mathbf{Q} = W_{Q,\text{core}}\, W_{Q,\text{factor}}\, x, \qquad \mathbf{K} = W_{K,\text{core}}\, W_{K,\text{factor}}\, x, \qquad \mathbf{V} = W_{V,\text{core}}\, W_{V,\text{factor}}\, x.\] \[\text{attn} = \texttt{softmax}\!\left(\frac{\mathbf{Q}\, \mathbf{K}^\top}{\sqrt{d_h}}\right) \mathbf{V}.\] Special cases: MHA \(\Leftrightarrow q_{\text{rank}} = k_{\text{rank}} = v_{\text{rank}} = d\); GQA \(\Leftrightarrow k_{\text{rank}} = v_{\text{rank}} < d\); MLA \(\Leftrightarrow k_{\text{rank}} = v_{\text{rank}} = r_{kv}\).

Algorithmic Pseudocode

Key Characteristics

• Training complexity: \(\mathcal{O}(n^2 \cdot d)\) on the reconstructed tensors.

• Inference complexity: \(\mathcal{O}(n)\) per token; cache size governed by \(k_{\text{rank}}\).

• Trainable: ranks are hyperparameters; the factorisation is a drop-in weight reparameterisation.

• Strengths: Unifies MHA/GQA/MLA; order-of-magnitude parameter reduction at matched metrics; FlashAttention/RoPE compatible.

• Limitations: Two matmuls per projection add latency; rank selection is a per-task tuning problem.

Interleaved Head Attention (IHA)

Interleaved Head Attention (IHA) constructs \(P\) pseudo-heads per original head (typically \(P = H\)) where each pseudo query/key/value is a learned linear combination of all \(H\) original heads. This induces up to \(P^2\) distinct attention patterns per head at an \(\mathcal{O}(H^2 P)\) overhead. Theoretically, IHA needs only \(\Theta(\sqrt{k}\, n^2)\) parameters versus \(\Theta(k\, n^2)\) for MHA on the Polynomial task; empirically it delivers \(+10\)–\(20\%\) on RULER multi-key retrieval, \(+5.8\%\) on GSM8K and \(+2.8\%\) on MATH-500.

Mathematical Formulation

\[q_{\text{pseudo}}[p] = \sum_{h=1}^{H} \text{mix}_q[p, h]\, q[h], \quad k_{\text{pseudo}}[p] = \sum_{h} \text{mix}_k[p, h]\, k[h], \quad v_{\text{pseudo}}[p] = \sum_{h} \text{mix}_v[p, h]\, v[h].\] \[\text{attn}_p = \texttt{softmax}\!\left(\frac{q_{\text{pseudo}}[p]\, k_{\text{pseudo}}[p]^\top}{\sqrt{d_h}}\right) v_{\text{pseudo}}[p], \qquad \mathbf{Y}[h] = \frac{1}{P} \sum_{p} \text{attn}_p[h].\]

Algorithmic Pseudocode

Key Characteristics

• Training complexity: \(\mathcal{O}(n^2 \cdot H \cdot P)\) with \(P \approx H\), i.e. \(\mathcal{O}(n^2 \cdot H^2)\).

• Inference complexity: \(\mathcal{O}(n)\) per token with KV cache; extra \(\mathcal{O}(H^2 P)\) mixing overhead.

• Trainable: mixing matrices are learned; \(P\) is a hyperparameter.

• Strengths: \(\Theta(\sqrt{k})\) parameter efficiency on Polynomial; large gains on retrieval and reasoning benchmarks.

• Limitations: Quadratic-in-heads mixing cost; more attention patterns to compute; added latency at large \(H\).

Grouped-head laTenT Attention (GTA)

Grouped-head laTenT Attention (GTA) combines two compressions: (1) a shared attention map — one attention-score tensor per group of heads, reused across all heads in the group, shrinking the key cache; and (2) a nonlinear value decoder — the value cache is compressed to a latent by a down-projection and reconstructed by a SiLU-gated up-projection before the output projection. GTA cuts attention FLOPs by up to 62.5% versus GQA, shrinks the KV cache by up to 70%, and yields a 2\(\times\) end-to-end inference speedup.

Mathematical Formulation

\[q_g = \texttt{mean}(q[\text{group } g]), \quad k_g = \texttt{mean}(k[\text{group } g]), \quad \text{attn}_g = \texttt{softmax}\!\left(\frac{q_g\, k_g^\top}{\sqrt{d_h}}\right) \text{ (reused across the group)}.\] \[v_{\text{latent}} = W_{DV}\, x, \qquad v = \text{SiLU}(W_{UV}\, v_{\text{latent}}), \qquad \mathbf{Y} = \text{attn}_g \cdot v.\]

Algorithmic Pseudocode

Key Characteristics

• Training complexity: \(\mathcal{O}(n^2 \cdot G \cdot d_h)\) where \(G\) is the number of groups (vs \(H\) heads for GQA).

• Inference complexity: \(\mathcal{O}(n)\) per token; value cache \(\mathcal{O}(r_v \cdot n)\), key cache reduced by group sharing.

• Trainable: group assignment and latent rank are trainable/hyperparameters.

• Strengths: up to 62.5% attention-FLOP reduction, 70% KV-cache shrink, 2\(\times\) end-to-end speedup.

• Limitations: Shared map reduces per-head diversity; nonlinear decoder adds latency; grouping must be chosen carefully.

Multi-head Temporal Latent Attention (MTLA)

Multi-head Temporal Latent Attention (MTLA) extends MLA along the temporal axis: a hyper-network dynamically merges \(m\) temporally adjacent KV-cache vectors into a single slot, shrinking the temporal length by a factor of \(m\). A stride-aware causal mask keeps parallel training consistent with autoregressive inference, yielding a 5.3\(\times\) decode speedup and 8.3\(\times\) GPU-memory reduction versus MHA on En–De speech translation.

Mathematical Formulation

\[c_{KV}^{(t)} = W_{DKV}\, x^{(t)} \in \mathbb{R}^{r_{kv}} \text{ (per-token latent)}.\] \[\bar c_{KV}^{(j)} = \sum_{i=0}^{m-1} \alpha_i\, c_{KV}^{(s j + i)}, \qquad \alpha_i = \texttt{sigmoid}(\text{HyperNet}(c_{KV}^{(s j + i)})),\] where \(s\) is the stride and \(m\) the merge window. Queries attend over merged slots \(\bar c_{KV}\) with a stride-aware causal mask \(M_{\text{stride}}\): \[\text{attn} = \texttt{softmax}\!\left(\frac{q\, \bar k^\top}{\sqrt{d_h}} + M_{\text{stride}}\right) \bar v, \quad \bar k = W_{UK}\, \bar c_{KV},\ \bar v = W_{UV}\, \bar c_{KV}.\]

Algorithmic Pseudocode

Key Characteristics

• Training complexity: \(\mathcal{O}(n^2 \cdot d / m)\) on the merged sequence (length \(\lceil n/m \rceil\)) plus hyper-network cost.

• Inference complexity: \(\mathcal{O}(n/m)\) per token over merged slots; \(\mathcal{O}(r_{kv} \cdot n / m)\) state.

• Trainable: hyper-network and merge window \(m\) are learned/chosen; stride-aware mask enforces consistency.

• Strengths: 5.3\(\times\) decode speedup, 8.3\(\times\) GPU-memory reduction vs MHA; temporal compression orthogonal to rank compression.

• Limitations: Merge window can blur fine temporal detail; hyper-network adds compute; mask design is non-trivial.

Architectural Comparison and Synthesis

Compressed Convolutional Attention (CCA)

Compressed Convolutional Attention (CCA) takes a fundamentally different approach to latent attention: rather than compressing the KV-cache and then up-projecting to full width before attending (as MLA/GQLA/MLRA do), CCA down-projects queries, keys, and values into a shared latent and performs the entire attention operation inside it. There are no Q/K/V up-projection matrices—only a single output up-projection \(\tilde{W}_O\) maps the latent output back to the residual stream. This simultaneously reduces parameters, KV-cache size, and attention FLOPs by the compression factor \(C = E/\tilde{e}\), whereas MLA only shrinks the cache.

To make attention in the fully compressed latent viable, CCA introduces three innovations: (1) two causal convolutions on the packed q/k tensor (a depth-wise sequence convolution and a head-wise grouped channel convolution), (2) q-k-mean (adds the pre-convolution mean of q and k to the post-convolution values), and (3) value-shift (each head receives half its values from the current token and half from the previous token via two independent projections). After these steps, QK L2-normalisation and a learnable key temperature \(\beta\) are applied, RoPE is applied directly in the latent (no separate RoPE head needed), and standard softmax attention is computed.

Mathematical Formulation

\[\tilde{q}_{\text{pre}} = \tilde{W}_Q\, x, \qquad \tilde{k}_{\text{pre}} = \tilde{W}_K\, x, \qquad \tilde{e} = E / C.\] \[\tilde{q} = \text{Conv}(\tilde{q}_{\text{pre}}) + \tfrac{1}{2}(\tilde{q}_{\text{pre}} + \tilde{k}_{\text{pre}}), \qquad \tilde{k} = \text{Conv}(\tilde{k}_{\text{pre}}) + \tfrac{1}{2}(\tilde{q}_{\text{pre}} + \tilde{k}_{\text{pre}}).\] \[\tilde{v} = [\tilde{W}_{V_1}\, x_t \;\|\; \tilde{W}_{V_2}\, x_{t-1}], \qquad \hat{q} = \frac{\tilde{q} \sqrt{d_h}}{\|\tilde{q}\|}, \quad \hat{k} = \frac{\tilde{k} \sqrt{d_h}}{\|\tilde{k}\|} e^{\beta}.\] \[\text{attn} = \texttt{softmax}\!\left(\frac{\text{RoPE}(\hat{q})\, \text{RoPE}(\hat{k})^\top}{\sqrt{d_h}}\right) \tilde{v}, \qquad \mathbf{Y} = \tilde{W}_O\, \text{attn}.\]

Algorithmic Pseudocode

Key Characteristics

• Training complexity: \(\mathcal{O}(n^2 \cdot d / C)\) (attention in compressed latent, \(C{\times}\) fewer FLOPs than MHA).

• Inference complexity: \(\mathcal{O}(n)\) per token, \(\mathcal{O}(\tilde{e} \cdot n)\) KV-cache space.

• Trainable: \(C\), conv kernel sizes, and all three tricks (convs, qk-mean, value-shift) are configurable.

• Strengths: Reduces params, cache, AND FLOPs simultaneously; RoPE integrates natively (no RoPE head); \(>\)2\(\times\) fewer params than MLA at same compression.

• Limitations: Does not remove quadratic \(S^2\) (only divides by \(C\)); conv/qk-mean/v-shift add inductive bias; fused kernel needed for theoretical speedup.

Compressed Convolutional Grouped Query Attention (CCGQA)

Compressed Convolutional Grouped Query Attention (CCGQA) extends CCA with GQA-style key/value head sharing applied inside the compressed latent, and decouples the query and KV compression rates: queries are compressed by \(C_1\) and keys/values by \(C_2 \geq C_1\). The per-head latent dimension \(d_h\) must match between query and key heads, enforcing the constraint \(C_2 / C_1 = n_h / n_{kv}\). KV heads are shared per group (replicated to match query heads), and the qk-mean bias uses broadcast/reduce operators \(B_{\text{group}}\) (replicate KV heads to query heads) and \(E_{\text{group}}\) (average query heads within each group). CCGQA achieves the best reported loss at 8\(\times\) KV-cache compression with no quality drop versus MHA.

Mathematical Formulation

Same pipeline as CCA, with decoupled latents: \[\tilde{q}_{\text{pre}} = \tilde{W}_Q\, x \in \mathbb{R}^{E/C_1}, \qquad \tilde{k}_{\text{pre}} = \tilde{W}_K\, x \in \mathbb{R}^{E/C_2}, \qquad C_2 \geq C_1.\] \[\text{qk\_\allowbreak mean}_q = \tfrac{1}{2}(\tilde{q}_{\text{pre}} + B_{\text{group}}(\tilde{k}_{\text{pre}})), \qquad \text{qk\_\allowbreak mean}_k = E_{\text{group}}(\text{qk\_\allowbreak mean}_q).\] \[\text{attn} = \texttt{softmax}\!\left(\frac{\text{RoPE}(\hat{q})\, \text{RoPE}(\hat{k})^\top}{\sqrt{d_h}}\right) \tilde{v}, \qquad \mathbf{Y} = \tilde{W}_O\, \text{attn}.\]

Key Characteristics

• Training complexity: \(\mathcal{O}(n^2 \cdot d / C_1)\) (\(S^2\) term scales by \(1/C_1\); projection terms by \(1/C_1 + 1/C_2\)).

• Inference complexity: \(\mathcal{O}(n)\) per token, \(\mathcal{O}(\tilde{e}_{kv} \cdot n)\) KV-cache space (\(C_2\)-compressed, group-shared).

• Trainable: \(C_1\), \(C_2\), \(n_{kv}\), and all CCA tricks are configurable. Constraint: \(C_2/C_1 = n_h/n_{kv}\).

• Strengths: Decoupled compression enables smooth Pareto frontier; same arithmetic intensity as GQA; best loss at 8\(\times\) cache reduction.

• Limitations: Same quadratic \(S^2\) (only reduced by \(C_1\)); fused kernel needed; two value projections add complexity.

The latent family shows a clear progression along two axes: (i) rank compression, moving from full heads (MHA) to a shared latent (MLA) to partitioned sub-head latents (MLRA), and (ii) deployment adaptivity, moving from a single decode path to hardware-adaptive dispatch (GQLA) and temporal cache merging (MTLA). Tucker Attention provides the unifying algebraic frame: MHA, GQA, and MLA are all special cases of a Tucker factorisation with appropriate rank choices. CCA and CCGQA take a fundamentally different approach: rather than compressing the KV-cache and then up-projecting for full-width attention (MLA/GQLA/MLRA), they perform attention entirely inside the compressed latent space, reducing not just the cache but also the attention FLOPs by the compression factor. CCGQA further decouples query and KV compression rates, achieving the best reported loss at 8\(\times\) KV-cache compression. Gaussian Mixture Attention (GMA) changes the routing philosophy entirely: instead of a low-rank linear bottleneck, it routes through \(K\) learned probabilistic Gaussian components, achieving linear-time \(\mathcal{O}(NK)\) sequence mixing for fixed \(K\) while retaining a normalised attention-style interpretation.

Gaussian Mixture Attention (GMA)

Gaussian Mixture Attention (GMA) replaces the explicit pairwise \(QK^\top\) interaction of standard dot-product attention with probabilistic routing through \(K\) learned Gaussian-mixture components per head. Queries and keys are mapped to posterior responsibility vectors over a shared latent routing space; their overlap defines an implicit responsibility-space affinity, while values are written into and read from a \(K\)-slot latent memory. By exploiting the associativity of matrix multiplication, GMA avoids materialising the induced \(N \times N\) affinity matrix and instead uses two \(N \times K\) responsibility matrices whose dominant activation storage scales as \(\mathcal{O}(NK)\) rather than \(\mathcal{O}(N^2)\) for fixed \(K\). The Gaussian mixture parameters—component means \(\mu\), diagonal covariances (reparameterised via \(\sigma^2 = \text{softplus}(\omega) + \epsilon_\sigma\)) and mixture-prior logits \(\alpha\) (with \(\pi = \text{softmax}(\alpha)\))—are fully differentiable learnable parameters optimised by standard backpropagation, with no EM loop or auxiliary clustering loss.

Mathematical Formulation

Per head, with routing dimension \(d_r\) and value dimension \(d_v\): \[Q_X = X W_Q \in \mathbb{R}^{N \times d_r}, \quad K_X = X W_K \in \mathbb{R}^{N \times d_r}, \quad V_X = X W_V \in \mathbb{R}^{N \times d_v}.\] The posterior responsibility of component \(k\) for a routing vector \(x_i\) is \[\gamma_{i,k} = p(z{=}k \mid x_i) = \frac{\pi_k\, \mathcal{N}(x_i \mid \mu_k, \Sigma_k)}{\sum_{\ell=1}^{K} \pi_\ell\, \mathcal{N}(x_i \mid \mu_\ell, \Sigma_\ell)}, \qquad \Sigma_k = \mathrm{diag}(\sigma_{k,1}^2, \dots, \sigma_{k,d_r}^2),\] yielding responsibility matrices \(\Gamma^Q, \Gamma^K \in \mathbb{R}^{N \times K}\). The bidirectional output is computed by the associative write–read pair \[\tilde V = (\Gamma^K)^\top V_X \in \mathbb{R}^{K \times d_v}, \qquad Z = (\Gamma^K)^\top \mathbf{1}_N \in \mathbb{R}^{K}, \qquad O = \frac{\Gamma^Q \tilde V}{\Gamma^Q Z + \epsilon},\] which never materialises the implicit normalised affinity \(A^{\text{GMA}}_{ij} = \frac{\sum_k \gamma^Q_{i,k} \gamma^K_{j,k}}{\sum_\ell \sum_k \gamma^Q_{i,k} \gamma^K_{\ell,k} + \epsilon}\). For autoregressive modelling the global write statistics are replaced by prefix-restricted cumulative sums \(\tilde V^{(i)}_k = \sum_{j \le i} \gamma^K_{j,k} V_{X,j}\) and \(Z^{(i)}_k = \sum_{j \le i} \gamma^K_{j,k}\), enforcing causality while preserving the fixed-\(K\) linear scaling.

Algorithmic Pseudocode

Key Characteristics

• Training complexity: \(\mathcal{O}(N K d_r + N K d_v)\) for fixed \(K\), linear in the sequence length \(N\).

• Inference complexity: \(\mathcal{O}(K)\) per token (responsibilities + latent read), \(\mathcal{O}(NK)\) activation storage.

• Trainable: fully differentiable GMM parameters \(\{\mu, \omega, \alpha\}\) per head, optimised end-to-end without EM.

• Strengths: Fixed-\(K\) linear-time sequence mixing; probabilistic interpretable routing; bidirectional and causal variants; non-negative low-rank affinity interpretation \(A^{\text{GMA}} = \mathrm{diag}(d)^{-1} \Gamma^Q (\Gamma^K)^\top\).

• Limitations: Quality depends on \(K\) and the routing dimension \(d_r\); causal GMA trails optimised SDPA and state-space models on WikiText-103 in the paper’s current implementation; not a universal replacement for softmax attention.

Kronecker Product Attention (KPA)

Kronecker Product Attention (KPA) is a KV-compression mixer that factorises the per-token key, value and query matrices as normalized sums of Kronecker products of compact contextual factors. Unlike Tensor Product Attention (TPA), whose CP decomposition operates on vectors, KPA acts on matrices and therefore encodes structure at both the head and the feature level; rank-one CP factorisation is recovered as the degenerate split \(h_1=h, d_1=1, h_2=1, d_2=d_h\), so KPA strictly generalises TPA at the same parameter budget.

Mathematical Foundation

Let \(h = h_1 h_2\) be the number of heads and \(d_h = d_1 d_2\) the head width. For a token with hidden state \(x_t \in \mathbb{R}^{d_{\text{model}}}\), the Kronecker product obeys the mixed-product identity \[(M \otimes N)(P \otimes Q) = (MP) \otimes (NQ),\] and the three projections are \[Q_t = \frac{1}{R_Q}\sum_{r=1}^{R_Q} A^Q_r(x_t)\otimes B^Q_r(x_t),\quad K_t = \frac{1}{R_K}\sum_{r=1}^{R_K} A^K_r(x_t)\otimes B^K_r(x_t),\quad V_t = \frac{1}{R_V}\sum_{r=1}^{R_V} A^V_r(x_t)\otimes B^V_r(x_t),\] with \(A^\cdot_r(x_t) = W^{A^\cdot}_r x_t \in \mathbb{R}^{h_1\times d_1}\) and \(B^\cdot_r(x_t) = W^{B^\cdot}_r x_t \in \mathbb{R}^{h_2\times d_2}\) produced by independent bias-free linear maps; the rank index is merged into a single batched projection per factor kind. Heads are indexed by pairs \((i_1,i_2)\) with \(i_1\in[h_1]\), \(i_2\in[h_2]\), and the mixed-product identity factorises every attention score into a head-level and a feature-level inner product: \[\mathrm{score}_{(i_1,i_2)}(t,s) = \frac{1}{R_Q R_K}\sum_{r=1}^{R_Q}\sum_{r'=1}^{R_K} \bigl[(A^Q_r)_{i_1}^{\top}(A^K_{r'})_{i_1}\bigr]\, \bigl[(B^Q_r)_{i_2}^{\top}(B^K_{r'})_{i_2}\bigr].\] Kronecker-Separable RoPE uses the operator \(R_{\mathrm{KPA}}(t) = I_{d_1}\otimes R^{(2)}_\theta(t)\) with frequencies \(\theta_j = 10000^{-2j/d_2}\) calibrated to \(d_2\). Because \((I_{d_1}\otimes R^{(2)}_\theta(t))((A_r)_{i_1}\otimes(B_r)_{i_2}) = (A_r)_{i_1}\otimes R^{(2)}_\theta(t)(B_r)_{i_2}\), only the \(B\) factor is rotated, the \(A\) factor is cached unrotated, and the relative-position property follows from \(R^{(2)}_\theta(t)^{\top}R^{(2)}_\theta(s) = R^{(2)}_\theta(s-t)\); no decoupled RoPE branch is needed. The cache stores \(\{A^K_r, B^K_r\}_{r=1}^{R_K}\cup\{A^V_r, B^V_r\}_{r=1}^{R_V}\) per token (with \(B^K\) pre-rotated), at a cost of \((R_K+R_V)(h_1d_1+h_2d_2)\) floats instead of \(2h_1h_2d_1d_2\) for MHA, i.e. a compression factor \[C_{\mathrm{KPA}} = \frac{2h_1h_2d_1d_2}{(R_K+R_V)(h_1d_1+h_2d_2)},\] which increases with factor asymmetry.

Algorithmic Pseudocode

Key Characteristics

• Training complexity: \(\mathcal{O}(n^2 d)\) attention plus \(\mathcal{O}\big(n\,d\,(R_Q+R_K+R_V)(h_1d_1+h_2d_2)\big)\) factor projections.

• Inference complexity: \(\mathcal{O}(n\,h\,d_h)\) attention per step and \(\mathcal{O}\big((R_K+R_V)(h_1d_1+h_2d_2)\big)\) cached state per token; \(6.4\times\)–\(10\times\) KV compression in the paper.

• Trainable: trained from scratch (nanoGPT + SwiGLU + RoPE) at 125M, 353M, 773M and 1.5B parameters on FineWeb-Edu for \(\sim\)50B tokens.

• Strengths: strict generalisation of TPA with richer block structure at matched parameters; RoPE-compatible without MLA’s decoupled branch; the head split \(h_1\times h_2\) is a design-time compression/accuracy knob with no architectural change.

• Limitations: the paper fixes \(d_1=d_2=8\) and evaluates only causal language models; an efficient factored decoding kernel and long-context evaluation are left as future work; materialising the Kronecker sums is quadratic in the sequence length like dense attention.

Tensor Product Attention (TPA)

Tensor Product Attention (TPA) is the canonical activation-space tensor decomposition for attention: rather than factorising the projection weights (as Tucker Attention does), it factorises the per-token query, key and value activations themselves into rank-one outer products of contextual vectors. It is the vector (CP) special case of KPA , recovered at \(h_1=h\), \(d_1=1\), \(h_2=1\), \(d_2=d_h\).

Mathematical Foundation

With \(h\) heads, head width \(d_h\) and CP ranks \(R_Q,R_K,R_V\), each token’s slice is \[Q_t = \frac{1}{R_Q}\sum_{r=1}^{R_Q} a^Q_r(x_t)\otimes b^Q_r(x_t),\quad K_t = \frac{1}{R_K}\sum_{r=1}^{R_K} a^K_r(x_t)\otimes b^K_r(x_t),\quad V_t = \frac{1}{R_V}\sum_{r=1}^{R_V} a^V_r(x_t)\otimes b^V_r(x_t),\] with \(a^\cdot_r(x_t)=W^{a^\cdot}_r x_t\in\mathbb{R}^{h}\) and \(b^\cdot_r(x_t)=W^{b^\cdot}_r x_t\in\mathbb{R}^{d_h}\) produced by bias-free linear maps. Merging the rank index into one projection per factor side, \(a_Q(x_t)=W^{aQ}x_t\in\mathbb{R}^{R_Q h}\) and \(b_Q(x_t)=W^{bQ}x_t\in\mathbb{R}^{R_Q d_h}\), and reshaping into \(A_Q(x_t)\in\mathbb{R}^{R_Q\times h}\), \(B_Q(x_t)\in\mathbb{R}^{R_Q\times d_h}\) gives the compact form \[Q_t = \frac{1}{R_Q}A_Q(x_t)^{\top}B_Q(x_t)\in\mathbb{R}^{h\times d_h},\] and attention is the ordinary scaled dot-product map \[\mathrm{head}_i = \mathrm{Softmax}\!\left(\frac{Q_iK_i^{\top}}{\sqrt{d_h}}\right)V_i,\qquad \mathrm{TPA}(Q,K,V)=\mathrm{Concat}(\mathrm{head}_1,\dots,\mathrm{head}_h)\,W^O .\] Because RoPE rotates only the head-width axis it commutes with the factorisation (Theorem 3.1 of the paper): \(B^\cdot(x_t)\mapsto B^\cdot(x_t)T_t\) row-wise and \(\widetilde{Q}_t=Q_tT_t=\frac{1}{R_Q}A_Q(x_t)^{\top}\widetilde{B}_Q(x_t)\), preserving the relative-position identity \(\widetilde{Q}_t\widetilde{K}_s^{\top}=Q_tT_{t-s}K_s^{\top}\). The cache holds \(\{A^K,\widetilde{B}^K,A^V,B^V\}\) per token, i.e. \[\mathrm{KV}_{\mathrm{TPA}}=(R_K+R_V)(h+d_h)\ \text{floats/token} \quad\text{vs.}\quad 2hd_h\ \text{for MHA},\] and the attention-layer parameter count (Table 1 of the paper) is \[P_{\mathrm{TPA}} = d_{\text{model}}\,(R_Q+R_K+R_V)(h+d_h) + d_{\text{model}}\,h\,d_h .\]

Algorithmic Pseudocode

A fused factor-space variant (FlashTPA, Appendix A of the paper) computes the head-shared feature dots \(P[b,q,k,r,s]=\langle B^Q_{q,r},B^K_{k,s}\rangle\) first, mixes them with \(A^Q,A^K\) into the logits, and aggregates values through \(A^V,B^V\) without ever materialising \(Q,K,V\); it reduces FLOPs whenever \(R_QR_KD + hR_QR_K + hR_VE < 2hd_h\). This implementation uses the materialised path and keeps the factor-space variant as a kernel-level optimisation.

Key Characteristics

• Training complexity: \(\mathcal{O}(n^2 d)\) attention plus \(\mathcal{O}\big(n\,d\,(R_Q+R_K+R_V)(h+d_h)\big)\) factor projections; the specialised path is \(\mathcal{O}(n^2R_QR_Kd_h + n^2hR_QR_K + n^2hR_Vd_h)\).

• Inference complexity: \(\mathcal{O}(M R_QR_Kd_h + MhR_QR_K + MhR_Vd_h)\) per decode step against a cache of length \(M\); \((R_K+R_V)(h+d_h)\) cached floats per token, an order of magnitude below MHA at \(R_K,R_V\le 2\).

• Trainable: trained from scratch; T6 composes TPA with RMSNorm and SwiGLU blocks and is pretrained with next-token prediction.

• Strengths: native RoPE (and generic position-encoding) compatibility without MLA’s decoupled branch; better validation perplexity and downstream scores than MHA, MQA, GQA and MLA at matched parameter counts; MHA/MQA/GQA are recovered as non-contextual rank-\(h\)/rank-1/rank-\(G\) special cases, which makes the rank a single design knob.

• Limitations: extra rank hyperparameters (\(R_Q,R_K,R_V\)) and a factorisation cost on the materialised path; quality depends on rank choice; the paper evaluates causal language models only; the full speed-up requires the FlashTPA kernel.

Multi-matrix Factorization Attention (MFA)

Multi-matrix Factorization Attention (MFA) is the capacity-oriented counterpart of MLA: instead of compressing a full multi-head key–value pair into one latent, it keeps many high-rank heads over a single shared key/value head. The paper reads every MHA variant as a low-rank factorisation of Fully Parameterized Bilinear Attention (FPBA) and shows that the head count and the per-head rank, not only the compression rank, drive capacity. MFA factorises the QK circuit as \(W_c\approx S_qQ_cS_k^{\top}\) and the VO circuit as \(U_c\approx S_vO_c^{\top}\), with \(S_q,S_k,S_v\) shared across heads and only \(Q_c,O_c\) head-specific. The Total Effective Rank is \(mC\) while the Shared Latent Subspace Dimension stays at \(C\), so scaling the head count costs about \(CH\) parameters per head and leaves the cache untouched.

Mathematical Foundation

The inference and factorisation forms of the paper are \[O_i=\sum_{c=1}^{m}\Big(\sum_{j=1}^{i}\phi\Big(\frac{x_iS_qQ_c\,(x_jS_k)^{\top}}{\sqrt{d}}\Big)x_jS_v\Big)O_c^{\top} =\sum_{c=1}^{m}\Big(\sum_{j=1}^{i}\phi\Big(\frac{w_i\,x_j^{\top}}{\sqrt{d}}\Big)x_jS_vO_c^{\top}\Big),\] with \(w_i=x_iS_qQ_cS_k^{\top}\) and \(\phi\) the softmax. Keys and values are single-head, \(k_j=x_jS_k,\;v_j=x_jS_v\in\mathbb{R}^{C}\), so decoding stores \[\mathrm{KV}_{\mathrm{MFA}}=2C\ \text{floats/token}\quad(\text{independent of }m),\] and the attention-layer parameter count (Table 1 of the paper) is \[P_{\mathrm{MFA}}=H(3C+mC)+mC^{2}=3HC+mC^{2}+mCH .\] The MFA-Key-Reuse variant reparameterises the value projection from the key projection, \[S_v=S_k+\alpha\odot NS_k=(I+\mathrm{diag}(\alpha)N)W_K,\qquad N\in\mathbb{R}^{C\times C},\ \alpha\in\mathbb{R}^{C},\] with \(\alpha\) initialised to zero so that \(S_v=S_k\) at the start of training, which the authors identify as the key stability ingredient (Table 4: a vanilla key-reuse variant loses 0.20 nats, an extra value projection is unstable at 7.88 ppl, the residual variant recovers 6.65, and the zero-initialised gate reaches 6.45 against MFA’s 6.35). Decoding then caches only \(k\), halving the footprint to \(C\) floats per token: \[\mathrm{KV}_{\mathrm{MFA\text{-}KR}}=C\ \text{floats/token},\qquad P_{\mathrm{MFA\text{-}KR}}=2HC+mC^{2}+mCH+C^{2}+C .\] Because \(S_k\) and \(S_v\) are ordinary single-head projections, RoPE applies row-wise to the \(C\)-dim queries and keys directly; unlike MLA there is no decoupled rotary branch, and the paper validates both RoPE and ALiBi (Table 5).

Algorithmic Pseudocode

At decode, \(k_t\) and \(v_t\) are appended to the shared cache (only \(k_t\) under MFA-KR), and the per-step cost is \(\mathcal{O}(m\,T\,C)\) against a cache of width \(C\).

Key Characteristics

• Training complexity: \(\mathcal{O}(n^{2}mC)\) attention plus \(\mathcal{O}(nH(3C+mC))\) projections.

• Inference complexity: \(\mathcal{O}(mTC)\) per decode step; \(2C\) cached floats per token (\(C\) under MFA-KR), independent of the head count.

• Trainable: trained from scratch; 6.9B-parameter MoE models over 1T tokens, with the dense ablation at 1B/20B tokens.

• Strengths: KV cache decoupled from head count; each head keeps rank \(C>d\), so capacity scales with \(mC\) at \(\approx CH\) parameters per extra head; native RoPE/ALiBi compatibility without MLA’s decoupled branch; outperforms MHA (49.9 vs. 49.0 average accuracy) with 87.5% less cache, and MFA-KR reaches 93.7% with a minor trade-off.

• Limitations: key and value share one \(C\)-dim bottleneck per token, so the shared latent subspace bounds quality; MFA-KR’s key reuse costs up to 0.1 nats against MFA; the paper evaluates causal language models only and does not measure system-level latency.

Comprehensive Sparse Attention Mechanisms

Executive Summary

Sparse attention mechanisms address the prohibitive \(\mathcal{O}(n^2)\) computational and memory complexity of standard scaled dot-product attention by restricting the set of attended key-value pairs. Modern sparse attention spans a rich design space distinguished by four orthogonal dimensions: (i) sparsity pattern (fixed geometric, data-dependent, or frequency-based), (ii) trainability (end-to-end trained or inference-only), (iii) sparsity unit (individual tokens, local windows, or blocks), and (iv) mechanism (masking, selection, filtering, or compression). This annex synthesizes eleven major sparse attention methods—Sparse Transformer, Longformer, SWA with attention sinks, BigBird, SparseK, NSA, FASA, SpargeAttn, MSA, SparDA, and LSA—providing mathematical formulations, algorithmic pseudocode, and comprehensive architectural comparisons. Fast Conformer, the limited-context variant of the convolution-augmented Conformer block, belongs to the convolution-augmented family and is documented alongside the Conformer in Appendix 19.

Sparse Transformer: Factorized Strided and Fixed Patterns

Mathematical Formulation

The Sparse Transformer reduces quadratic complexity to \(\mathcal{O}(n\sqrt{n})\) by factorizing sparse attention into two complementary sparse heads. Let \(n\) be the sequence length and \(l = \lfloor\sqrt{n}\rfloor\) be the stride.

Strided attention head

: Each position \(i\) attends to every \(l\)-th previous position: \[\mathcal{A}_i^{(\text{strided})} = \{j : j \leq i,\; (i - j) \bmod l = 0\}\]

Fixed attention head

: Each position \(i\) attends to positions within its local block plus fixed summary columns: \[\mathcal{A}_i^{(\text{fixed})} = \{j : \lfloor j/l \rfloor = \lfloor i/l \rfloor\} \cup \{j : j \bmod l \in \{l-c, \ldots, l-1\}\}\] where \(c\) is a hyperparameter controlling the number of summary columns (typically \(c=1\) or \(c=2\)).

The attention computation for each head \(h\) follows standard scaled dot-product rules restricted to the sparse set: \[\text{Attn}_h(Q, K, V)_i = \text{softmax}\left(\frac{q_i^h K_{\mathcal{A}_i}^h\top}{\sqrt{d_k}}\right) V_{\mathcal{A}_i}^h\]

The key insight is that with the two factorized heads operating in parallel across a depth of \(L\) transformer layers, every position can reach every other position through a path of length at most \(L + 1\), preserving long-range reachability with a fraction of dense attention cost.

Algorithmic Pseudocode

Key Characteristics

• Complexity: \(\mathcal{O}(n\sqrt{n} \cdot d)\) during training; \(\mathcal{O}(n)\) per token during generation.

• Trainable: Yes; full backpropagation through selected positions.

• Sparsity pattern: Fixed and data-agnostic; patterns do not adapt to content.

• Strength: Structured reachability guarantees long-range dependencies; proven effective on long sequences (Enwik8, text images).

• Limitation: Fixed patterns may miss important but non-contiguous relationships; requires custom CUDA kernels for practical efficiency.

Longformer: Sliding Window, Dilation, and Global Tokens

Mathematical Formulation

Longformer achieves linear \(\mathcal{O}(n \cdot w)\) complexity by combining three complementary attention patterns.

Sliding window attention

: Local neighborhood of size \(w\): \[\mathcal{A}_i^{(\text{window})} = \{j : |i - j| \leq w/2\}\]

Dilated sliding window

: Windowed attention with dilation \(d\), skipping stride \(d+1\): \[\mathcal{A}_i^{(\text{dilated})} = \{j : |i - j| \leq (w/2)(d+1) \text{ and } (i-j) \bmod (d+1) = 0\}\]

Global attention

: Designated global tokens (e.g., \([\text{CLS}]\)) attend to all positions and are attended by all: \[\mathcal{A}_i^{(\text{global})} = \begin{cases} \{1, \ldots, n\} & \text{if } i \in \mathcal{G} \\ \mathcal{A}_i^{(\text{window})} \cup \mathcal{G} & \text{otherwise} \end{cases}\]

The three patterns are applied via separate projections. The local and global attention can use either the same projections or task-specific separate ones.

Algorithmic Pseudocode

Key Characteristics

• Complexity: \(\mathcal{O}(n \cdot w)\) where \(w\) is the window size (linear when \(w \ll n\)).

• Trainable: Yes; drop-in replacement for standard attention.

• Coverage: With \(L\) layers and window size \(w\), receptive field grows to \(L \times w\), covering entire sequence with shallow networks.

• Strength: Practical for document-level tasks; dilation expands local receptive fields with minimal overhead; global tokens provide query-document correspondence.

• Limitation: Window size remains a design hyperparameter; suboptimal for tasks requiring dense attention patterns.

SWA with Attention Sinks: Sliding Window Plus a Sink Prefix

Mathematical Formulation

SWA with attention sinks is the simplest KV-cache eviction policy: instead of compressing the cache or replacing attention with a recurrent state, it restricts every query to a small, fixed set of keys. Write the window size as \(w\) and the number of sinks as \(s\), i.e. \(\mathrm{SWA}(w,s)\). In decoder mode, query \(i\) attends to \[\mathcal{A}_i = \underbrace{\{j : j < s,\; j \leq i\}}_{\text{attention sinks}} \;\cup\; \underbrace{\{j : i - (w-s) + 1 \leq j \leq i\}}_{\text{sliding window of } w-s},\] so each query sees at most \(w\) keys. The causal clip \(j \leq i\) on the sink term is essential: tokens before position \(s\) must not see future sinks.

Why sinks are indispensable

Transformers learn to assign disproportionately large attention mass to the first few tokens even when they are not semantically relevant. These attention sinks act as repositories for surplus attention. Once the prefix leaves a sink-free sliding window, the model loses its dump target and quality collapses catastrophically; re-including the first \(s\) tokens restores most of full-attention performance. The paper fixes \(s = 4\) and evaluates \(w \in \{64, 128, 256, 512\}\).

Encoder extension

The paper defines \(\mathrm{SWA}(w,s)\) for causal language models. For bidirectional encoders the local window is made symmetric, \(\{j : |i-j| \leq \lfloor (w-s)/2 \rfloor\}\), and unioned with the sink prefix—the same convention used by Longformer.

Because the intervention is a pure attention-mask change, it requires no additional parameters and no new QKV math: the standard projections \(W_q, W_k, W_v, W_o\) are reused. It can be applied to a pretrained checkpoint at inference time with zero training tokens, or trained from scratch as a normal mixer.

Algorithmic Pseudocode

Key Characteristics

• Complexity: \(\mathcal{O}(n \cdot w \cdot d)\) time and \(\mathcal{O}(w)\) KV storage per layer at inference (fixed once the sequence exceeds \(w\)).

• Trainable: Both; training-free at inference on pretrained checkpoints, or trained from scratch.

• Receptive field: After \(L\) layers, \(L(w-s)\) positions plus the always-visible sink prefix.

• Strength: Recovers \(\approx 99\%\) of the average full-attention score across models from 1.3B to 72B with zero training tokens, and beats retrofitted linear attention by 2–10\(\times\) on long-context reasoning (S-NIAH, BABILong) while decoding faster and using less memory.

• Limitation: The horizon is fixed and task-dependent—\(w=64\) suffices for log-likelihood benchmarks but collapses on generative math, where \(w=512\) is needed; there is no content-adaptive selection.

BigBird: Random, Local, and Global Sparse Graph

Mathematical Formulation

BigBird preserves universal approximation and Turing completeness using a principled mix of three sparsity patterns.

Random connections

: Each position connects to \(r\) randomly sampled positions: \[\mathcal{A}_i^{(\text{random})} = \text{RandomSample}(\{1, \ldots, n\}, r)\]

Local window

: Sliding window neighborhood: \[\mathcal{A}_i^{(\text{local})} = \{j : |i-j| \leq w/2\}\]

Global tokens

: Task-specific global anchors: \[\mathcal{A}_i^{(\text{global})} = \begin{cases} \{1, \ldots, n\} & \text{if } i \in \mathcal{G} \\ \mathcal{A}_i^{(\text{random})} \cup \mathcal{A}_i^{(\text{local})} \cup \mathcal{G} & \text{otherwise} \end{cases}\]

The composite sparse mask \(M\) is: \[M_{ij} = \mathbf{1}[j \in \mathcal{A}_i^{(\text{random})} \cup \mathcal{A}_i^{(\text{local})} \cup \mathcal{A}_i^{(\text{global})}]\]

In practice, BigBird groups tokens into blocks of size \(b\) and applies the sparsity decision at the block level, enabling efficient block-sparse implementations.

Theoretical Guarantee

The authors prove that with \(O(1)\) random connections and global tokens, the sparse graph maintains the same universal approximation and Turing completeness properties as dense attention. Formally, any computable function can be represented and any sequence of operations can be simulated within the BigBird connectivity graph.

Algorithmic Pseudocode

Key Characteristics

• Complexity: \(\mathcal{O}(n)\) or near-linear in optimal block-sparse implementation.

• Trainable: Yes.

• Theoretical grounding: Provably maintains universal approximation and Turing-completeness with sparse connectivity.

• Strength: Strong long-document performance on question answering, summarization, and genomics tasks.

• Limitation: Random sampling introduces variance and non-determinism; tuning of \(r, w, g\) hyperparameters remains task-dependent.

SparseK Attention: Differentiable Top-\(k\) Selection

Mathematical Formulation

SparseK enables learnable sparsity via a differentiable top-\(k\) operator. A scoring network \(\phi_\theta\) evaluates key-value pair importance:

\[u_j = \phi_\theta(\mathbf{k}_j, \mathbf{q}_i) \in \mathbb{R}\]

The SparseK operator selects the top-\(k\) scores while remaining differentiable. It computes a threshold \(\tau(u)\) such that the sum of active scores equals \(k\):

\[\text{SparseK}(u, k)_j = \max(u_j - \tau(u), 0), \quad \text{where} \quad \sum_j \max(u_j - \tau, 0) = k\]

The threshold \(\tau\) is found via bisection. The attention becomes:

\[m_j = \text{SparseK}(u, k)_j, \quad \text{Attn}_i = \text{softmax}\left(\frac{\mathbf{q}_i K_{\text{sel}}^\top}{\sqrt{d_k}}\right) V_{\text{sel}}\]

where \(K_{\text{sel}}, V_{\text{sel}}\) contain only the top-\(k\) entries (those with \(m_j > 0\)).

Algorithmic Pseudocode

Key Characteristics

• Complexity: \(\mathcal{O}(n \cdot d)\) during training (linear in sequence length); \(\mathcal{O}(k)\) per token during generation.

• Trainable: Yes; end-to-end gradient flow through the SparseK operator.

• Incremental generation: Supports efficient constant-memory autoregressive generation.

• Strength: Seamlessly integrates into existing LLM architectures; minimal fine-tuning needed.

• Limitation: Scattered memory access from top-\(k\) selection may limit cache efficiency on some hardware; scoring network adds overhead.

NSA (Native Sparse Attention): Hardware-Aligned Hierarchical Branches

Mathematical Formulation

NSA decomposes sparse attention into three parallel branches that are combined through learned gates.

Compression branch

: A learnable MLP \(\varphi\) compresses key-value blocks:

\[\tilde{K}_t^{\text{cmp}} = \left[\varphi(k_{id+1:id+l})\right]_{\substack{1 \leq i \leq \lfloor(t-l)/d\rfloor}} \quad, \quad \tilde{V}_t^{\text{cmp}} = \left[\varphi(v_{id+1:id+l})\right]\]

Selection branch

: High-importance blocks are selected based on compressed scores:

\[p_t^{\text{cmp}} = \text{softmax}(q_t^\top \tilde{K}_t^{\text{cmp}} / \sqrt{d}), \quad I_t = \mathrm{TopK}(p_t^{\text{cmp}}, n), \quad \tilde{K}_t^{\text{sel}} = \mathrm{Gather}(K, I_t)\]

Sliding window branch

: Fixed local context:

\[\tilde{K}_t^{\text{win}} = K_{t-w:t}, \quad \tilde{V}_t^{\text{win}} = V_{t-w:t}\]

Gated combination

: The three attention outputs are combined via learned gates:

\[\alpha_t^{(\text{cmp})} = \sigma(\text{Linear}_\text{cmp}(q_t)), \quad \alpha_t^{(\text{sel})} = \sigma(\text{Linear}_\text{sel}(q_t)), \quad \alpha_t^{(\text{win})} = \sigma(\text{Linear}_\text{win}(q_t))\]

\[y_t = \alpha_t^{(\text{cmp})} \cdot \text{Attn}(q_t, \tilde{K}^{\text{cmp}}, \tilde{V}^{\text{cmp}}) + \alpha_t^{(\text{sel})} \cdot \text{Attn}(q_t, \tilde{K}^{\text{sel}}, \tilde{V}^{\text{sel}}) + \alpha_t^{(\text{win})} \cdot \text{Attn}(q_t, \tilde{K}^{\text{win}}, \tilde{V}^{\text{win}})\]

Algorithmic Pseudocode

Key Characteristics

• Complexity: \(\mathcal{O}(t/d + n \cdot l + w)\) tokens attended per step (significantly reduced versus \(\mathcal{O}(t)\) for dense).

• Trainable: Yes; direct training with all branches differentiable.

• Hardware alignment: Designed for efficient tensor core utilization; compression and selection reduce memory bandwidth pressure.

• Strength: 11.6\(\times\) decoding speedup and 9\(\times\) forward speedup on 64k sequences; outperforms full attention on many tasks.

• Limitation: Requires custom Triton/CUDA kernels; complex multi-branch architecture increases engineering overhead.

FASA: Frequency-Aware Sparse Attention

Mathematical Formulation

FASA is a training-free inference-time method that exploits rotary position embedding (RoPE) structure. Under RoPE, the token embedding is rotated by \(\theta_i = B^{-2(i-1)/d}\), decomposing the \(d\)-dimensional space into \(d/2\) frequency chunks (FCs).

Key insight

: Only a small subset (\(< 1\%\)) of FCs matter for contextual awareness; most encode positional patterns.

Dominant frequency chunk identification

: For each layer \(l\) and head \(h\), identify dominant FCs via contextual agreement (CA):

\[\text{CA}^{l,h,i} = \frac{|\text{TopK}(\alpha^{l,h}) \cap \text{TopK}(\alpha^{l,h,i})|}{K}\]

where \(\alpha^{l,h}\) is the full attention mask and \(\alpha^{l,h,i}\) is the mask using only FC \(i\). Dominant FCs have high CA—they contribute meaningfully to the final attention pattern.

Token importance prediction (TIP) stage

: Using only dominant FCs, compute lightweight importance per token:

\[S_t^{l,h} = \sum_{i \in \mathcal{I}_{\text{dom}}^{l,h}} \alpha^{l,h,i}(\mathbf{q}_t, \mathbf{K}_{1:t}), \quad \mathcal{T}_t = \text{TopK-Indices}(S_t, N_{\text{fac}})\]

Focused attention computation (FAC) stage

: Compute full-precision attention on selected tokens:

\[\hat{\alpha}_{\text{FAC}} = \text{softmax}\left(\frac{\mathbf{q}_t \mathbf{K}_{\mathcal{T}_t}^\top}{\sqrt{d}}\right), \quad \mathbf{o}_t = \hat{\alpha}_{\text{FAC}} \mathbf{V}_{\mathcal{T}_t}\]

Algorithmic Pseudocode

Key Characteristics

• Complexity: TIP is \(\mathcal{O}(t \cdot N_{\text{tip}})\); FAC is \(\mathcal{O}(N_{\text{fac}} \cdot d)\).

• Trainable: No; training-free inference optimization.

• Applicability: Requires RoPE-based models; extended to ALiBi and MLA with modifications.

• Strength: Near-oracle accuracy with \(\leq 256\) tokens out of millions; up to 2.56\(\times\) decoding speedup; orthogonal to quantization.

• Limitation: Offline offline calibration required; dominant FC identification adds preprocessing overhead.

Implementation note (training vs. inference).

In this repository the sparse TIP/FAC path is inference-only. While the model is in training mode, the enclosing hybrid layer substitutes standard dense attention over the same \(\mathbf{q}/\mathbf{k}/\mathbf{v}/\mathbf{o}\) projections, so the mixer can be trained from scratch and the weights optimized during training are exactly the ones the sparse path reads at evaluation.

SpargeAttn: Two-Stage Block-Level Filtering

Mathematical Formulation

SpargeAttn is a universal training-free method that predicts which blocks of the attention matrix will contain negligible values and skips computation for those blocks.

Stage 1 — Sparse prediction

: Compute low-cost proxy importance \(\hat{s}_{ij}\) for block \((i,j)\):

\[\hat{s}_{ij} = f_{\text{pred}}(\mathbf{Q}_i, \mathbf{K}_j; \text{hyperparams}), \quad \text{skip if } \hat{s}_{ij} < \epsilon_1\]

Common predictors include block-mean similarity or self-similarity statistics.

Stage 2 — Softmax-aware filtering

: After computing \(\tilde{P}_{ij} = \text{softmax}(Q_i K_j^\top/\sqrt{d})\), check if the block’s maximum probability is negligible relative to the running softmax maximum:

\[\text{skip } P_{ij} V_j \quad \text{if} \quad \max(\tilde{P}_{ij}) < e^{m_{\text{old}} - m_{\text{new}}} \cdot \epsilon_2\]

where \(m_{\text{old}}, m_{\text{new}}\) are online softmax running maxima (computed via FlashAttention-style numerics).

Algorithmic Pseudocode

Key Characteristics

• Complexity: Empirical \(\mathcal{O}(n^2 \cdot s)\) where \(s\) is the fraction of non-skipped blocks (typically 0.2–0.5).

• Trainable: No; plug-and-play acceleration for existing models.

• Universality: Works on language models, image diffusion models, and video generation.

• Strength: 2.5–5\(\times\) speedup compared to dense or previous sparse methods; compatible with quantization (SageAttention integration).

• Limitation: Speedup depends on inherent sparsity; block-level granularity may miss finer patterns; threshold tuning required per model family.

Implementation note (training vs. inference).

In this repository the two-stage block filter is inference-only. While the model is in training mode, the enclosing hybrid layer substitutes standard dense attention over the same \(\mathbf{q}/\mathbf{k}/\mathbf{v}/\mathbf{o}\) projections, so the mixer can be trained from scratch and the weights optimized during training are exactly the ones the sparse path reads at evaluation.

MSA (MiniMax Sparse Attention): Blockwise Sparse Attention with Lightweight Index Branch

Mathematical Formulation

MSA is a blockwise sparse attention mechanism built on grouped-query attention (GQA) that combines a lightweight Index Branch for block selection with a Main Branch for exact block-sparse softmax attention over the selected blocks. Let \(N\) be the sequence length, \(B_k = 128\) the block size, and \(k = 16\) the number of blocks selected per query (yielding a fixed per-query budget of \(k \cdot B_k = 2048\) tokens regardless of \(N\)).

Index Branch

: A low-dimensional projection produces index queries \(\mathbf{q}_t^{\text{idx}} \in \mathbb{R}^{d_{\text{idx}}}\) and block representatives \(\mathbf{r}_b \in \mathbb{R}^{d_{\text{idx}}}\) for each KV block \(b\). The index scores per GQA group are: \[s_{t,b}^{(g)} = \mathbf{q}_t^{\text{idx},(g)} \cdot \mathbf{r}_b, \qquad \mathcal{I}_t^{(g)} = \mathrm{TopK}\left(\{s_{t,b}^{(g)}\}_{b=1}^{N/B_k},\; k\right)\] The index input is detached (stop-gradient), so gradients do not flow back through the indexer’s input; the Index Branch is trained instead by a KL alignment loss between its attention distribution and the Main Branch’s distribution. Because softmax is order-preserving, the raw index scores feed directly into the top-\(k\) operator without an explicit softmax (exp-free top-\(k\)), reducing numerical cost.

Forced local block

: The block containing the query position \(t\) is always included in \(\mathcal{I}_t^{(g)}\), guaranteeing local context coverage.

Main Branch

: Standard scaled dot-product attention is computed exactly over the union of selected blocks: \[\mathbf{o}_t^{(h)} = \mathrm{softmax}\left(\frac{\mathbf{q}_t^{(h)} \mathbf{K}_{\mathcal{I}_t}^{(h)\top}}{\sqrt{d_h}}\right) \mathbf{V}_{\mathcal{I}_t}^{(h)}\] where \(\mathbf{K}_{\mathcal{I}_t}, \mathbf{V}_{\mathcal{I}_t}\) gather the full-precision key/value tensors of the selected blocks. Selection is performed independently per GQA group, so different query groups within the same layer may attend to different block subsets.

Algorithmic Pseudocode

Key Characteristics

• Complexity: Index branch \(\mathcal{O}(H_{\text{kv}} \cdot d_{\text{idx}} \cdot N^2)\); Main branch \(\mathcal{O}(H_q \cdot d_h \cdot N \cdot k \cdot B_k)\). Per-query budget fixed at \(k \cdot B_k = 2048\) tokens regardless of \(N\).

• Trainable: Yes; end-to-end (the Index Branch is trained via a KL alignment loss, not eval-only).

• State: \(\mathcal{O}(N)\) KV cache (GQA-sized); no expansion of the recurrent/state dimension.

• Strength: Deployed on the 109B MiniMax-M3, MSA achieves a \(28.4\times\) per-token attention compute reduction at 1M context, with \(14.2\times\) prefill and \(7.6\times\) decode speedup on H800 GPUs.

• Limitation: Selection granularity is fixed at the block level (\(B_k = 128\)); the forced local block and per-group independent selection increase engineering complexity and require custom block-sparse kernels.

SparDA (Sparse Decoupled Attention): Forecast-Guided Lookahead Block Selection

Mathematical Formulation

SparDA introduces a decoupled sparse attention scheme that augments the standard per-layer projections \(Q, K, V\) with a fourth projection, the Forecast, which predicts the KV blocks that will be needed by the next layer. This lookahead enables overlapping CPU-to-GPU prefetch of the next layer’s selected blocks with the current layer’s execution, hiding the offload latency that plagues sparse-attention deployments on memory-constrained hardware.

Forecast projection

: For input \(\mathbf{x} \in \mathbb{R}^{N \times d}\), \[\mathbf{F} = \mathbf{x} \mathbf{W}_F \in \mathbb{R}^{N \times H_{\text{kv}} \times d_f}\] with one Forecast head per GQA group, which reduces selection overhead relative to multi-head selectors used by prior work.

Block scoring and selection

: For each query \(i\) and KV block \(b\) with representative \(\mathbf{r}_b\): \[S_{i,b} = \mathbf{F}_i \cdot \mathbf{r}_b, \qquad \mathcal{I}_i = \mathrm{TopK}(\{S_{i,b}\}_b, k)\]

Block-sparse attention

: Standard scaled dot-product attention is then computed over the selected blocks: \[\mathbf{o}_i = \mathrm{softmax}\left(\frac{\mathbf{q}_i \mathbf{K}_{\mathcal{I}_i}^\top}{\sqrt{d_h}}\right) \mathbf{V}_{\mathcal{I}_i}\]

Training

: SparDA adds fewer than \(0.5\%\) additional parameters. Only the Forecast projections are trained, by matching the original (frozen) sparse selector’s attention distribution—a lightweight distillation that preserves the base model’s behaviour while enabling lookahead prefetch.

Algorithmic Pseudocode

Key Characteristics

• Complexity: Forecast \(\mathcal{O}(H_{\text{kv}} \cdot d_f \cdot N \cdot N_{\text{blocks}})\); Main attention \(\mathcal{O}(H_q \cdot d_h \cdot N \cdot k \cdot B_k)\).

• Trainable: Yes; only the Forecast projections are trained (\(<0.5\%\) extra parameters).

• Strength: On two 8B sparse-pretrained models, SparDA achieves up to \(1.25\times\) prefill and \(1.7\times\) decode speedup over the sparse-attention offload baseline, and up to \(5.3\times\) decode throughput versus a non-offload sparse baseline.

• Limitation: Requires a CPU–GPU offload pipeline and a frozen base selector; the lookahead benefit is contingent on next-layer block demand being predictable from the current Forecast.

LSA (Latent Sparse Attention): Semantic Key Clustering with Top-k Selection

Mathematical Formulation

Latent Sparse Attention replaces the per-key (or per-block) scan of a sparse indexer with semantic clustering: keys are clustered by content in a learned latent space, and each query scores the bucket centroids instead of every key. The paper defines two trainable variants: Distilled-LSA (D-LSA), which adds small index projections to a frozen pre-trained model, and Native-LSA (N-LSA), trained from scratch on the next-token loss alone. The implementation in this codebase follows N-LSA.

Semantic clustering

: Each index key is first L2-normalised, \[\hat{\mathbf{k}}_j = \frac{\mathbf{k}_j}{\lVert \mathbf{k}_j \rVert_2},\] and every \(n_{\text{pool}}\) consecutive keys are summed into one pool key, so clustering runs over \(T/n_{\text{pool}}\) pool keys rather than \(T\) keys. Keys are processed in cluster blocks of \(n_{\text{cluster}}\) consecutive tokens; once a block is complete, its pool keys are merged bottom-up with Ward’s criterion into \(m = n_{\text{cluster}}/\bar{n}_b\) buckets whose unit-norm centroids are \[\mathbf{c}_b = \frac{\mathbf{k}^{\text{sum}}_b}{\lVert \mathbf{k}^{\text{sum}}_b \rVert_2},\] where \(\mathbf{k}^{\text{sum}}_b\) is the sum of the keys in bucket \(b\) and \(\bar{n}_b\) is the average number of keys per bucket. In every round all pairs that are each other’s cheapest merge are merged in parallel; with the Gram matrix \(G_{bb'} = \mathbf{k}^{\text{sum}}_b \cdot \mathbf{k}^{\text{sum}}_{b'}\) the merge cost is \[a_b = \sqrt{G_{bb}}, \qquad a_{b \cup b'} = \sqrt{G_{bb} + G_{b'b'} + 2G_{bb'}}, \qquad \Delta(b,b') = a_b + a_{b'} - a_{b \cup b'},\] and a merge is skipped when the combined bucket would exceed \(n_{\max}^b\) keys. The bucket cache stores every centroid together with the positions of its keys. Once a cluster block is clustered, its buckets never change, which is what makes incremental decoding exact in this implementation.

Semantic search

: Each query is L2-normalised, \(\hat{\mathbf{q}}_i = \mathbf{q}_i/\lVert \mathbf{q}_i \rVert_2\), and scores the centroids of every bucket in the cluster blocks before its own. For N-LSA the score is the best similarity any query head gives the centroid, \[s_{i,b} = \max_h \hat{\mathbf{q}}_{i,h} \cdot \mathbf{c}_b,\] so a single index list \(\mathcal{J}_i\) is shared by every query head. The top-\(p\) buckets are expanded best-first until \(k\) long-range keys are selected; the expansion always fills \(k\) because \(p\) is chosen such that \(p\, n_{\text{pool}} \ge k\), and keys already covered by the sink prefix or the window are skipped. Every query additionally attends to \(n_{\text{sink}}\) sink keys and the \(n_w = n_{\text{cluster}}\) most recent keys (its own, not-yet-clustered block), giving a set \[\mathcal{J}_i \subseteq \{\text{sinks}\} \cup \{\text{top-}k \text{ long-range keys}\} \cup \{\text{own-block window}\}, \qquad |\mathcal{J}_i| \le n_{\text{sink}} + k + n_w .\]

Attention

: Standard scaled dot-product attention is computed only over \(\mathcal{J}_i\), \[z_{i,j} = \frac{\mathbf{q}_i \cdot \mathbf{k}_j}{\sqrt{d_h}}, \qquad \mathbf{o}_i = \mathrm{softmax}_{j \in \mathcal{J}_i}(z_{i,j})\, \mathbf{V}_{\mathcal{J}_i},\] where the softmax is renormalised per head over the shared selection. N-LSA adds no parameters, uses no positional encoding (RoPE could be supported by clustering only the RoPE-free part of the keys), and stops the gradient through the selection: the clustering, centroid scoring and expansion run without gradients, so the next-token loss trains the attention projections only through the keys each query attends to.

Algorithmic Pseudocode

Key Characteristics

• Complexity: Flat search scores \(\mathcal{O}(T/\bar{n}_b)\) centroids per query, i.e. \(\mathcal{O}(T^2/\bar{n}_b + T(k+n_w+n_{\text{sink}})d_h)\) for prefill and \(\mathcal{O}(T/\bar{n}_b + k+n_w+n_{\text{sink}})\) per decoded token; the tree variant of the paper makes the indexer \(\mathcal{O}(T \log T)\).

• Trainable: Yes, natively from scratch on the next-token loss (N-LSA); the selection is non-differentiable and the gradient is stopped through it. No teacher, no index projections, no auxiliary loss—and no extra parameters.

• Strength: In the paper’s 0.5B hybrid, N-LSA matches the full-attention validation loss at 128k tokens (\(1.085\) against \(1.085\); PISA \(1.100\), NSA \(1.117\), SWA \(1.154\)) while training up to \(1.6\times\) faster per attention layer at 128k. Because every cluster block is immutable once complete, cached decoding in this codebase is exact rather than an approximation.

• Limitation: The flat indexer remains quadratic in context length (with a \(1/\bar{n}_b\) constant); tree search is left to future work. The paper’s single index list assumes a single shared KV head (\(n_{kv}=1\)), and N-LSA deliberately uses no positional encoding.

ProbSparse Attention: Query-Sparsity Top-\(u\) Selection

Mathematical Formulation

ProbSparse self-attention , adapted from Informer , attacks the quadratic cost from the query side instead of the key side. For a sequence of length \(L\) it keeps a same-shape tensor \(\bar{\mathbf{Q}}\) containing only the top-\(u\) dominant queries under a sparsity measure \(\mathcal{M}(\mathbf{q}, \mathbf{K})\), so that every key is attended to by only those queries: \[A(\mathbf{Q},\mathbf{K},\mathbf{V}) = \sum_j \left[\frac{k(\mathbf{Q},\mathbf{k}_j)}{\sum_l k(\mathbf{Q},\mathbf{k}_l)}\right]\mathbf{v}_j = \mathbb{E}_{p(\mathbf{k}_j\mid\mathbf{Q})}[\mathbf{v}_j] \approx \mathrm{Softmax}\!\left(\frac{\bar{\mathbf{Q}}\mathbf{K}^{\top}}{\sqrt{d}}\right)\mathbf{V},\] with the exponential kernel \(k(\mathbf{q}_i,\mathbf{k}_j)=\exp(\mathbf{q}_i\mathbf{k}_j^{\top}/\sqrt{d})\). The measure follows from the Kullback–Leibler divergence between the query’s attention distribution and the uniform distribution: a query is “dominant” when its distribution is far from uniform, and the max–mean approximation \[\bar{\mathcal{M}}(\mathbf{q}_i,\mathbf{K}) = \max_j\left\{\frac{\mathbf{q}_i\mathbf{k}_j^{\top}}{\sqrt{d}}\right\} - \frac{1}{L}\sum_{j=1}^{L}\frac{\mathbf{q}_i\mathbf{k}_j^{\top}}{\sqrt{d}}\] ranks the queries. Computing it over all keys would itself be quadratic, so the paper samples only \(L'_K = c_1\lceil\ln L_K\rceil\) keys for the measure and keeps \(L'_Q = c_2\lceil\ln L_Q\rceil\) queries, with \(c_1=c_2=5.0\); this is what yields \(\mathcal{O}(L\log L)\) time and memory. Multi-head attention is the usual \[\mathrm{MHSA\text{-}Sparse}(\mathbf{X}) = \mathrm{Concat}(\mathrm{head}_1,\dots,\mathrm{head}_h)\mathbf{W}_O, \qquad \mathrm{head}_i = A(\mathbf{X}\mathbf{W}_Q^i,\mathbf{X}\mathbf{W}_K^i,\mathbf{X}\mathbf{W}_V^i).\]

Relative position terms

: The paper follows Transformer-XL and scores the selected queries with two per-head bias vectors \(\mathbf{U}_1,\mathbf{U}_2\) and a projected relative-position tensor \(\mathbf{P}=\mathbf{X}_P\mathbf{W}_P\), \[\mathbf{S} = \left(\bar{\mathbf{Q}}_{U_1}\mathbf{K}^{\top} + \bar{\mathbf{Q}}_{U_2}\mathbf{P}^{\top}\right)/\sqrt{d_q}, \qquad \mathbf{Q}_{U_1} = \mathbf{Q}+\mathbf{U}_1,\; \mathbf{Q}_{U_2} = \mathbf{Q}+\mathbf{U}_2,\] where \(\mathbf{X}_P\in\mathbb{R}^{b\times L_K\times d}\) encodes the keys’ positions and \(W_P\in\mathbb{R}^{d\times d}\) is trained. In this codebase the encoding is sinusoidal (as in Transformer-XL) and only the projection is learned, so the layer costs \(5d^2+2d\) attention parameters per layer: four projections, \(W_P\), and the two vectors \(\mathbf{U}_1,\mathbf{U}_2\) (shape \(h\times d_q\) each). The remaining rows, i.e. the non-dominant queries, keep the head-mean of \(\mathbf{V}\), which is what Algorithm 1 of the paper does when it overwrites only the selected rows.

Decoder mode

: The selection is a block-global decision, and filling discarded rows with a bidirectional statistic would leak future tokens. Decoder mode therefore scores every query causally with the same \(\mathbf{U}_1\), \(\mathbf{U}_2\) and \(W_P\) terms, and caches K/V like standard attention; because sparsity is query-side, single-query decoding makes the top-\(u\) rule degenerate, so attending the full cached prefix is exact.

Algorithmic Pseudocode

Key Characteristics

• Complexity: \(\mathcal{O}(L_Q L_K d)\) for the max–mean probe over sampled keys plus \(\mathcal{O}(L'_Q L_K d)\) for the selected rows, i.e. \(\mathcal{O}(L\log L)\) time and memory; \(\mathcal{O}(n_{\text{ctx}})\) per decoded token over the KV cache.

• Trainable: Yes, natively from scratch; the CSJ models are trained end-to-end with the sparse path active. The selection itself is non-differentiable, so gradients reach the projections through the selected rows and through the value mean of the discarded ones.

• Strength: On CSJ-500h, the 12-layer deep sparse Conformer reaches CERs of 5.52/4.03/4.50 against 6.97/4.65/5.29 for Conformer-Large at the same 135M parameters, trains in 80% of the time and decodes \(1.23\times\) faster; the same blocks scale to 100 layers with DeepNorm.

• Limitation: The sparsity is query-side, so it does not reduce the key-side cost of a single query—decoding a token still attends the whole cache. The paper defines the mechanism for bidirectional encoders only; the paper’s second ingredient, DeepNorm (\(\alpha_{1..4}=0.81(N^4M)^{1/16}\)), is a residual-scaling strategy that lives in the residual family of this codebase, not in this mixer.

Design Space and Selection Criteria

The twelve sparse attention methods occupy complementary positions in a multidimensional design space:

• Geometric/fixed patterns (Sparse Transformer, Longformer): Simple to implement and analyze, but patterns do not adapt to content.

• Fixed-memory KV eviction (SWA with sinks): The cheapest possible policy—keep a sink prefix and a recent window, discard everything else. No extra parameters, no training required, bounded memory.

• Learned selection (SparseK, NSA): Enable adaptation through trainable selection networks, at the cost of higher implementation complexity.

• Training-free acceleration (FASA, SpargeAttn): Enable drop-in speedup of existing models without retraining, suited for deployed systems.

• Semantic selection (LSA): Clusters keys by content and caches one centroid per bucket, so the indexer scans a small multiple of \(T/\bar{n}_b\) instead of \(T\); the only member whose per-block statistics are immutable, which makes exact cached decoding possible.

• Theoretical guarantees (BigBird): Provide formal proofs of expressive completeness, important for understanding safety margins.

• Query-side sparsity (ProbSparse): Selects which queries deserve a full score row instead of which keys a query sees, so cost falls to \(\mathcal{O}(L\log L)\) without changing the attended set of any kept query; the natural fit for long bidirectional encoders rather than for decoding.

Selection guidance:

1. New model training where resources permit: NSA or SparseK for maximal inference speedup; BigBird for theoretical assurance; LSA when semantic clustering and exact cached decoding at ultra-long context matter; SWA with sinks as the strongest low-cost baseline.

2. Long-context document understanding: SWA with sinks (window 256–512) for the best accuracy per byte of KV cache; Longformer or FASA for practical simplicity; NSA for extreme scale; LSA when the indexer scan itself dominates.

3. Accelerating existing models: SWA with sinks as the default training-free baseline (beat most linear-attention retrofits in the original study); FASA for RoPE-based LLMs; SpargeAttn for any architecture.

4. Constrained environments: SWA with sinks for a fixed memory budget; Sparse Transformer for simplicity and proven efficiency at moderate scale.

Gated Attention Families—Complete Literature Analysis

Executive Summary: Gating for Memory Control

The emerging field of gated attention mechanisms addresses a fundamental challenge in sequence modeling: how should information be retained, forgotten, and updated as the model processes new tokens? Traditional additive recurrent states accumulate all information without erasure, leading to memory saturation and inability to adapt to changing context. Softmax attention provides full expressiveness but uses \(\mathcal{O}(Ld)\) KV cache during inference, limiting context window size. Gated mechanisms provide a middle ground: selective control over information survival.

The unifying principle across gated architectures is that gating controls what information survives. Gates operate at three distinct levels:

1. Recurrent state decay (GLA, HGRN2): Multiplicative gates on the memory matrix \(S_t\) control how much previous state persists into the next step.

2. Write strength in recurrent updates (DeltaNet, Gated DeltaNet, Gated DeltaNet-2): Gates control the magnitude and channel-wise direction of new information written into memory.

3. Softmax logit biasing (FoX): Gates inject token-level recency bias into attention logit computation.

4. Post-attention sparsity (Gated Softmax): Gates selectively scale SDPA output channels.

This appendix synthesizes eight major gated-attention architectures published 2023–2026, analyzing their mathematical foundations, hardware efficiency characteristics, and empirical trade-offs.

1. Gated Linear Attention (GLA)

Mathematical Foundation

Gated Linear Attention augments vanilla linear attention (which uses a matrix-valued recurrent state) with data-dependent multiplicative decay. Standard linear attention reformulates the traditional softmax mechanism as a recurrence over hidden states: \[\mathbf{S}_t = \mathbf{S}_{t-1} + \mathbf{v}_t \mathbf{k}_t^\top, \quad \mathbf{o}_t = \mathbf{S}_t \mathbf{q}_t\] where state \(\mathbf{S}_t \in \mathbb{R}^{d_v \times d_k}\) accumulates outer products. This purely additive formulation suffers from “memory overload”: information accumulates monotonically and cannot be erased, degrading performance on tasks requiring context selection.

GLA introduces a data-dependent diagonal gating matrix \(\mathbf{G}_t \in [0,1]^{d_k \times d_k}\): \[\mathbf{S}_t = \mathbf{G}_t \odot \mathbf{S}_{t-1} + \mathbf{v}_t \mathbf{k}_t^\top, \quad \mathbf{o}_t = \mathbf{S}_t \mathbf{q}_t\] The gate \(\mathbf{G}_t\) is parameterized with an outer-product structure: \[\mathbf{G}_t = \mathbf{1} \boldsymbol{\alpha}_t^\top, \quad \boldsymbol{\alpha}_t = \text{logsigmoid}\left(\mathbf{W}_{gk} \mathbf{x}_t\right) / c\] where \(\mathbf{W}_{gk}\) is a low-rank projection (\(\mathbb{R}^{d \to d_k}\)) and \(c \approx 16\) is a normalizer. The logsigmoid activation maps \(\alpha_t \in \mathbb{R}\) to approximately \([-1/2, 0]\), and division by \(c\) further contracts this to near-identity for initialization. The resulting gate \(G_t = 1 + \alpha_t \approx 1\) near initialization, then learns to decay \((\alpha_t \to -\infty)\) historical information.

Hardware-Efficient Training

For parallel training, GLA uses a chunkwise block-wise algorithm that groups tokens into chunks \(C\) and computes recurrent transitions in parallel: \[\mathbf{O}_{[t]} = \overleftarrow{\mathbf{Q}}_{[t]} \mathbf{S}_{[t]}^\top + \left(\mathbf{Q}_{[t]} \mathbf{K}_{[t]}^\top \odot \mathbf{\Gamma}_{[t]}\right) \mathbf{V}_{[t]}\] where \(\overleftarrow{\mathbf{Q}}_{[t]}\) are attending to the state at the chunk boundary (lookback), and \(\mathbf{\Gamma}_{[t]}\) is the causal mask with decay-aware scaling: \[\mathbf{\Gamma}_{[t],ij} = \begin{cases} \prod_{l=j}^{i-1} \alpha_l & \text{if } i > j \text{ (lookback)} \\ \prod_{l=1}^{i-j} \alpha_l & \text{if } i \le j \text{ (in-chunk)} \end{cases}\] This design maximizes matmul operations susceptible to tensor-core acceleration, achieving sub-quadratic total complexity \(\mathcal{O}(nLd^2/C)\) where \(L\) is the number of attention heads and \(C\) is chunk size.

Strengths and Limitations

2. DeltaNet: Error-Correcting Linear Attention

Mathematical Foundation

DeltaNet applies a classical delta learning rule to linear attention state updates. Instead of additive accumulation, DeltaNet computes the difference between the predicted value (retrieved from memory) and the target value, then performs an error-correcting update: \[\mathbf{S}_t = \mathbf{S}_{t-1}(\mathbf{I} - \beta_t \mathbf{k}_t \mathbf{k}_t^\top) + \beta_t \mathbf{v}_t \mathbf{k}_t^\top\] where \(\beta_t \in (0,1)\) is a data-dependent writing strength scalar. This update can be decomposed as an erase-write operation: \[\mathbf{v}_t^\text{old} = \mathbf{S}_{t-1} \mathbf{k}_t \quad \text{(current prediction)}, \quad \mathbf{v}_t^\text{updated} = \beta_t \mathbf{v}_t + (1-\beta_t)\mathbf{v}_t^\text{old} \quad \text{(blend old/new)}\] so that: \[\mathbf{S}_t = \mathbf{S}_{t-1} - \mathbf{v}_t^\text{old}\mathbf{k}_t^\top + \mathbf{v}_t^\text{updated}\mathbf{k}_t^\top\] The key insight is that this update rule minimizes an online MSE loss at each timestep: \[\mathcal{L}_t(\mathbf{S}) = \tfrac{1}{2}\|\mathbf{S}\mathbf{k}_t - \mathbf{v}_t\|^2 \Rightarrow \nabla_{\mathbf{S}} \mathcal{L}_t = (\mathbf{S}\mathbf{k}_t - \mathbf{v}_t)\mathbf{k}_t^\top\] One step of SGD with learning rate \(\beta_t\) yields the delta rule update. This connection to test-time training (TTT) provides theoretical grounding: DeltaNet is directly optimizing for value prediction at inference time.

Efficient Parallel Training

DeltaNet’s transition matrix \(\mathbf{I} - \beta_t \mathbf{k}_t \mathbf{k}_t^\top\) is a generalized Householder matrix (rank-1 update to identity). For chunkwise training, DeltaNet uses the WY representation, which represents the product of \(m\) such rank-1 updates compactly: \[\prod_{i=1}^{m} (\mathbf{I} - \beta_i \mathbf{k}_i \mathbf{k}_i^\top) = \mathbf{I} - \mathbf{W}\mathbf{Y}^\top\] where \(\mathbf{W}, \mathbf{Y} \in \mathbb{R}^{d \times m}\) are constructed recursively. This factorization enables matmul-rich parallelism without materializing the full product, keeping chunkwise training efficient.

Strengths and Limitations

3. Gated DeltaNet: Synthesis of Gating and Error Correction

Mathematical Foundation

Gated DeltaNet synthesizes the strengths of GLA and DeltaNet by combining a decay gate \(\alpha_t\) (from GLA) with the writing strength gate \(\beta_t\) (from DeltaNet): \[\mathbf{S}_t = \mathbf{S}_{t-1}\left(\alpha_t (\mathbf{I} - \beta_t \mathbf{k}_t \mathbf{k}_t^\top)\right) + \beta_t \mathbf{v}_t \mathbf{k}_t^\top\] Rewriting for clarity: \[\mathbf{S}_t = \alpha_t \mathbf{S}_{t-1}(\mathbf{I} - \beta_t \mathbf{k}_t \mathbf{k}_t^\top) + \beta_t \mathbf{v}_t \mathbf{k}_t^\top\] The two gates are complementary:

• \(\alpha_t \to 0\) rapidly erases historical state: \(\mathbf{S}_t \approx \beta_t \mathbf{v}_t \mathbf{k}_t^\top\) (context reset).

• \(\alpha_t \to 1\) recovers pure delta-rule behavior: \(\mathbf{S}_t \approx \mathbf{S}_{t-1}(\mathbf{I} - \beta_t \mathbf{k}_t \mathbf{k}_t^\top) + \beta_t \mathbf{v}_t \mathbf{k}_t^\top\) (targeted updates).

• \(\beta_t \to 0\) skips writing but decays: \(\mathbf{S}_t \approx \alpha_t \mathbf{S}_{t-1}\) (pure forgetting).

• \(\beta_t \to 1\) replaces memory aggressively: \(\mathbf{S}_t \approx \alpha_t \mathbf{S}_{t-1} + \mathbf{v}_t \mathbf{k}_t^\top\) (replacement).

From an online learning perspective, Gated DeltaNet corresponds to: \[\min_{\mathbf{S}_t} \|\mathbf{S}_t - \alpha_t \mathbf{S}_{t-1}\|_F^2 - 2\langle \mathbf{S}_t \mathbf{k}_t, \beta_t(\mathbf{v}_t - \alpha_t \mathbf{S}_{t-1} \mathbf{k}_t)\rangle\] which introduces an adaptive weight decay \(\alpha_t\) into an SGD-like update—analogous to decoupled weight decay in deep learning optimization.

Empirical Performance and Trade-offs

Gated DeltaNet has been integrated into Alibaba’s Qwen3-Next and Qwen3.5 production models. Empirically: \[\begin{array}{l|cccc} \textbf{Task} & \textbf{Gated DeltaNet} & \textbf{Mamba2} & \textbf{DeltaNet} & \textbf{GLA} \\ \hline \text{Language Modeling} & \mathbf{1.0\times} & 1.08\times & 1.05\times & 1.12\times \\ \text{Associative Recall} & \mathbf{1.0\times} & --(no perfect) & \mathbf{1.0\times} & 0.75 \\ \text{Length Extrapolation} & \mathbf{1.0\times} & 0.98\times & 0.96\times & 0.94\times \\ \text{Throughput (tokens/sec)} & 0.85\times & \mathbf{1.0\times} & 0.82\times & 0.88\times \end{array}\] Gated DeltaNet achieves the best balance across diverse tasks at the cost of slightly reduced throughput due to richer transition matrices.

4. HGRN2: Hierarchical Gating with Outer-Product Expansion

Mathematical Foundation

HGRN2 uses an outer-product-based state expansion with hierarchically lower-bounded forget gates. The state update is: \[\mathbf{S}_t = \text{diag}(\mathbf{g}_t) \cdot \mathbf{S}_{t-1} + \mathbf{v}_t \mathbf{k}_t^\top, \quad \mathbf{o}_t = \mathbf{S}_t \mathbf{q}_t\] where \(\mathbf{g}_t \in [b_\ell, 1]^d\) is a lower-bounded forget gate for layer \(\ell\). The bounds satisfy: \[0 \le b_{\ell_{\text{shallow}}} < b_{\ell_{\text{middle}}} < b_{\ell_{\text{deep}}} \le 1\] i.e., bounds increase (become less restrictive) in the deeper layers. The gate itself is computed as: \[\mathbf{g}_t = b_\ell + (1 - b_\ell) \cdot \sigma(W_g \mathbf{x}_t)\] where \(\sigma\) is sigmoid. When \(W_g\) outputs weakly negative logits, \(\mathbf{g}_t \approx b_\ell\) (forced retention at shallow layers). When outputs are strongly positive, \(\mathbf{g}_t \approx 1\) (memory refresh at any layer).

The hierarchical structure encourages different layers to specialize to different timescales: shallow layers model local dependencies (high retention due to binding \(b_\ell\)), while deeper layers capture long-range structure (flexible gating with high upper bound).

Scaling and Empirical Results

At 3B scale on 100B tokens, HGRN2 slightly outperforms Mamba2 and LLaMA-architecture Transformers on language modeling. The hierarchical gating provides multi-scale temporal modeling without explicit layer-wise architectural variants. However, vector-valued diagonal gating (as opposed to element-wise selection) provides less per-item flexibility than delta-rule variants.

5. Forgetting Transformer (FoX): Gating in Softmax Logit Space

Mathematical Foundation

Forgetting Transformer (FoX), proposed by Lin et al. , embeds a forget gate directly into softmax attention logits. Rather than replacing softmax with linear recurrence, FoX preserves full softmax expressiveness while adding recency control through a data-dependent logit bias.

For each token position \(t\) and all keys \(j \le t\), compute a scalar forget gate: \[f_t = \sigma(w_f^\top \mathbf{x}_t + b_f)\] The logit bias at position \((i,j)\) is the cumulative log-forget product: \[d_{ij} = \sum_{l=j+1}^{i} \log f_l = \log \prod_{l=j+1}^{i} f_l\] The full attention output becomes: \[\mathbf{O} = \text{softmax}\left(\frac{\mathbf{Q}\mathbf{K}^\top}{\sqrt{d_k}} + \mathbf{D}\right) \mathbf{V}\] where \(\mathbf{D}_{ij} = d_{ij}\). Equivalently: \[\text{Attention}(i,j) = \frac{\exp(\mathbf{q}_i^\top \mathbf{k}_j / \sqrt{d_k} + d_{ij})}{\sum_{j'=1}^{i} \exp(\mathbf{q}_i^\top \mathbf{k}_{j'} / \sqrt{d_k} + d_{ij'})}\] This weighting down-weights older tokens multiplicatively: a token from 10 steps ago is scaled by \(\prod_{l=j+1}^{i} f_l\), which is typically much less than 1 if \(f_t < 1\) on average.

The mechanism is mathematically equivalent to a data-dependent learnable variant of ALiBi (Attention with Linear Biases), where earlier work used fixed \(d_{ij} = -\infty \cdot |i-j|\) or \(d_{ij} = -(i-j)\).

Integration with FlashAttention

The key advantage of FoX is compatibility with FlashAttention and existing optimized implementations. The forget biases are added in the softmax logit computation, which is already a core operation in FlashAttention’s block-wise algorithm. The overhead is: \[\text{Overhead} \approx 0.5--2\% \text{ wall-clock time}\] since the bias addition is amortized across matmul operations.

Strengths and Limitations

FoX achieves state-of-the-art results on length extrapolation, near-perfect needle-in-haystack retrieval, and superior long-context understanding compared to Transformers. However, it remains \(\mathcal{O}(L^2 d)\) at training and \(\mathcal{O}(L d)\) per step at inference with KV cache, limiting extreme sequence lengths.

6. Gated Attention (Post-SDPA Sigmoid Gating)

Mathematical Foundation

The NeurIPS 2025 Best Paper by the Qwen team proposes applying a sigmoid gate after scaled dot-product attention (SDPA): \[\mathbf{Y} = \text{SDPA}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \text{softmax}\left(\frac{\mathbf{Q}\mathbf{K}^\top}{\sqrt{d_k}}\right) \mathbf{V}\] \[\mathbf{Y}' = \mathbf{Y} \odot \sigma\left(\mathbf{X} \mathbf{W}_\theta\right)\] where the gate score \(\sigma(\mathbf{X} \mathbf{W}_\theta)\) can have shape \(\mathbb{R}^{n \times q \times d_k}\) (element-wise gating per query head) or \(\mathbb{R}^{n \times q}\) (head-wise fixed gating). Across 30+ variants and 15B MoE + 1.7B dense models trained on 3.5T tokens, the team found:

• Headwise gating adds only \(\sim 1.6M\) parameters to a 15B model (negligible overhead).

• Effective gates are sparse: mean activation \(\approx 0.116\) (i.e., 88% are zero or near-zero).

• Gates act as query-dependent filters, suppressing low-value channels while preserving high-value signal.

• Gates demonstrably eliminate the attention sink phenomenon, where the attention pattern becomes dominated by a single early token.

Key Findings

• Attention-sink suppression: Gating breaks the feedback loop where softmax concentrates on the first token, enabling more balanced attention.

• Training stability: Models with gating tolerate larger learning rates (\(+30\%\) higher stable \(\eta_{max}\)).

• Minimal overhead: Latency impact is only \(1.6\%\) on H100 GPUs; throughput drops at most \(2-3\%\).

• Scaling law improvement: Improves loss across model scales from 800M to 110B parameters consistently.

7. Gated DeltaNet-2: Decoupled Channel-wise Erase and Write

Mathematical Foundation

Gated DeltaNet-2 addresses a shared limitation of Gated DeltaNet and Kimi Delta Attention (KDA): both use a single scalar delta gate \(\beta_t\) to simultaneously control (i) how much old content is erased along the key-side read direction and (ii) how much new value is written along the value axis. Gated DeltaNet-2 decouples these two roles into a channel-wise erase gate \(\mathbf{b}_t\) (key axis) and a channel-wise write gate \(\mathbf{w}_t\) (value axis), while inheriting the channel-wise decay \(\boldsymbol{\alpha}_t\) from KDA. With state \(\mathbf{S}_t \in \mathbb{R}^{d_k \times d_v}\), decay \(\mathbf{D}_t = \mathrm{Diag}(\boldsymbol{\alpha}_t)\), erase direction \(\mathbf{e}_t = \mathbf{b}_t \odot \mathbf{k}_t\), and write target \(\mathbf{z}_t = \mathbf{w}_t \odot \mathbf{v}_t\), the state update is: \[\mathbf{S}_t = \left(\mathbf{I} - \mathbf{k}_t \mathbf{e}_t^\top\right) \mathbf{D}_t \, \mathbf{S}_{t-1} + \mathbf{k}_t \mathbf{z}_t^\top, \qquad \mathbf{o}_t = \mathbf{S}_t^\top \mathbf{q}_t\] or, expanding the gate products: \[\mathbf{S}_t = \left(\mathbf{I} - \mathbf{k}_t (\mathbf{b}_t \odot \mathbf{k}_t)^\top\right) \mathbf{D}_t \, \mathbf{S}_{t-1} + \mathbf{k}_t (\mathbf{w}_t \odot \mathbf{v}_t)^\top\] The two gates are produced by independent linear projections followed by sigmoid: \[\mathbf{b}_t = \sigma(\mathbf{W}_b \mathbf{x}_t), \qquad \mathbf{w}_t = \sigma(\mathbf{W}_w \mathbf{x}_t)\] and the channel-wise decay follows a log-decay parameterization (computed in fp32 to avoid precision loss in long cumulative products): \[\mathbf{g}_t = \mathbf{a} - \mathrm{softplus}(\boldsymbol{\delta}), \qquad \boldsymbol{\alpha}_t = \exp(\mathbf{g}_t)\] The structural roles of the two gates are complementary:

• \(\mathbf{b}_t\) (key axis) makes the read direction channel-selective: it determines which key channels are removed before writing.

• \(\mathbf{w}_t\) (value axis) makes the write target channel-selective: it determines which value channels are deposited into memory.

• \(\boldsymbol{\alpha}_t\) (per key channel) provides global decay, as in KDA.

Reductions. Setting \(\mathbf{b}_t = \beta_t \mathbf{1}_{d_k}\) and \(\mathbf{w}_t = \beta_t \mathbf{1}_{d_v}\) recovers KDA; further collapsing \(\boldsymbol{\alpha}_t\) to a scalar recovers Gated DeltaNet. From the fast-weight / online-learning perspective, the update is the minimizer of: \[\mathcal{L}_t(\mathbf{S}) = \tfrac{1}{2}\|\mathbf{S} - \mathbf{D}_t \mathbf{S}_{t-1}\|_F^2 - 2\left\langle \mathbf{S}^\top \mathbf{k}_t,\; \mathbf{z}_t - (\mathbf{D}_t \mathbf{S}_{t-1})^\top \mathbf{e}_t \right\rangle\]

The chunkwise training algorithm extends KDA’s WY representation by absorbing the channel-wise decay into an asymmetric delta recurrence on a decay-normalized state, enabling matmul-rich parallelism on tensor cores. Because \(\mathbf{b}_t\) and \(\mathbf{w}_t\) are per-channel diagonal operators varying per row, the scalar post-scaling shortcut of KDA breaks: the backward pass requires a gate-aware WY inversion that bakes the gates into each dot product accumulating the gradient.

Empirical Performance and Trade-offs

At 1.3B parameters trained on 100B FineWeb-Edu tokens, Gated DeltaNet-2 outperforms Mamba-2, Mamba-3, Gated DeltaNet, and KDA on language modeling, commonsense reasoning, and—most markedly—real-world and RULER long-context retrieval (especially multi-key needle-in-a-haystack). Its throughput on H100 stays near-flat (\(38.0 \to 36.1\) Kt/s) as sequence length grows from 2K to 16K, where a Transformer’s drops sharply. Ablations show the erase gate \(\mathbf{b}_t\) accounts for most of the gain; the write gate \(\mathbf{w}_t\) contributes a smaller additional improvement. The hybrid variant (recurrent Gated DeltaNet-2 + sliding-window attention) further closes the gap with pure-attention models on tasks requiring local evidence aggregation.

8. Kimi Delta Attention (KDA): Channel-wise Decay Gating for the Delta Rule

Mathematical Foundation

Kimi Delta Attention is the core module of Kimi Linear and extends Gated DeltaNet with a finer-grained channel-wise decay gate. Where Gated DeltaNet uses a single scalar per-head decay \(\alpha_t\) alongside the scalar delta write gate \(\beta_t\), KDA replaces the scalar decay with a per-key-channel decay \(\boldsymbol{\alpha}_t \in \mathbb{R}^{d_k}\), parameterised through a log-decay formulation and applied as a diagonal matrix \(\mathbf{D}_t = \mathrm{diag}(\boldsymbol{\alpha}_t)\) before the delta-rule update. The delta rule itself retains a scalar per-head write gate \(\beta_t\): \[\mathbf{S}_t = (\mathbf{I} - \beta_t \mathbf{k}_t \mathbf{k}_t^\top)\, \mathbf{D}_t\, \mathbf{S}_{t-1} + \beta_t \mathbf{v}_t \mathbf{k}_t^\top, \qquad \mathbf{o}_t = \mathbf{S}_t^\top \mathbf{q}_t\] The channel-wise decay follows a log-decay parameterisation (computed in fp32 for numerical stability over long cumulative products): \[\mathbf{g} = \mathbf{a} - \mathrm{softplus}(\boldsymbol{\delta}), \qquad \boldsymbol{\alpha}_t = \exp(\mathbf{g})\] where \(\mathbf{a}\) and \(\boldsymbol{\delta}\) are learned per-channel parameters. This gives each key channel an independent forgetting timescale, allowing the model to retain some channels longer than others within the same head.

Reduction to Gated DeltaNet

: When \(\boldsymbol{\alpha}_t\) collapses to a scalar (i.e., all channels share the same decay), KDA reduces exactly to Gated DeltaNet. The channel-wise decay is thus a strict generalisation that adds per-channel selectivity on top of the scalar decay/write gating of Gated DeltaNet.

Chunkwise training

: A bespoke chunkwise algorithm uses a specialised Diagonal-Plus-Low-Rank (DPLR) variant of the transition matrix to remain hardware-efficient. The product of per-step transitions \((\mathbf{I} - \beta_t \mathbf{k}_t \mathbf{k}_t^\top)\mathbf{D}_t\) admits a DPLR decomposition that allows matmul-rich parallelism on tensor cores while preserving the channel-wise decay structure.

Empirical Performance and Trade-offs

KDA is deployed in Kimi Linear, a hybrid model (3B activated / 48B total parameters) combining KDA layers with Multi-head Latent Attention (MLA). Kimi Linear outperforms a full-MLA baseline across all evaluated tasks, reduces the KV cache by up to \(75\%\), and achieves up to \(6\times\) decoding throughput for 1M-token context. The channel-wise decay is the principal driver of the gain: it allows fine-grained control over which key channels persist, improving long-context recall relative to scalar-decay Gated DeltaNet.

9. First Passage Attention (FPA): Forgetting by Path Extrema

Mathematical Foundation

First Passage Attention keeps the softmax attention of Vaswani et al.  but replaces FoX’s path sum of forget gates with a path extremum. A causal head of dimension \(d\) with queries \(\mathbf{q}_t\), keys \(\mathbf{k}_s\) and values \(\mathbf{v}_s\) computes \[\mathbf{o}_t=\sum_{s\le t}w_{ts}\mathbf{v}_s,\qquad w_{ts}\propto e^{\ell_{ts}},\qquad \ell_{ts}=\frac{\mathbf{q}_t^\top\mathbf{k}_s}{\sqrt{d}}+b_{ts},\] where the first-passage bias is \[b_{ts}=\beta\min\Bigl(0,\ \min_{s<u\le t}\phi_u-\psi_s\Bigr),\qquad b_{tt}=0,\] with a learned positive strength \(\beta\) per head, a scalar field \(\phi_u\) at every position and a scalar level \(\psi_s\) for every key. We write \(\bar{\phi}_{ts}=\min_{s<u\le t}\phi_u\) for the range minimum. A key is alive at time \(t\) when \(b_{ts}=0\), and dead otherwise; a dead key’s bias is \(-\beta\) times the depth \(\psi_s-\bar{\phi}_{ts}\) of its deepest dip. Death is permanent: \(b_{t's}\le b_{ts}\) for all \(t'<t\) (Proposition 2.3(a)), so the decoder can drop a dead key once its bias passes a threshold.

Each head derives the three scalars per position from the layer input \(\mathbf{x}_t\) through one linear map, with a walk \(W_t\) and the field/level \[W_t=\sum_{u\le t}\delta_u,\qquad \phi_u=W_u+g_u,\qquad \psi_s=W_s+h_s .\] The increments \(\delta\) move the reference of every later comparison, while \(g\) and \(h\) let a token dip or set a level without moving it. A layer of width \(D\) with \(H\) heads therefore gains \(3H(D+1)+H\) parameters — \(27{,}696\) per layer at \(D=768\), \(H=12\), i.e. \(332{,}352\) over twelve layers, and \(+1.56\%\) over keys and values in the cache.

Special cases (Proposition 2.2).

The linear parameterisation contains several families:

• (a) If \(g=h=0\), then \(b_{ts}=-\beta L_{ts}\) with \(L_{ts}=\max\bigl(0,\max_{s<u\le t}(W_s-W_u)\bigr)\), the regulator of a queue that starts empty at time \(s\) (the stack form).

• (b) If in addition \(\delta_u\le 0\) for all \(u\), then \(b_{ts}=\sum_{s<u\le t}\log f_u\) with \(f_u=e^{\beta\delta_u}\in(0,1]\) — exactly FoX with forget gates \(f_u\). A rising walk is what lets FPA keep an old key while removing newer ones, which no forget gate can do.

• (c) If \(\delta=0\), \(h=0\), \(g\le 0\) and \(\beta=1\), then \(b_{ts}=\min_{s<u\le t}g_u\), a bottleneck bias.

• (d) If \(\delta_u=-1\), \(g=0\) and \(\psi_s=W_s-e_s\) with \(e_s\ge 0\), then \(b_{ts}=\beta\min(0,e_s-(t-s))\), the expiry rule of Expire-Span.

In the stack form the levels of the alive keys do not decrease with the time of writing, and a dip kills a suffix in time of the alive keys: the alive set evolves as a stack (Proposition 2.3(b)).

Depth, not duration.

Selective Attention subtracts a sum of per-pair selections, so its suppression grows with the number of intervening steps, confounding the depth of an event with its duration. FPA’s bias is \(-\beta\) times the deepest dip and is unchanged when a dip no deeper than the deepest is held longer or inserted (Proposition 3.4); among aggregations of dips, functions of the maximum are the only ones with this duration invariance (Proposition 3.4(iii)).

Training (straight-through).

The bias is piecewise linear in the field and the levels, and on the closed side of the outer minimum it passes no gradient to the field; in pilot runs that minimum could close for every key of a head early in training. FPA therefore trains with a straight-through estimator: the forward pass takes the hard range minimum and the backward pass takes, in its place, the gradient of the soft minimum \(-\tau\log\sum_{s<u\le t}e^{-\phi_u/\tau}\) with \(\tau=1\). The outer minimum acts on the hard value, and evaluation uses the hard bias unchanged.

Cost and exact decoding.

In a blocked computation with block size \(B\), the bias between a query block and an earlier key block needs only the suffix minima of \(\phi\) within the key block, the minimum over the blocks strictly between, and the prefix minima of \(\phi\) within the query block: \(O(B^2)\) operations per pair of blocks, as an element-wise bias such as FoX’s, plus \(O(T/B)\) block minima (Proposition 2.4(a)). In autoregressive decoding \(\bar{\phi}_{ts}=\min(\bar{\phi}_{t-1,s},\phi_t)\) for \(s<t-1\) and \(\bar{\phi}_{t,t-1}=\phi_t\), so each step updates the bias row in \(O(t)\) and stores two scalars per cached key and head (Proposition 2.4(b)) — an exact plain-KV incremental branch, implemented in this repository as slot extras.

Empirical Performance and Trade-offs

The paper trains four 110M-parameter decoders (12 layers, width 768, 12 heads of dimension 64, SwiGLU-2048, RMSNorm, context 2048, 2.7B SlimPajama tokens) that differ only in attention. FPA (without positional encoding) matches or beats FoX in perplexity and extrapolates further:

The cost is close to standard attention: one FPA layer takes \(0.97\times\) the forward-plus-backward time of PyTorch SDPA and \(1.50\times\) that of FoX, whose bias needs no range minimum; at batch size 1 and length 8192 decoding takes 22.0 ms per token against 20.2 ms for FoX and 23.9 ms for softmax. Because the bias of a key never rises again, a decoder can drop a key once its bias falls below a threshold: at \(\theta=10\) the FPA model loses \(0.00044\) nats at length 2048 and keeps 34% of its keys (26% at 8192), and it still retrieves every passkey at 8192 and 16,384 tokens with 24–27% of its cache.

Unified Gating Principle

All nine architectures share a common principle: gating is a mechanism to control information flow and selective memory retention. The specific design choices diverge along several dimensions:

1. Location of gating: Recurrent state updates (GLA, DeltaNet, Gated DeltaNet, Gated DeltaNet-2, HGRN2) versus softmax logit/output space (FoX, Gated Softmax, FPA, RetNet).

2. Granularity of control: Dimension-wise (GLA, HGRN2 diagonal), key-specific (DeltaNet), token-wise (FoX), head-wise (Gated Softmax), channel-wise decoupled erase/write (Gated DeltaNet-2), or a per-key level compared against a head-specific field (FPA).

3. Data dependence: Fully data-dependent (GLA, DeltaNet, Gated DeltaNet, Gated DeltaNet-2, FoX, Gated Softmax, FPA) versus fixed schedules (RetNet with exponential decay).

4. Aggregation along the key–query path: No path term (Gated Softmax), a sum of log-forget gates (FoX), or a range minimum of a learned field (FPA) — sums accumulate with duration, extrema depend only on the deepest dip.

5. Expressiveness trade-off: Linear recurrent efficiency (GLA, DeltaNet, Gated DeltaNet, Gated DeltaNet-2, HGRN2, RetNet) versus full softmax expressiveness (FoX, Gated Softmax, FPA).

Implementation and Practical Guidance

When to Use Each Architecture

1. GLA: Fast inference with moderate-to-long contexts (2K–20K tokens), when custom CUDA kernels are acceptable, and retrieval performance is not critical.

2. DeltaNet: Strong associative recall and in-context learning tasks; good for synthetic tasks (MQAR) and key-value association testing.

3. Gated DeltaNet: Production models requiring the best balance of recall, forgetting, and throughput (proven in Qwen3-Next; ICLR 2025).

4. Gated DeltaNet-2: Long-context retrieval workloads (multi-key NIAH) where decoupled channel-wise erase/write yields the strongest recall; when a gate-aware WY kernel is available.

5. RetNet: Extremely efficient inference without KV caching; suitable for edge/mobile deployments or toy models.

6. HGRN2: Multi-scale temporal hierarchies; when layer-specific forget bounds are desirable.

7. FoX: Length extrapolation and long-context understanding; when quadratic training is acceptable and replacing softmax is not desired.

8. Gated Softmax: Quick improvement to existing softmax models with minimal code changes; best for practitioners wanting immediate gains without major refactoring.

9. FPA: Long-context and structured sequences where a later event invalidates recent tokens but not older ones (delimiters, backspaces, reassignment); when softmax quality and exact KV decoding are required and a small per-head scalar state is acceptable.

Hardware Considerations

SSOG: Separable Sum of Gaussians Attention

This annex covers SSOG (Separable Sum of Gaussians), a geometric-field attention mixer. SSOG replaces the content-scored interaction of scaled dot-product attention with a learned mixture of Gaussians over relative position: each head owns a handful of Gaussian atoms—five numbers per atom, a center offset \((\mu_y, \mu_x)\), a width per axis \((\sigma_y, \sigma_x)\) and a mixture weight \(\lambda\)—forming a fixed field over the displacement between tokens. There are no query–key dot products anywhere: content never scores content, it only steers the field through small bounded residuals.

Because a 2D Gaussian factorises into a product of two 1D Gaussians, applying the field is two 1D filter passes per atom (rows, then columns of a \(\sqrt{N} \times \sqrt{N}\) token raster) followed by a \(\lambda\)-mix contraction. The induced \(N \times N\) attention matrix never exists, so the attention-specific cost is \(\mathcal{O}(R \cdot N \sqrt{N} \cdot d)\) for \(R\) atoms instead of \(\mathcal{O}(N^2 \cdot d)\). SSOG also skips the query and key projections entirely, paying \(2d^2\) of projections per layer instead of the \(4d^2\) of SDPA. In the reference recipe the full mechanism beats its matched SDPA twin on ImageNet-1k (65.28% vs 64.34% at \(\sim\)3M params; 72.02% vs 71.84% at \(\sim\)12M) while being \(\sim\)20% smaller and \(\sim\)30% cheaper per forward, and by +17 points on CIFAR-100 at matched recipe.

Mathematical Formulation

Let tokens be the row-major raster of a \(G_h \times G_w\) grid. The attention weight from token \(p\) to token \(q\) is the softmax-tempered log-sum-exp of the atom log-weights evaluated at their displacement: \[A(p, q) = \operatorname{softmax}_q\!\left( \tfrac{1}{\tau}\, s(p, q) \right), \qquad s(p, q) = \operatorname{logsumexp}_r \big( \log \lambda_r + \log \mathcal{N}(p - q;\ \mu_r, \sigma_r) \big),\] where \(\tau\) is a learnable temperature and \(\mathcal{N}(\cdot\,;\mu_r, \sigma_r)\) is a 2D Gaussian with diagonal covariance. Since the 2D Gaussian factorises, the softmaxed field applies separably: per axis, \(a_y[i, j] = \operatorname{softmax}_j(\log \mathcal{N}(\Delta y_{ij}; \mu_y, \sigma_y) / \tau)\) and analogously \(a_x\), followed by \[Y = \sum_{r=1}^{R} \lambda_r\, \big( a_x^{(r)} \big( a_y^{(r)}\, V \big) \big), \qquad V = x W_V,\] i.e. three einsum contractions per layer. With content steering (lookat) enabled, zero-initialised linear probes predict bounded per-token residuals on the field parameters behind cold-started softplus gates: \[\mu \leftarrow \mu_0 + s_\mu \cdot m \cdot \tanh(W_\mu x), \qquad \sigma \leftarrow \sigma_0 \, e^{s_\sigma \tanh(W_\sigma x)}, \qquad \lambda \leftarrow \operatorname{softmax}\big( \log \lambda_0 + s_\lambda \tanh(W_\lambda x) \big),\] with \(m\) the maximum offset in grid cells and gate scales \(s_\bullet = \operatorname{softplus}(\cdot)\) initialised at \(-8\) (\(\approx 3 \times 10^{-4}\)), so the model begins as a frozen geometric field and learns how far to open each content tap. Widths are reparameterised as \(\sigma = \operatorname{softplus}(\omega) + \epsilon + \sigma_{\text{floor}}\) to stay strictly positive.

Algorithmic Pseudocode

Key Characteristics

• Training complexity: \(\mathcal{O}(R \cdot N \sqrt{N} \cdot d)\) for fixed \(R\) atoms—near-linear by geometry alone, not by kernel approximation.

• Inference complexity: same per-forward cost; encoder-only (the separable field has no causal formulation), so there are no KV-cache semantics.

• Trainable: atoms \(\{\mu, \omega, \log\lambda\}\), temperature, and the steering probes are fully differentiable; cold-started gates open during optimisation.

• Strengths: strong geometric prior (+17 pts over SDPA on small data at matched recipe; beats SDPA on ImageNet-1k at 3M and 12M params with \(\sim\)20% fewer parameters and \(\sim\)30% fewer FLOPs); inherently interpretable (each head is a few plottable blobs—early layers become convolutions in disguise, middle layers strip detectors, late layers go global); the field lives on coordinates, granting zero-shot resolution transfer.

• Limitations: encoder-only; requires a fixed \(G_h \times G_w\) row-major raster (ViT needs cls_token=false; sequences need an explicit \(1 \times L\) grid); the per-query steering kernels add a memory term; language results are unproven—content scoring is likely load-bearing there.

Fast-Weight Attention for Continual Learning

This annex covers the Falcon family , a fast-weight attention framework that treats the recurrent fast-weight state transition as an online continual-learning rule under read-after-write (RAW) autoregressive semantics. The framework’s central observation is one of temporal alignment: for the prefix-prediction objective, the local fast-memory example revealed at step \(t\) is the prefix-aligned pair \((x_t, y_t) = (\phi(k_{t-1}),\, v_t)\) — one step shifted relative to the same-step association \((\phi(k_t), v_t)\) used by DeltaNet and standard linear attention. The same-step pairing remains causal, but it optimizes a different internal objective. The framework derives normalized first-order updates for two objectives (squared-error regression and negative inner product) and three step-size structures (scalar, per-column, sliding window), giving six variants: Falcon-1/2/3 (regression family) and Falcon-1A/2A/3A (inner-product family). All six are implemented in the codebase under model.attention.falcon with a shared configuration block, each providing both a recurrent reference scan and an SSD-style chunk-parallel kernel.

Preliminaries: Alignment, Normalization, and Stability

Read-after-write semantics.

With feature maps \(\phi(k_t)\) (QK-RMSNorm by default, \(\ell_2\) norm optional), the prefix-aligned stream sets \(x_t = \phi(k_{t-1})\) for \(t \ge 2\) with boundary sentinel \(x_1 = 0\); the readout is \(o_t = S_t^\top \phi(q_t)\). The write at step \(t\) therefore updates the memory before the query at step \(t+1\) reads it, so no token’s key ever associates with its own value.

Normalized first-order step sizes.

A gradient step on the local objective \(\ell_t(S)\) with the objective-matched denominator gives \[\eta_t = \frac{\beta_t}{\text{statistic}_t + \lambda_t + \varepsilon}, \qquad \beta_t \in (0, 2),\] where the statistic is the squared key energy \(\|x_t\|^2\) (Falcon-1/1A/2/2A), the window spectral statistic \(\mu_t^{(B)} = \lambda_{\max}(X^\top X)/B_t\) (Falcon-3), or the window mean energy \(\bar{E}_t^{(B)}\) (Falcon-3A). Objective-matched normalization guarantees per-step descent and removes the scale sensitivity of fixed learning rates. The plasticity gain \(\beta_t\), the forgetting ridge \(\lambda_t\), the temporal alignment, and the rehearsal window \(B\) are thereby four previously entangled knobs, separated.

Positive-decay renormalization.

When \(\lambda_t > 0\), the per-step shrinkage \(\alpha_t = \min(\eta_t \lambda_t,\, 1 - \varepsilon_\gamma)\) must be applied stably. The implementation clamps \(\alpha_t\) strictly below \(1\), computes chunk-local log-prefixes in fp32 via \(\log(1 - \alpha_t)\) and cumulative sums, rescales the drivers to the no-decay domain (\(\tilde{\eta}_t = \eta_t / \gamma_t\), \(\tilde{v}_t = v_t / c_{t-1}\) with \(c_{t-1} = \prod_{r \le t-1} \gamma_r\)), and rescales outputs and the emitted state back by the accumulated \(\delta\). Boundary sentinels are \((\alpha_1, \gamma_1) = (0, 1)\) and \(\eta_1 = 0\), so the first step is a pure read.

The Regression Family (Falcon-1/2/3)

The squared-error objective on the shifted pair is \[\ell_t(S) = \tfrac{1}{2}\|v_t - S^\top x_t\|^2 + \tfrac{\lambda_t}{2}\|S\|_F^2, \qquad \nabla_S \ell_t = x_t (S^\top x_t - v_t)^\top + \lambda_t S,\] which is \(L_t\)-smooth with \(L_t = \|x_t\|^2 + \lambda_t\).

Falcon-1: Scalar NLMS Regression

The scalar delta-rule variant — the classical NLMS recursion of adaptive filtering upgraded with the shifted write stream and an explicit ridge: \[S_t = (1 - \eta_t \lambda_t)\, S_{t-1} + \eta_t\, x_t \big(v_t - S_{t-1}^\top x_t\big)^\top, \qquad o_t = S_t^\top \phi(q_t).\] With \(\varepsilon = 0\), \(\lambda_t = 0\) and the unshifted assignment \(x_t = \phi(k_t)\) this recovers DeltaNet exactly. The eigenstructure of the transition \(A_t = \gamma_t I - \eta_t x_t x_t^\top\) (with \(\gamma_t = 1 - \eta_t\lambda_t\)) shows that the write direction contracts by \(1 - \beta_t\) while the orthogonal subspace carries \(\gamma_t\): \(\beta_t > 1\) yields a stable sign-flip along the write direction, useful for state tracking. The chunk-parallel form uses the WY representation \(\prod_s (I - \eta_s x_s x_s^\top) = I - U T U^\top\) with a single-residual unit-lower-triangular solve \(L = \operatorname{tril}(G, -1) + I\) where \(G = X^\top X\) (Algorithm 8 of the paper), after the positive-decay renormalization.

Falcon-2: Per-Column NLMS Regression

The ridge loss decomposes across value coordinates, so each value channel \(j\) receives its own step size while sharing the NLMS normalizer: \[\eta_{j,t} = \frac{\beta_{j,t}}{\|x_t\|^2 + \lambda_t + \varepsilon}, \qquad S_t = S_{t-1}\big(I - \lambda_t \operatorname{Diag}(\boldsymbol{\eta}_t)\big) + x_t \big(\boldsymbol{\eta}_t \odot \mathbf{r}_t\big)^\top .\] The chunk-parallel kernel forms the shared key Gram \(G = X^\top X\) and solves the batch of per-channel unit-lower-triangular systems \(L_j = I + \operatorname{tril}\big(G \odot (\Sigma_j \Sigma_j^\top), -1\big)\) with \(\Sigma = \sqrt{\tilde{\boldsymbol{\eta}}}\) (Algorithm 1). Defined but not separately benchmarked in the paper.

Falcon-3: Sliding-Window Minibatch Regression

Bounded rehearsal: each step regresses the fast memory on the last \(B_t \le B\) prefix-aligned pairs, with all window residuals evaluated at the pre-update state \(S_{t-1}\): \[\mu_t = \frac{\lambda_{\max}\big(X_{\mathcal{I}_t}^\top X_{\mathcal{I}_t}\big)}{B_t}, \qquad \eta_t = \frac{\beta_t}{\mu_t + \lambda_t + \varepsilon}, \qquad S_t = (1 - \eta_t \lambda_t)\, S_{t-1} + \frac{\eta_t}{B_t} \sum_{j \in \mathcal{I}_t} x_j \big(v_j - S_{t-1}^\top x_j\big)^\top ,\] with \(\mathcal{I}_t = \{\max(2, t - B + 1), \dots, t\}\). Window-averaging keeps the update magnitude invariant to \(B\): each pair is replayed in exactly \(B\) consecutive updates at weight \(1/B\), so its cumulative injection is \(\eta\)-independent of \(B\). The state alone is not Markov: exact continuation requires a FIFO tail of the last \(B - 1\) pairs, and the chunk kernels gather an extended slice with \(|\mathcal{J}| \le C + B - 1\). The chunk-parallel form is a rank-\(B\) affine recurrence solved by block forward substitution over the (time \(\times\) rank) extended index space, with same-time rank components uncoupled (Algorithm 3).

The Inner-Product Family (Falcon-1A/2A/3A)

The negative inner-product objective \[\ell_t^{\mathrm{ip}}(S) = -\langle S^\top x_t,\, v_t \rangle + \tfrac{\lambda_t}{2}\|S\|_F^2 \qquad \Rightarrow \qquad \nabla_S \ell_t^{\mathrm{ip}} = -x_t v_t^\top + \lambda_t S\] yields purely additive (Hebbian) writes with energy-normalized gains.

Falcon-1A and Falcon-2A

\[\begin{gathered} \text{Falcon-1A:}\quad S_t = \gamma_t S_{t-1} + \eta_t\, x_t v_t^\top, \quad \gamma_t = 1 - \min(\eta_t \lambda_t,\, 1 - \varepsilon_\gamma);\\ \text{Falcon-2A:}\quad S_t = S_{t-1}\operatorname{Diag}(\boldsymbol{\gamma}_t) + x_t (\boldsymbol{\eta}_t \odot v_t)^\top . \end{gathered}\] With \(\lambda_t = 0\) and the unshifted assignment this is exactly standard linear attention / Mamba-2 accumulation; the shifted stream makes it the one-step-shifted next-latent variant. The A-family’s denominator is not required by curvature (the objective is linear) — it retains the energy-normalized write as a magnitude stabilizer. The chunk-parallel form is decay-mask linear attention with \(M_{t,i} = \eta_i \prod_{r=i+1}^{t} \gamma_r\) (per-channel masks for Falcon-2A).

Falcon-3A: Sliding-Window Inner-Product

The windowed additive variant — the family’s best length-extrapolation configuration: \[\bar{N}_t^{(B)} = \frac{1}{B_t} \sum_{j \in \mathcal{I}_t} x_j v_j^\top, \quad \bar{E}_t^{(B)} = \frac{1}{B_t} \sum_{j \in \mathcal{I}_t} \|x_j\|^2, \quad \eta_t = \frac{\beta_t}{\bar{E}_t^{(B)} + \lambda_t + \varepsilon}, \qquad S_t = \gamma_t S_{t-1} + \eta_t\, \bar{N}_t^{(B)} .\] The window-banded decay mask factors as \(M = D \cdot \operatorname{Diag}(\boldsymbol{\eta}) \cdot A\), where \(A\) is the \(B\)-banded window operator and \(D\) the causal decay kernel, evaluated over the extended slice (Algorithm 4).

Configuration Interface

All six variants share one nested block, model.attention.falcon, selected through the layer_pattern entries falcon1, falcon2, falcon3, falcon1a, falcon2a, falcon3a:

Implementation notes: the recurrent and chunk-parallel kernels are verified numerically equivalent (chunk \(\equiv\) recurrent to \(\le 3.6 \times 10^{-7}\) across all six variants, including decay, multi-chunk, and windowed paths); gradients flow through the fp32 triangular-solve islands (feature drivers take only \(\sqrt{\tilde{\eta}}\) while value targets take \(\sqrt{\tilde{\eta}}\, e^{-u_{t-1}}\), with the square root clamped at \(10^{-12}\) for NaN-free backward at the \(\eta = 0\) boundary); positive exponents are clamped at \(80\).

Published Results

Relation to Prior Work

• DeltaNet: the delta rule is the squared-error one-gradient-descent step with the same-step pairing; Falcon recovers it with the critical index shift plus objective-matched normalization and explicit \(\lambda_t\) shrinkage.

• Gated DeltaNet: adds a free learned carry gate; Falcon derives the carry \(\gamma_t = 1 - \eta_t \lambda_t\) from the local objective and the normalized step-size parameterization instead of adding an independent gate, and defaults to QK-RMSNorm rather than \(\ell_2\) normalization.

• Linear attention / Mamba-2: the additive numerator is one gradient step on the inner-product objective with the unshifted assignment; the chunkwise kernels reuse the SSD scheme.

• TTT / Titans / ATLAS: fast-weight updates as test-time training; Falcon-3 is a sliding-window specialization of the internal-objective view with strict next-latent alignment.

• Adaptive filtering (LMS/NLMS/RLS): Falcon imports the stability and normalization principles of NLMS into fast-weight attention under strict causality.

When to Use Each Variant

1. Falcon-1: best language-modeling perplexity of the family (FineEdu 17.10, beating Gated DeltaNet 17.32); principled forgetting tied to the local objective.

2. Falcon-2: finer-grained plasticity than Falcon-1 at the cost of heavier per-channel chunk kernels; unbenchmarked in the paper.

3. Falcon-3: controlled rehearsal horizon with exact window statistics; most complex kernel.

4. Falcon-1A: simple and stable additive variant; second-best length extrapolation (85.9%).

5. Falcon-2A: channel-wise forgetting granularity; unbenchmarked.

6. Falcon-3A: best length extrapolation of the paper (87.2% on 33–48-digit addition); carries arithmetic tasks.

Convolution-Augmented Attention: Conformer and Fast Conformer

Executive Summary

The Conformer block is the toolkit’s convolution-augmented sequence mixer, exposed as conformer. It keeps the global receptive field of scaled dot-product attention and adds a depthwise-convolution branch that models local, narrow-band structure—the inductive bias speech recognition needs and the one plain self-attention lacks. The block sandwiches multi-head self-attention (MHSA) with relative sinusoidal positional encodings between two Macaron-style half-step feed-forward layers , and inserts the convolution module between the attention and the second FFN. In the reference recipe it reaches \(1.9\%\)/\(3.9\%\) WER on LibriSpeech test-clean/test-other without an external language model (Conformer-L, 118.8M parameters), and it is the mixer recommended for the audio model class of Appendix 33. Its long-form sibling, Fast Conformer , keeps the convolution module and the block layout but replaces global self-attention with limited-context attention—a fixed local window plus a global token pinned at the sequence head—so hours-long audio decodes with bounded memory; it is documented below. The toolkit runs both feed-forward layers at full step (the paper’s “full-step” variant of its Table 5) and lets the outer layer own the second one; Table [tab:conformer-mapping] lists the deviations.

The mixers are available to every model class: the text encoder, the causal decoder (where depthwise padding is shifted to stay streaming-safe), the Vision Transformer, and frankenstein_audio. Because the blocks live inside HybridLayer, they compose with all 48 other mixers, the 6 norms, the 4 residual strategies, MoE, Mixture-of-Depths, mHC, AttnRes, BitNet and the 38 optimizer families.

Conformer: Convolution-augmented Attention

Mathematical Formulation

For a block input \(x \in \mathbb{R}^{B \times n \times d}\) the reference block of Fig. 1 of is \[\begin{aligned} \tilde{x} &= x + \tfrac{1}{2}\,\mathrm{FFN}_1\big(\mathrm{LN}(x)\big), \label{eq:conformer-ffn1}\\ x' &= \tilde{x} + \mathrm{MHSA}\big(\mathrm{LN}(\tilde{x})\big), \label{eq:conformer-mhsa}\\ x'' &= x' + \mathrm{Conv}\big(\mathrm{LN}(x')\big), \label{eq:conformer-conv}\\ y &= \mathrm{LN}\Big(x'' + \tfrac{1}{2}\,\mathrm{FFN}_2\big(x''\big)\Big). \label{eq:conformer-ffn2} \end{aligned}\] The two feed-forward terms carry the Macaron half-step coefficient: each contributes half of a full residual update, one before and one after the attention–convolution pair.

Feed-forward module.

With hidden width \(d_{ff}\), \[\mathrm{FFN}(x) = W_2\,\mathrm{Dropout}\big(\mathrm{Swish}(W_1\,\mathrm{LN}(x) + b_1)\big) + b_2, \qquad \mathrm{Swish}(a) = a\,\sigma(a), \label{eq:conformer-ffn}\] where \(W_1 \in \mathbb{R}^{d_{ff} \times d}\), \(W_2 \in \mathbb{R}^{d \times d_{ff}}\) and \(\sigma\) is the logistic sigmoid. The paper uses \(d_{ff} = 4d\).

Multi-head self-attention with relative positional encoding.

Queries, keys and values are linear projections of the normalized block input, \(Q = \mathrm{LN}(\tilde{x})W_Q\), \(K = \mathrm{LN}(\tilde{x})W_K\), \(V = \mathrm{LN}(\tilde{x})W_V\), split into \(H\) heads of width \(d_h = d/H\). Following the Transformer-XL relative scheme , the score of query \(i\) against key \(j\) adds a sinusoidal function of the displacement \(i-j\) instead of an absolute position: \[e_{ij} = q_i^{\top} k_j + (q_i + v)^{\top} R_{i-j}, \qquad A_{ij} = \operatorname{softmax}_j\!\Big(\frac{e_{ij}}{\sqrt{d_h}}\Big), \qquad O = A V W_O, \label{eq:conformer-relpe}\] with \(v\) a learned bias vector and, for \(m = 0, \dots, d_h/2 - 1\), \[R_{\delta,\,2m} = \sin\!\big(\delta\,/\,10000^{2m/d_h}\big), \qquad R_{\delta,\,2m+1} = \cos\!\big(\delta\,/\,10000^{2m/d_h}\big). \label{eq:conformer-sinusoid}\] Because the encoding depends only on \(\delta = i-j\), the block generalizes across utterance lengths. The toolkit routes positional information through the model-wide machinery that every other mixer uses: model.positional_encoding (RoPE by default) is applied to \(Q/K\), and conformer_use_pe (default true) lets a layer opt out. The paper’s relative sinusoidal scheme is closest to sinusoidal_rotary (sinusoidal frequencies applied as a rotation on \(Q/K\), relative by construction); alibi/bam supply a distance bias instead of a rotation. A dedicated Transformer-XL-style relative-sinusoidal module is not implemented. In decoder mode the score matrix receives the usual causal mask before the softmax, and the fixed-softmax path of Appendix 26.3 is used when numerics.fixed_softmax is enabled.

Convolution module.

The module applies a gated depthwise convolution over the time axis: \[\begin{aligned} h &= \mathrm{LN}(x'), \\ [u,\, g] &= \mathrm{Conv1d}_{d \to 2d}(h), \qquad \tilde{h} = u \odot \sigma(g), \label{eq:conformer-glu}\\ z_t &= \sum_{k=0}^{K-1} w_k\, \tilde{h}_{t + k - \lfloor K/2 \rfloor}, \label{eq:conformer-depthwise}\\ c &= \mathrm{Dropout}\Big(\mathrm{Conv1d}_{d \to d}\big(\mathrm{Swish}(\mathrm{BN}(z))\big)\Big), \label{eq:conformer-conv-out} \end{aligned}\] where \(\odot\) is the elementwise product (a GLU that halves the \(2d\) channels back to \(d\)), \(w \in \mathbb{R}^{K}\) is the depthwise kernel and \(\mathrm{BN}\) normalizes each channel over the batch and time axes, \[\mathrm{BN}(z)_{b,t,c} = \gamma_c\,\frac{z_{b,t,c} - \mu_c}{\sqrt{\sigma_c^2 + \epsilon}} + \beta_c . \label{eq:conformer-bn}\] In decoder mode the sum in Equation [eq:conformer-depthwise] is shifted to \(z_t = \sum_{k=0}^{K-1} w_k \tilde{h}_{t-K+1+k}\), so that no future frame enters the convolution; encoder and ViT modes keep the symmetric “same” padding of the paper.

Algorithmic Pseudocode

Implementation Mapping and Deviations

Table [tab:conformer-mapping] maps the paper block onto the toolkit. The substantive deviations are the two residual scalings and the second feed-forward layer: the paper keeps both \(\tfrac12\) Macaron coefficients and wraps \(\mathrm{FFN}_2\) in the block, whereas conformer uses full steps and returns the state after the convolution module, letting the outer HybridLayer apply its shared \(\mathrm{norm}_2 + \mathrm{FFN}\) (also at full step). The decomposition is deliberate—it keeps MoE routing, Mixture-of-Depths, mHC/Hyperloop streams, AttnRes, FlashNorm fusion, BitNet projections and the per-parameter-group optimizer routing available for the second FFN, none of which would survive a private FFN inside the mixer.

Why the full-step FFN is safe.

Dropping the Macaron \(\tfrac12\) changes the per-branch residual scale, not the expressive form of the block: the full-step variant is exactly what ablate in their Table 5, where it moves dev-clean/dev-other WER from \(1.9/4.4\) (half step) to \(1.9/4.5\) (full step), i.e. one tenth of a point. The toolkit therefore uses full steps for both feed-forward layers; the first stays private to the mixer, and the second is the shared one, which preserves the composability described above.

Published Ablations

Table [tab:conformer-ablations] summarizes the ablations that motivate the default configuration. Values are the ones reported by ; the qualitative ordering is what the defaults encode.

Reuse of Existing Infrastructure

Because ConformerAttention returns a delta and inherits the mixer(config, pos_encoder=...) constructor, the entire stack is reused without special cases:

• Model classes: text encoder (mlm/sbert/text_classification), causal decoder (causal_lm), Vision Transformer (patch_prediction/classification/segmentation) and frankenstein_audio (asr, audio_classification).

• Normalization: all 6 variants, including partial RMSNorm and FlashNorm fusion on \(\mathrm{norm}_1\)/\(\mathrm{norm}_2\); the convolution module keeps its own BatchNorm unless conformer.conv_norm selects group or layer normalization.

• Routing and residuals: MoE, Mixture-of-Depths, mHC, Hyperloop, AttnRes and the BitNet ternary projections (the private FFN and the convolution projections are BitLinear-aware).

• Optimizers: all 23 families with per-group prefixes; mixer parameters, including the depthwise kernel and the convolution norms, route to the attention group.

• KV cache: conformer is registered as not cacheable (supported=False), because the depthwise-convolution state and the BatchNorm running statistics would have to be carried across steps. generate() falls back to the full-recompute path, which is exact.

When to Use It

• Speech: the default mixer for asr and audio_classification; at 40 ms frames after \(4\times\) subsampling the kernel of 32 spans \(\sim\)1.28 s of context, roughly the width of a phone sequence.

• Text and vision: a strong drop-in for standard when the data has local structure (character-level models, code, OCR, patch sequences); it costs one depthwise convolution and two extra norms per layer over a plain attention block.

• Streaming decoding: usable in mode: decoder with causal convolution padding, but without the incremental cache of Appendix 34; the depthwise recurrence is the missing piece.

Fast Conformer: Limited-Context Attention with a Head-Pinned Global Token

Mathematical Formulation

Fast Conformer attacks the quadratic cost of Conformer self-attention in two ways: it raises the front-end downsampling to \(8\times\), and, for long-form inputs, it replaces global attention with limited-context attention. The paper describes the mixer in prose rather than in display equations: “we use a single global attention token, which attends to all other tokens and has all other tokens attend to it, using a separate set of query, key and value linear projections, while others attend in a fixed-size window surrounding each token” (§2.2), while the evaluated long-form model “has shared query, key and value projection layers that are used for global and local attention” (§3.4) and sets “the size of the limited context …to 128 steps on each side of a token”. This implementation encodes exactly that description and inherits the global-attention dataflow from Longformer , which the paper explicitly cites for the mechanism and for its overlapping-chunk kernel.

Let \(X \in \mathbb{R}^{n \times H}\) be the layer input, \(h\) the number of heads and \(d = H/h\). The shared projections are \[\mathbf{Q} = X\mathbf{W}_Q, \qquad \mathbf{K} = X\mathbf{W}_K, \qquad \mathbf{V} = X\mathbf{W}_V,\] with \(\mathbf{W}_\ast \in \mathbb{R}^{H \times H}\) and no bias. The optional global projections of the paper’s §2.2 variant are a second set, \[\mathbf{Q}^{g} = X\mathbf{W}_Q^{g}, \qquad \mathbf{K}^{g} = X\mathbf{W}_K^{g}, \qquad \mathbf{V}^{g} = X\mathbf{W}_V^{g},\] which adds \(3H^2\) parameters; the §3.4 evaluated model sets \(\mathbf{W}_\ast^{g} = \mathbf{W}_\ast\).

The global positions are the head of the sequence, spaced \(s\) apart, \[\mathcal{G} = \{\,0, s, 2s, \dots, (g-1)s\,\},\] so a global token is a designated input position rather than an extra learned embedding (the paper’s “addition of a global token”). The attention pattern is \[M_{ij} = \big[\,|i-j| \le w/2\,\big] \;\vee\; \big[\,j \in \mathcal{G}\,\big] \;\vee\; \big[\,i \in \mathcal{G}\,\big],\] where the first term is the local window, the second lets every query attend to the global tokens, and the third lets a global token attend to every position; \(M\) is intersected with the causal mask in decoder mode. With the effective keys and values \[\tilde{\mathbf{k}}_j = \begin{cases} \mathbf{k}^{g}_j & j \in \mathcal{G}\ \text{and separate}, \\ \mathbf{k}_j & \text{otherwise}, \end{cases} \qquad \tilde{\mathbf{v}}_j = \begin{cases} \mathbf{v}^{g}_j & j \in \mathcal{G}\ \text{and separate}, \\ \mathbf{v}_j & \text{otherwise}, \end{cases}\] the output is standard scaled dot-product attention over the allowed pairs, \[z_{ij} = \frac{\mathbf{q}_i \cdot \tilde{\mathbf{k}}_j}{\sqrt{d}}, \qquad \mathbf{o}_i = \sum_{j:\,M_{ij}} \operatorname{softmax}_j(z_{ij})\,\tilde{\mathbf{v}}_j, \qquad \mathbf{Y} = \mathbf{O}\mathbf{W}_O .\] In the separate variant a global query row is computed with the global projections over the whole causal prefix, \(\mathbf{o}_i = \sum_{j \le i} \operatorname{softmax}_j(\mathbf{q}^{g}_i \cdot \mathbf{k}^{g}_j/\sqrt{d})\,\mathbf{v}^{g}_j\), matching Longformer’s global-attention design. With shared projections all three cases collapse to one \(\mathbf{Q}/\mathbf{K}/\mathbf{V}\) set: the global tokens are simply positions with unmasked rows and columns.

Algorithmic Pseudocode

In incremental decoding the retained state is a bounded ring: the pinned head region covers \(\mathcal{G}\) and the ring keeps the last \(w/2+1\) keys, so each step costs \(\mathcal{O}(w+g)\). The exact pattern is evaluated against the cached absolute positions (unlike a retention-equals-pattern shortcut, because a spaced global prefix may contain non-global positions).

Key Characteristics

• Complexity: \(\mathcal{O}(n\,(w+g)\,d)\) training and \(\mathcal{O}(w+g)\) per decoded token with the paper’s overlapping-chunk kernel; the reference implementation here masks a dense score matrix, which is mathematically identical but \(\mathcal{O}(n^2)\) in compute. Memory is bounded by \(w+g\).

• Trainable: Yes, natively from scratch (the pattern runs in both training and evaluation). The paper fineTunes a full-context checkpoint for 10k steps at peak LR \(10^{-6}\) with a 1k-step warmup and cosine decay, but the mixer’s weights are ordinary attention projections.

• Strength: On long-form benchmarks the global token is the decisive ingredient: normalized greedy WER improves from 9.18 (Conformer) to 8.25 (limited context) to 7.5 (plus global token) on TED-LIUM v3, and from 18.26 to 16.08 to 11.85 on Earnings-21, while the single-pass maximum duration grows from 15 to 675 minutes on one A100. The encoder is \(2.8\times\) faster and uses \(2.9\times\) fewer multiply-accumulates than the original Conformer.

• Limitation: The paper gives no closed-form equations for the global-token path (the formulation above is a documented reconstruction from its prose, Longformer and the authors’ NeMo recipe). The default is a single global token with shared projections; the separate-projection variant costs \(3H^2\) extra parameters and is only exact in the cached path when the global queries are pre-filled from position zero. The efficient chunked kernel is not implemented.

Optimizer Families

The Evolution of Optimization in Neural Networks

The optimizer survey frames transformer optimization as a response to three structural pressures: non-convex loss landscapes, severe curvature heterogeneity across parameter blocks, and the memory cost of storing optimizer state for very large models. The report argues that the field has diverged into several trajectories: adaptive first-order baselines, variance-reduction methods, memory-efficient methods, structured second-order preconditioners, schedule-free methods, and orthogonality-oriented updates.

Tutorial: From Gradient Descent to Optimizer Families

The rest of this appendix is a reference: one entry per optimizer, with its update rule, pseudocode, and memory profile. This subsection is the on-ramp. It introduces the five ideas that every entry reuses—the loss surface, the gradient, the step size, the batch, and the two statistics that later methods add—and then compares the families as families, without deriving any single method. A reader who finishes it should be able to open any of the following sections and recognise which of the five ideas it is modifying.

One problem, many recipes

Training a neural network is an optimization problem wrapped in bookkeeping. Write \(\theta\) for the vector of all trainable numbers (weights and biases, tens of millions for a small transformer), \(L(\theta)\) for the loss, and \(f(\theta)\) for the objective actually minimised, usually \(L\) plus a regulariser. The architecture only decides which surface \(L\) is: the number of layers, the attention pattern, the normalisation, the activation. The optimizer decides how the walk over that surface proceeds. Every optimizer in this appendix can be written as one line, \[\theta_{t+1} \;=\; \theta_t \;-\; \eta_t\, u_t, \qquad u_t = u\!\left(g_t, g_{t-1}, \dots,\, \theta_t\right), \label{eq:opt-tut-generic}\] where \(g_t=\nabla f(\theta_t)\) is the gradient and \(u_t\) is a direction built from the current gradient and, possibly, a small amount of remembered state. Gradient descent takes \(u_t=g_t\). Momentum takes \(u_t\) to be a decaying average of past gradients. Adaptive methods divide \(u_t\) coordinate by coordinate. Second-order methods rotate it with a curvature estimate. Orthogonality methods reshape it as a matrix. Nothing else changes: the sign in front of \(\eta_t\) is always negative, because the gradient points uphill.

The usual picture is a hiker descending a mountain in fog . The hiker feels the slope underfoot but cannot see the valley, so the strategy reduces to three questions: how confident is the slope reading (one example or a thousand?), how long is the stride, and how much memory of the terrain is worth carrying. Figure 75 turns those questions into the map of the families that follow.

Gradient descent: walk downhill

For a single weight \(w\), the derivative \(f'(w)\) is the slope: positive means the loss grows to the right, so the useful move is to the left. In several dimensions the gradient \(\nabla f\) collects all partial derivatives and points in the direction of steepest increase; its negative is the direction of steepest decrease. Taking a step of length \(\eta\) along it gives \[\theta_{t+1} = \theta_t - \eta\,\nabla f(\theta_t), \label{eq:opt-tut-gd}\] the oldest algorithm in the appendix, usually credited to Cauchy in the nineteenth century and analysed by Curry in 1944 . The scalar \(\eta>0\) is the learning rate or step size, and it is the only tunable number.

Two facts explain almost everything that follows. First, the step size has a hard ceiling. If the gradient changes at most \(L\) times faster than the parameters change (the gradient is \(L\)-smooth), then [eq:opt-tut-gd] decreases the loss for \(\eta<2/L\) and diverges beyond it; the safe monotone choice is \(\eta\le 1/L\) . Second, contour lines of a quadratic loss are ellipses whose axes have very different lengths, and the negative gradient is perpendicular to the contour but generally does not point at the minimum—except when the ellipses are circles . The walk therefore zig-zags across the steep direction while creeping along the flat one. Figure 76 shows the three regimes on the same bowl, and Figure 77 shows what they look like as a loss curve.

When the curvature is unknown, the classical remedy is a line search: try \(\eta=1\), halve it until the loss actually decreases, and repeat . Deep learning rarely does this, because every trial costs a full pass over the data; it schedules \(\eta\) by hand instead. Theory supplies the pace: for a convex \(L\)-smooth objective, gradient descent reaches \(f(\theta_t)-f^\star\le 2L\|\theta_0-\theta^\star\|^2/(t+1)\), that is, \(O(1/t)\) . Halving the distance to the optimum costs a doubling of iterations—a fair summary of why first-order methods need so many steps.

Where the gradient comes from: backpropagation

Nothing above said how to obtain \(g_t\). For a network with millions of weights there is no formula to differentiate by hand, and finite differences would cost one forward pass per parameter. Backpropagation computes all partial derivatives in a single backward sweep, using only the chain rule.

In a feed-forward network each layer alternates an affine map and an activation, \(a^{(\ell)}=W^{(\ell)}z^{(\ell-1)}+b^{(\ell)}\) and \(z^{(\ell)}=\sigma(a^{(\ell)})\). The forward pass stores every \(a\) and \(z\). The backward pass then propagates one number per neuron, \(\delta^{(\ell)}=\partial e/\partial a^{(\ell)}\), from the loss towards the input, \[\delta^{(L)} = \frac{\partial e}{\partial z^{(L)}}\odot\sigma'(a^{(L)}), \qquad \delta^{(\ell)} = \sigma'(a^{(\ell)})\odot \left(W^{(\ell+1)\top}\delta^{(\ell+1)}\right), \qquad \frac{\partial e}{\partial W^{(\ell)}} = \delta^{(\ell)}\,z^{(\ell-1)\top}, \label{eq:opt-tut-backprop}\] where \(\odot\) is the element-wise product. The gradient of a weight is therefore the product of a backward number and a forward number—the two quantities drawn in Figure 78. The recursion is linear in the layer index, so the whole sweep costs about the same as one extra forward pass, roughly twice the forward cost per step, and it must keep the activations of every layer in memory until they are used .

The important consequence for this appendix is that backpropagation is not an optimizer. It is a gradient factory, and it always returns the exact gradient of whatever batch it was given. Every difference between the families that follow lives after the factory, in how \(u_t\) in [eq:opt-tut-generic] is formed.

Stochastic and mini-batch gradient descent

With \(n\) examples the loss is an average, \(f(\theta)=\frac1n\sum_i f_i(\theta)\), so the exact gradient is an average of \(n\) per-example gradients and one step costs \(n\) backward passes. For a dataset of a billion tokens that is unaffordable, and it is also wasteful: neighbouring examples usually say almost the same thing about the parameters.

Stochastic gradient descent, in the line of the Robbins–Monro stochastic approximation , replaces the average by one term, \(\hat g_t=\nabla f_{i_t}(\theta_t)\) with \(i_t\) drawn uniformly, giving \[\theta_{t+1} = \theta_t - \eta_t\,\nabla f_{i_t}(\theta_t). \label{eq:opt-tut-sgd}\] The estimate is unbiased, \(\mathbb{E}[\hat g_t]=g_t\), so on average it points the right way; its error falls as \(\sigma/\sqrt{b}\) when a batch of \(b\) examples is used instead of one . Modern libraries call this mini-batch method “SGD” and almost never use \(b=1\): one pass over the data (an epoch) then contains \(\lfloor n/b\rfloor\) updates rather than one . Figure 79 shows both faces of the trade: with the same number of steps and the same step size, small batches wander, and the noise they inject is precisely the quantity that shrinks as \(1/\sqrt{b}\).

Noise is not only a cost. Because each step is a random approximation, the iterate can leave a shallow basin or slide across a saddle that a full-batch method would treat as a stopping point . The price appears in the rate. For a convex \(L\)-smooth objective, SGD with \(\eta_t\propto1/\sqrt t\) satisfies \(f(\theta_t)-f^\star=O(\log t/\sqrt t)\), and with a constant step it converges to a neighbourhood of size \(O(\eta)\) and then stalls . Full-batch gradient descent is therefore more accurate per iteration and much slower per iteration; SGD is the opposite. Figure 80 puts the two statements on the same axes, once per iteration and once per unit of wall-clock, which is the comparison that actually matters when a step costs \(n/b\) times less.

Two knobs that every later method turns

Everything after this point is a refinement of two observations about equation [eq:opt-tut-gd], and it is worth separating them before meeting the families.

Memory of the past.

A ball rolling down the same valley keeps its velocity: it accumulates speed along a consistent slope and cancels the oscillation when the slope keeps changing sign. Momentum implements exactly that, replacing the gradient by a decaying sum of gradients, \[v_t = \beta\,v_{t-1} - \eta\,g_t, \qquad \theta_{t+1} = \theta_t + v_t, \label{eq:opt-tut-momentum}\] with \(\beta\in[0,1)\) the decay. The effective step is \(\eta/(1-\beta)\) along a steady direction, which is how a small \(\eta\) can still move quickly. Nesterov momentum makes the same move but reads the gradient at the position the velocity is heading towards, so the correction arrives before the overshoot does . Figure 81 shows the composition of the two terms and its effect on the valley.

One step size per coordinate.

The second observation is that the valley of Figure 76(b) has no single good \(\eta\): what is stable along the steep axis is negligible along the flat one. Adaptive methods keep a running average of squared gradients and divide by its square root, so a coordinate that has been receiving large gradients gets a smaller step and a quiet coordinate gets a larger one, \[s_t = \gamma\,s_{t-1} + (1-\gamma)\,g_t^2, \qquad \theta_{t+1} = \theta_t - \eta\,\frac{g_t}{\sqrt{s_t}+\epsilon}. \label{eq:opt-tut-rms}\] The scheme grew out of rprop, which used only the sign of the gradient, and its mini-batch version, rmsprop ; AdaGrad is the cumulative rather than decaying version . Figure 82 shows the effect: the two axes of the valley receive comparable displacements, at the cost of a per-coordinate state and a step that no longer shrinks on its own—the residual jitter of size \(\approx\eta\) around the minimum is what later motivates decaying or schedule-free step sizes.

The two knobs compose, and their composition is the reason AdamW is the default in practice: a first moment for direction, a second moment for scale, a bias correction for the empty first steps, and a decoupled weight decay for regularisation . Nothing in equations [eq:opt-tut-momentum] and [eq:opt-tut-rms] is specific to transformers; what differs across the families is which statistics are kept, how they are combined, and how many bytes per parameter that costs. A compact survey of the resulting zoo is given by .

How the families differ

The appendix that follows contains thirty-eight methods, which is too many to hold in mind at once. They are easier to remember as eight families, each answering one of the questions in Figure 75 and each distinguished by what it remembers and what it pays.

• First-order baselines (SGD with momentum, AdamW, RAdam) keep one or two statistics per parameter and assume the landscape is smooth enough for a locally linear model. They are the reference against which everything else is measured.

• Momentum and variance reduction (Adan, ADOPT, AdEMAMix, MARS, Cautious AdamW) keep more than two buffers, or a gradient difference, to reduce the variance of the update itself. They buy stability and speed at the cost of memory and of a second gradient evaluation or a longer history.

• Memory-efficient and low-rank (Adafactor, GaLore, APOLLO, Lion, AdamPM) accept a worse estimate of the second moment—a factored one, a projected one, or none—precisely to shrink the state. Lion goes as far as keeping only the sign of the update.

• Schedule-free and parameter-free (Schedule-Free AdamW, Prodigy, LAMB for large batches) move the burden of choosing \(\eta_t\) from the practitioner to the algorithm, usually by averaging iterates or by estimating the step from the loss itself.

• Curvature and second-order (Sophia, Shampoo, SOAP) estimate how the gradient changes, not just its size, and precondition the step with that estimate. They converge in fewer iterations and cost more per iteration.

• Geometric and orthogonality (Muon, Turbo-Muon, AF-Muon, ROOT, NorMuon, Dion3) treat a weight matrix as a matrix: they orthogonalise or otherwise whiten the update so that all singular directions move at the same rate, which is a cheap surrogate for curvature on matrix-shaped parameters.

• Adaptivity-tunable (Anon) keeps the standard state but adapts its own coefficients, so the amount of adaptivity is a hyperparameter rather than a fixed design choice.

• Evolutionary and population (EvE, EA4LLM) do not follow a gradient at all: they sample perturbations or maintain a population of models, which is attractive when gradients are noisy, unavailable, or when many hyperparameters must be tuned at once, and expensive when the model is large.

Table [tab:opt-tut-families] states the same comparison in the two currencies that decide a real training run: what is stored per parameter and what each step costs. Figure 83 turns it into a fingerprint: reading down a row tells you which ingredients a family uses, reading across a column tells you which families share one.

Algorithm [alg:opt-tut-generic] makes the comparison operational: the families differ only in which slots they fill. Reading it side by side with Table [tab:opt-tut-knobs] is enough to place any method from the rest of the appendix without memorising its algebra.

Reading a loss curve

The practical value of the five ideas is diagnostic. A single loss curve already says which of them is wrong, and the answer is usually visible within a few hundred steps (Figure 84).

Five rules of thumb close the tutorial, all of them consequences of the figures above rather than of any particular method.

1. Start from a first-order baseline with both moments (AdamW). It is the cheapest configuration that survives heterogeneous curvature.

2. If the optimizer state does not fit in memory, change family before changing hyper parameters: the memory-efficient and low-rank families exist for exactly that constraint.

3. If tuning the schedule consumes the experiment budget, the schedule-free and parameter-free families remove the knob instead of tuning it.

4. If most parameters live in matrices (attention and feed-forward projections), geometric methods exploit a structure that element-wise methods ignore; if the model is small and the search is over many configurations, population methods become affordable.

5. Change one slot at a time. Because the families are orthogonal compositions, a run that changes two of them cannot attribute its result.

The rest of this appendix is the reference behind these statements: one entry per method, each stating its update rule, its pseudocode, its memory profile, and the trade it makes.

Standard Baseline and Adaptive Optimizers

Vanilla SGD (Mini-Batch Gradient Descent).

The plainest first-order method: update the weights with the gradient of the current mini-batch instead of the full batch . On redundant datasets the mini-batch gradient is nearly as informative as the full one, and computing it with matrix–matrix products makes it markedly cheaper than online updates, especially on GPUs. At iteration \(t\) with mini-batch \(B\): \[\begin{aligned} g_t &= \frac{1}{|B|} \sum_{i \in B} \nabla_\theta E_i(\theta_t) \\ \theta_{t+1} &= \theta_t - \eta\, g_t \end{aligned}\] where \(\eta\) is the only hyperparameter of the method. The deck’s own tuning recipe is heuristic: reduce the rate when the error oscillates or diverges, increase it when the error falls consistently but slowly, and turn it down towards the end of training . Momentum is a separate extension (next paragraph); this entry is deliberately the zero-state baseline of the catalog.

Cons.

• A single global rate cannot suit heterogeneous curvature: the method zig-zags across narrow ravines and stalls along flat directions.

• Strong dependence on hand-tuned learning-rate schedules; no per-parameter adaptivity.

Pros.

• Zero optimizer state: the memory floor of the catalog and the cheapest step.

• Fully deterministic given the batch order; the reference point against which every other entry is measured.

SGD with Momentum.

The classical update accumulates a momentum buffer and then applies a fixed learning rate. At each iteration \(t\), the algorithm maintains an exponential moving average \(m_t\) of past gradients: \[\begin{aligned} m_t &= \beta m_{t-1} + g_t \\ \theta_{t+1} &= \theta_t - \eta m_t \end{aligned}\] where \(g_t = \nabla f(\theta_t)\), \(\beta \in [0.8, 0.99]\) controls the momentum decay, and \(\eta\) is the fixed learning rate. The momentum term acts as velocity: it accelerates movement in consistent directions while dampening oscillations in variable directions.

Cons.

• Poor robustness to heterogeneous curvature: a single global rate cannot suit every parameter group.

• Strong dependence on carefully engineered learning-rate schedules.

Pros.

• Low memory overhead: a single momentum buffer per parameter.

• Strong generalization when tuned carefully, and a simple update to implement.

RMSProp.

RMSProp is the mini-batch-safe version of rprop: instead of dividing by the size of the current gradient, force the divisor to be nearly identical for adjacent mini-batches by keeping a moving average of the squared gradient for each weight : \[\begin{aligned} v_t &= 0.9\, v_{t-1} + 0.1\, g_t^2 \\ \theta_{t+1} &= \theta_t - \eta\, \frac{g_t}{\sqrt{v_t} + \epsilon} \end{aligned}\] where \(g_t = \partial E/\partial w(t)\), the \(0.9/0.1\) split is the exponential decay of the slides (exposed as the alpha hyperparameter), and \(\epsilon\) is a numerical guard added by the implementation (the slides define none). No bias correction and no first moment appear in the original method ; combining it with momentum is explicitly left open by the deck (“momentum does not help as much as it normally does”). The running average is the only state buffer and the direct predecessor of the adaptive family: AdaGrad’s diagonal idea  made mini-batch-safe, later extended into Adam.

Cons.

• No bias correction: the divisor is inflated early in training, when \(v_t\) is still a poor estimate.

• Still tuning-sensitive (rate, decay, epsilon) and largely superseded by Adam for transformers.

Pros.

• One buffer per parameter and an \(\mathcal{O}(n)\) step — the cheapest adaptive method in the catalog.

• Per-coordinate scaling keeps the step robust to heterogeneous gradient magnitudes across layers.

Adam and AdamW.

Adam tracks exponential moving averages of both first moment (mean) and second moment (uncentered variance) to achieve element-wise adaptive learning rates. The update rule is: \[\begin{aligned} m_t &= \beta_1 m_{t-1} + (1-\beta_1) g_t \\ v_t &= \beta_2 v_{t-1} + (1-\beta_2) g_t^2 \\ \hat{m}_t &= \frac{m_t}{1-\beta_1^t}, \quad \hat{v}_t = \frac{v_t}{1-\beta_2^t} \quad \text{(bias correction)} \\ \theta_{t+1} &= \theta_t - \eta \left(\frac{\hat{m}_t}{\sqrt{\hat{v}_t}+\epsilon}\right) \end{aligned}\] AdamW introduces a crucial modification: weight decay is applied directly to parameters (decoupled from gradients): \(\theta_t \gets \theta_t(1-\eta\lambda)\) rather than adding \(\lambda\theta_t\) to the gradient. This decoupling prevents the adaptive scaling from interfering with regularization strength.

Cons.

• Stores first and second moment tensors for every parameter, doubling optimizer-state memory relative to SGD.

• Element-wise adaptive scaling can interfere with regularization unless weight decay is decoupled.

Pros.

• Practical baseline for transformer fine-tuning: fast convergence and forgiving to hyperparameter variation.

• Decoupled weight decay keeps the regularization strength independent of the adaptive scaling.

RAdam (Rectified Adam).

RAdam addresses Adam’s instability in early training by dynamically rectifying the adaptive learning rate. The key observation is that the second moment \(v_t\) has very high variance in the first few steps, causing unreliable adaptive scaling. RAdam computes the effective simple moving average (SMA) window length: \[\rho_t = \rho_{\infty} - \frac{2t\beta_2^t}{1-\beta_2^t}, \quad \text{where} \quad \rho_{\infty} = \frac{2}{1-\beta_2}-1\] When \(\rho_t > 4\) (sufficient samples for variance estimation), RAdam applies adaptive scaling with a rectification term: \[r_t = \sqrt{\frac{(\rho_t-4)(\rho_t-2)\rho_{\infty}}{(\rho_{\infty}-4)(\rho_{\infty}-2)\rho_t}}, \quad \theta_{t+1} = \theta_t - \eta r_t \frac{\hat{m}_t}{\sqrt{\hat{v}_t}+\epsilon}\] When \(\rho_t \leq 4\), RAdam falls back to SGD with momentum: \(\theta_{t+1} = \theta_t - \eta\hat{m}_t\).

Cons.

• Adds rectification bookkeeping on top of Adam’s moment state.

• The benefit shrinks when a well-tuned warmup schedule is already in use.

Pros.

• Eliminates the need for manual learning-rate warmup by rectifying the adaptive rate.

• Falls back gracefully to momentum SGD while variance estimates are unreliable, improving robustness.

LaProp.

LaProp targets an unnecessary coupling in Adam rather than a new preconditioner . Writing the adaptive family as \(\theta_{t+1} = \theta_t - \lambda_t G_t \nabla\ell(\theta_t) + \mu_t K_t(\theta_t - \theta_{t-1})\), Adam’s parameter-space momentum is multiplied by \(K_t = H_t^{-1}H_{t-1}\), i.e. the previous motion is rescaled by the current preconditioner. A single pathologically large gradient—noise, an out-of-distribution batch, or an explosion—can therefore erase the accumulated momentum. LaProp simply sets \(K_t = I\) and instead accumulates the momentum of the preconditioned gradient: \[\begin{aligned} n_t &= \nu n_{t-1} + (1-\nu) g_t^2 \\ m_t &= \mu m_{t-1} + (1-\mu)\,\frac{g_t}{\sqrt{n_t/(1-\nu^t)} + \epsilon} \\ \theta_{t+1} &= \theta_t - \lambda\,\frac{m_t}{1-\mu^t} \end{aligned}\] with \(g_t = \nabla\ell(\theta_t)\) and \(0 \le \mu, \nu < 1\). The second moment \(n_t\) is Adam’s uncentered variance; the first moment \(m_t\) averages the already-normalized direction, so the parameter-space momentum stays additive like SGD. The influence of any single gradient on the accumulated momentum is bounded by \((1-\mu)/\sqrt{1-\nu}\). LaProp has exactly the state (two buffers), per-step cost and hyperparameter count of Adam, and Adam’s settings carry over unchanged; \(\nu \to 0\) recovers signed-gradient methods and \(\mu = 0\) recovers RMSProp with bias correction (the paper describes it as RMSProp with momentum plus both correction factors).

Cons.

• Does not dominate Adam on every task; its gains are concentrated in noisy or unstable optimization regimes.

• Couples the momentum decay to the preconditioned scale, so the effective step size at initialization differs from Adam until the bias correction settles.

Pros.

• Adam-identical memory (\(m\), \(v\)), per-step cost and hyperparameter count; existing Adam settings transfer directly.

• Keeps parameter-space momentum additive, avoiding the \(H_t^{-1}H_{t-1}\) reweighting that lets one bad gradient wipe the accumulated direction.

• Interpolates continuously between adaptive and signed-gradient updates through \(\nu\).

Advanced Momentum and Variance Reduction (2024–2025)

Adan (Adaptive Nesterov Momentum).

Adan reformulates Nesterov acceleration without the extra gradient computation required by classical Nesterov SGD. The algorithm maintains three momentum buffers: \[\begin{aligned} m_t &= (1-\beta_1)m_{t-1} + \beta_1 g_t \quad\text{(first moment)} \\ v_t &= (1-\beta_2)v_{t-1} + \beta_2(g_t - g_{t-1}) \quad\text{(velocity/gradient difference)} \\ n_t &= (1-\beta_3)n_{t-1} + \beta_3[g_t + (1-\beta_1)(g_t-g_{t-1})]^2 \quad\text{(Nesterov second moment)} \end{aligned}\] The Nesterov Momentum Estimation (NME) term \(\bar{g}_t = g_t + (1-\beta_1)(g_t-g_{t-1})\) estimates the gradient at a future position without evaluating it. The update combines momentum with velocity: \[\bar{m}_t = m_t + (1-\beta_1)v_t, \quad \theta_{t+1} = \theta_t - \eta\frac{\bar{m}_t}{\sqrt{n_t}+\epsilon}\]

Cons.

• Three momentum-like buffers increase memory overhead relative to Adam.

• Additional gradient-difference bookkeeping per step.

Pros.

• Nesterov-style acceleration without an extra gradient evaluation.

• Fast convergence across diverse architectures (CNNs, GANs, transformers).

ADOPT (Adam with Optimal Pruned Tuning).

ADOPT fixes a fundamental theoretical issue in Adam: the gradient appears in both the first moment and second moment estimates, creating circularity. ADOPT decouples by using the previous step’s second moment for the denominator: \[\begin{aligned} v_t &= \beta_2 v_{t-1} + (1-\beta_2)g_t^2 \\ m_t &= \beta_1 m_{t-1} + (1-\beta_1)\frac{g_t}{\sqrt{v_{t-1}}+\epsilon} \quad\text{(uses $v_{t-1}$, not $v_t$)} \\ \theta_{t+1} &= \theta_t - \eta m_t \end{aligned}\]

Cons.

• Requires keeping the previous second moment as an explicit denominator.

• Empirical gains are architecture- and tuning-dependent.

Pros.

• Removes Adam’s gradient circularity and attains the optimal convergence rate \(O(1/\sqrt{T})\) with any \(\beta_2 \in (0,1)\), without bounded-noise assumptions.

• Drop-in Adam replacement with comparable or superior empirical performance across vision, NLP, RL, and generative modeling.

AdEMAMix (Exponential Moving Average Mixture).

AdEMAMix replaces Adam’s single EMA of gradients with a mixture of two EMAs: one fast-decaying and one slow-decaying. This addresses the observation that gradients remain informative over tens of thousands of steps, not just hundreds: \[\begin{aligned} m_{1,t} &= \beta_1 m_{1,t-1} + (1-\beta_1)g_t \quad\text{(fast EMA, $\beta_1 \approx 0.9$)} \\ m_{2,t} &= \beta_3 m_{2,t-1} + (1-\beta_3)g_t \quad\text{(slow EMA, $\beta_3 \approx 0.9999$)} \\ v_t &= \beta_2 v_{t-1} + (1-\beta_2)g_t^2 \\ m_t &= m_{1,t} + \alpha_t m_{2,t} \quad\text{(mixture with scheduled weight $\alpha_t$)} \\ \theta_{t+1} &= \theta_t - \eta\frac{m_t}{\sqrt{v_t}+\epsilon} \end{aligned}\] The mixture weight \(\alpha_t\) typically increases during training, allowing the fast EMA to provide immediate adaptation while the slow EMA accumulates long-range gradient correlations.

Cons.

• Two gradient EMAs plus the second moment roughly triple Adam’s optimizer state.

• The mixture schedule \(\alpha_t\) adds a hyperparameter to tune.

Pros.

• Exploits gradient information that remains informative over tens of thousands of steps.

• A 1.3B model on 101B tokens matches AdamW on 197B tokens (95% data-efficiency gain), suggesting reduced forgetting.

MARS (Make Adaptive learning Rates Shine).

MARS combines preconditioned gradient methods (e.g., AdamW) with variance reduction via scaled stochastic recursive momentum. The core innovation is a variance-reduced gradient estimate: \[\begin{aligned} c_t &= g_t + \gamma(c_{t-1} - g_{t-1}) \quad\text{(SVRG-style recursive estimate)} \\ m_t &= \beta_1 m_{t-1} + (1-\beta_1)c_t \\ v_t &= \beta_2 v_{t-1} + (1-\beta_2)c_t^2 \\ \theta_{t+1} &= \theta_t - \eta\frac{m_t}{\sqrt{v_t}+\epsilon} \end{aligned}\] where \(\gamma \in [0.01, 0.1]\) controls how much historical gradient information is retained. The variance-reduced gradient \(c_t\) acts as an implicit noise filter, amplifying consistent signal directions and dampening contradictory noise.

Cons.

• Stores the variance-reduced gradient and the previous gradient in addition to Adam’s moments.

• The variance-reduction coefficient \(\gamma \in [0.01, 0.1]\) must be tuned.

Pros.

• Acts as an implicit noise filter, amplifying consistent directions and dampening contradictory noise.

• MARS-AdamW consistently outperforms bare AdamW by significant margins on GPT-2 pretraining.

Cautious Optimizers (Cautious AdamW, Cautious Lion).

The Cautious framework applies a one-line modification to any momentum-based optimizer: mask the update so that only dimensions where momentum and gradient directions agree are applied: \[\begin{aligned} \text{mask}_i &= \begin{cases} 1 & \text{if } m_t[i] \cdot g_t[i] > 0 \quad\text{(agreement)} \\ 0 & \text{otherwise} \end{cases} \\ u_t &= \left(\frac{m_t}{\sqrt{v_t}+\epsilon}\right) \odot \text{mask} \quad\text{(element-wise masking)} \\ \theta_{t+1} &= \theta_t - \eta u_t \end{aligned}\] The intuition: momentum \(m_t\) estimates the gradient direction from history; the current gradient \(g_t\) is instantaneous signal. When both agree, the optimizer is confident and should update aggressively. When they disagree, the optimizer is conflicted—the historical trend points toward a region that may have been good before, but current evidence contradicts it. By masking conflicted dimensions, Cautious becomes more conservative and avoids corrupted steps.

Cons.

• Masking can discard useful updates when momentum and gradient temporarily disagree.

• Requires a base momentum optimizer plus the element-wise mask.

Pros.

• A one-line modification applicable to any momentum-based optimizer.

• Up to 1.47\(\times\) speedup on Llama and MAE pretraining while preserving convergence guarantees.

Large-Batch, Memory-Efficient, and Parameter-Free Optimizers

LAMB (Layer-wise Adaptive Moments optimizer for Batch training).

LAMB extends Adam with layer-wise adaptive rate scaling (inspired by LARS), enabling stable training with extreme batch sizes (e.g., 64K on BERT): \[\begin{aligned} m_t^L &= \beta_1 m_{t-1}^L + (1-\beta_1)g_t^L \quad\text{(layer-wise first moment)} \\ v_t^L &= \beta_2 v_{t-1}^L + (1-\beta_2)(g_t^L)^2 \\ u_{\text{adam}}^L &= \frac{m_t^L}{\sqrt{v_t^L}+\epsilon} + \lambda\theta_{t-1}^L \quad\text{(base adaptive step with decay)} \\ \phi^L &= \frac{\|\theta_{t-1}^L\|_2}{\|u_{\text{adam}}^L\|_2} \quad\text{(trust ratio: layer normalization)} \\ \theta_t^L &= \theta_{t-1}^L - \eta \cdot \phi^L \cdot u_{\text{adam}}^L \end{aligned}\] The trust ratio \(\phi^L = \|\theta^L\|_2 / \|u_{\text{adam}}^L\|_2\) normalizes the effective update step relative to weight magnitude. In very large batches, stochastic gradient noise can exceed signal, destabilizing layer-wise learning rates. By scaling updates proportionally to weight norms, LAMB ensures relative changes rather than absolute ones, preventing small confident updates in large vectors from dominating.

Cons.

• Layer-wise norm bookkeeping and trust-ratio computation add per-step overhead.

• The trust ratio is unstable when layer or update norms are close to zero.

Pros.

• Enables stable training with extreme batch sizes (e.g., 64K on BERT).

• Preserves small-batch generalization benefits while scaling the batch.

Schedule-Free AdamW.

Schedule-free methods remove explicit scheduler design from the optimization recipe. The algorithm maintains two parameter streams: \(z_t\) (exploration) and \(x_t\) (smooth average): \[\begin{aligned} v_t &= \beta_2 v_{t-1} + (1-\beta_2)g_t^2 \quad\text{(variance)} \\ z_t &= z_{t-1}(1-\eta\lambda) - \eta\frac{g_t}{\sqrt{v_t}+\epsilon} \quad\text{(adaptive step with decay)} \\ c_t &= \frac{1}{t+1} \quad\text{(averaging coefficient, decreases over time)} \\ x_t &= (1-c_t)x_{t-1} + c_t z_t \quad\text{(iterate averaging)} \\ y_t &= (1-\beta)z_t + \beta x_t \quad\text{(interpolation for evaluation point)} \end{aligned}\] The averaging coefficient \(c_t = 1/(t+1)\) implements an implicit learning-rate schedule without explicitly specifying total steps \(T\). This framework unifies scheduling and iterate averaging: the algorithm explores via \(z_t\) while accumulating stable direction via \(x_t\). The evaluation point \(y_t\) (used for gradient computation) interpolates between exploration and stability.

Cons.

• Maintains two parameter streams (\(z\) and \(x\)) and requires explicit handling of the evaluation point.

• Behavior departs from scheduled baselines, which complicates comparisons.

Pros.

• Removes the explicit learning-rate schedule and its total-step hyperparameter.

• State-of-the-art convergence across convex optimization, large-scale deep learning, and reinforcement learning.

Adafactor.

Adafactor reduces optimizer memory by factorizing second-moment statistics for matrix-shaped parameters, storing only row and column accumulators rather than dense variance states. For a gradient matrix \(G_t \in \mathbb{R}^{m \times n}\): \[\begin{aligned} R_t &= \beta_2 R_{t-1} + (1-\beta_2)(G_t^2) \mathbf{1}_n^T \quad\text{(row variance)} \\ C_t &= \beta_2 C_{t-1} + (1-\beta_2) \mathbf{1}_m^\top (G_t^2) \quad\text{(column variance)} \\ \hat{V}_t &= \frac{R_t C_t}{\mathbf{1}_n^\top R_t} \quad\text{(reconstructed variance via outer product)} \\ U_t &= \frac{G_t}{\sqrt{\hat{V}_t}+\epsilon} \quad\text{(normalized adaptive step)} \\ \hat{U}_t &= \frac{U_t}{\max(1, \text{RMS}(U_t))} \quad\text{(stability clipping)} \end{aligned}\]

Cons.

• Factorized variance is an approximation and can be unstable on some tasks.

• Reduced optimization expressiveness relative to full Adam.

Pros.

• Memory drops from \(O(n_{\text{params}})\) to \(O(\sqrt{n_{\text{params}}})\) for matrix parameters.

• Most attractive when VRAM is dominated by optimizer state rather than activations.

GaLore (Gradient Low-Rank Projection).

GaLore projects 2D gradients into a low-rank subspace before optimization, reducing optimizer-state memory while maintaining adaptive step scaling. For a 2D parameter matrix \(W \in \mathbb{R}^{m \times n}\): \[\begin{aligned} U, S, V &= \text{SVD}(G_t) \quad\text{(compute singular decomposition)} \\ P &\in \mathbb{R}^{n \times r} \quad\text{or} \quad P^\top \in \mathbb{R}^{r \times m} \quad\text{(select top-$r$ singular vectors)} \\ G_{\text{low}} &= P^\top G_t \quad\text{(project to low-rank space)} \\ \Delta_{\text{low}} &= \text{Adam}(G_{\text{low}}) \quad\text{(optimize in compressed space)} \\ \Delta &= P \Delta_{\text{low}} \quad\text{(reconstruct in original space)} \end{aligned}\] By projecting into a rank-\(r\) subspace (typically \(r \ll \min(m,n)\)), optimizer state is reduced from \(O(mn)\) to \(O(r(m+n))\) for 2D parameters. This complementary memory-saving approach is especially relevant when the model is too large for full-rank optimizer state.

Cons.

• Periodic SVD adds compute and introduces a projection-refresh interval to tune.

• Low-rank projection can discard gradient information.

Pros.

• Optimizer state reduced from \(O(mn)\) to \(O(r(m+n))\) for 2D parameters.

• Works synergistically with other techniques and has shown strong results on billion-parameter models.

APOLLO, APOLLO-Mini, and Q-APOLLO.

The APOLLO family begins from the observation that AdamW’s element-wise denominator can be coarsened into a structured learning-rate update. Rather than storing dense moments for every parameter entry, APOLLO projects a matrix gradient \(G_t \in \mathbb{R}^{m \times n}\) into a compact random subspace and tracks Adam-style moments there: \[\begin{aligned} R_t &= P_t G_t \quad\text{or}\quad R_t = G_t P_t^\top \\ M_t^R &= \beta_1 M_{t-1}^R + (1-\beta_1)R_t \\ V_t^R &= \beta_2 V_{t-1}^R + (1-\beta_2)R_t^2 \\ \widetilde{R}_t &= \frac{M_t^R}{\sqrt{V_t^R}+\epsilon} \end{aligned}\] The projected state is not expanded back into a dense low-rank update as in SVD-based methods. Instead, APOLLO estimates a structured scaling tensor \(S_t\) in the original space. In the standard APOLLO variant, that scaling is channel-wise: each row or channel receives its own norm ratio. In APOLLO-Mini, the scaling is reduced to a single tensor-wise scalar, corresponding to the rank-1 extreme described in the paper. The resulting parameter update is therefore Adam-like in adaptation but much closer to SGD in state cost: \[W_t \gets (1-\eta\lambda)W_{t-1} - \eta\,\alpha\,(G_t \odot S_t)\] where \(\alpha\) is the extra scale factor used to stabilize highly compressed variants.

Cons.

• Structured scaling is a coarser approximation of AdamW’s element-wise denominator.

• Projectors must be refreshed periodically, and very small ranks make the scaling noisy.

Pros.

• Replaces repeated SVD with periodically refreshed Gaussian random projection, so the systems cost is ordinary matrix multiplication rather than spectral decomposition.

• Unusually tolerant to extreme compression: even APOLLO-Mini, which keeps only rank-1 auxiliary state, remains competitive with or better than AdamW while approaching SGD-level memory.

APOLLO-Mini.

APOLLO-Mini is the family member optimized for extreme memory efficiency. The paper motivates it by arguing that, in a rank-1 compact space, channel-wise scaling becomes too noisy, so the update is coarsened further to a single tensor-wise scale. The loss of granularity is partially offset by the explicit scaling factor \(\alpha\) (the paper discusses values such as \(128\) in this regime), giving a useful point on the optimization Pareto frontier: very small state, no SVD overhead, and still strong pre-training behavior.

Cons.

• A single tensor-wise scale loses channel granularity, and rank-1 state can be noisy.

• Depends on the extra scale factor \(\alpha\) (e.g., \(128\)) to remain stable.

Pros.

• Smallest state in the APOLLO family, with no SVD overhead.

• Strong pre-training behavior at near-SGD memory cost.

Q-APOLLO.

The APOLLO paper also stresses that the family combines naturally with quantization for ultra-low-memory training. This repository turns that systems observation into a concrete optimizer variant, q_apollo. In the implementation, APOLLO’s low-rank first and second moments are stored in quantized form together with per-tensor scales and offsets, and are dequantized only when needed for the next step. Q-APOLLO therefore preserves APOLLO’s projected-gradient logic while reducing the precision of the remaining optimizer state.

Cons.

• Quantized moments introduce approximation error and dequantization overhead.

• Requires managing per-tensor scales and offsets.

Pros.

• Reduces the remaining optimizer state to ultra-low precision.

• Most aggressive memory-saving member of the local optimizer stack while preserving APOLLO’s projected-gradient logic.

AdamPM (Pooled-Moment Adam).

AdamPM  is a minimal, model-agnostic modification of Adam: it keeps the element-wise first moment untouched and compresses only the second moment by group-mean pooling of adjacent entries along each row (the default, rows being output features) or each column (the paper’s AdamPM-col). For a weight matrix \(W \in \mathbb{R}^{m \times n}\) and group size \(s\), every row is partitioned into \(q = \lceil n/s \rceil\) groups, the squared gradient is averaged within each group before the exponential moving average, and only the compressed statistic \(\bar{V}_t \in \mathbb{R}^{m \times q}\) is persisted: \[\begin{aligned} M_t &= \beta_1 M_{t-1} + (1-\beta_1)\, G_t && \text{(first moment, element-wise, Eq.~1)} \\ \bar{S}_t &= \frac{1}{s}\sum_{k=1}^{s} \mathrm{reshape}(G_t \odot G_t,\, m, q, s)[:,:,k] && \text{(within-group means, Eq.~2)} \\ \bar{V}_t &= \beta_2 \bar{V}_{t-1} + (1-\beta_2)\, \bar{S}_t && \text{(compressed second-moment EMA, Eq.~3)} \\ W_t &= W_{t-1} - \eta_t \, \frac{M_t/(1-\beta_1^t)} {\sqrt{\hat{\bar{V}}_t \otimes \mathbf{1}_{1\times s}} + \epsilon}, \qquad \hat{\bar{V}}_t = \frac{\bar{V}_t}{1-\beta_2^t} && \text{(update, Eq.~4)} \end{aligned}\] When \(n\) is not divisible by \(s\) the final group is padded with its own mean, so its statistic averages only the entries it actually contains; \(s{=}1\) recovers Adam exactly, and \(s \ge n\) collapses each row (or column) to one scalar, the paper’s AdamPM-inf limit. AdamPM pools the second moments of all weight matrices—embeddings, the LM head, attention and MLP weights—and treats one-dimensional parameters such as RMSNorm scales as single-row or single-column matrices, so no parameter group is left on a fallback path. The persistent state is \(mn(1+1/s)\) elements against Adam’s \(2mn\).

Empirically, AdamPM matches a carefully tuned Adam on LLaMA pre-training from 60M to 1.3B parameters while cutting second-moment memory by 87.5% (\(s{=}8\)) to 93.75% (\(s{=}16\)), with the same pattern on GLUE fine-tuning of RoBERTa-base and on TinyLLaVA-1B supervised fine-tuning; group sizes between 2 and 8 are reported as robust, and three or five—not only powers of two—work because the incomplete final group is averaged over its own entries. Theorem 3.1 gives a stationarity bound independent of the group size \(s\), and under a relative-noise condition AdamPM attains the same \(\mathcal{O}(\log T/\sqrt{T})\) rate as Adam. Algorithm 1 of the paper defines no weight decay; this implementation applies the repository’s decoupled wd_ decay before the adaptive update, and the shared per-group betas/eps suffixes parameterize the bias corrections and the denominator.

Cons.

• The diagonal preconditioner becomes coarser as \(s\) grows; the paper’s AdamPM-inf limit shows a measurable perplexity loss at 350M and 1.3B.

• The saving covers optimizer state only: the first moment stays full-size, so optimizer memory is \(1+1/s\) of Adam’s, not \(1/s\).

• Pooling is one-dimensional (rows or columns); two-dimensional groups and selective per-layer pooling remain unexplored (paper limitations).

Pros.

• Roughly halves optimizer-state memory at the default \(s{=}8\) (second moment at \(1/8\)) with Adam-level perplexity through 1.3B parameters and negligible per-step overhead—no SVD, no projection, no extra state.

• Inherits Adam’s learning rate and decay rates directly, so existing Adam presets transfer without retuning and \(\beta_1{=}0.9\), \(\beta_2{=}0.95\) match the repository defaults.

• Model-agnostic and applies uniformly to every parameter group; the preconditioned updates grow more similar to Adam’s as model scale grows.

Anon (Adaptivity Non-restricted Optimizer with Novel convergence technique).

Anon introduces continuously tunable adaptivity via a single parameter \(\gamma \in \mathbb{R}\), bridging the gap between SGD-like (\(\gamma=0\)) and Adam-like (\(\gamma=1\)) behaviors, and extrapolating beyond both. The core innovation is the Incremental Delay Update (IDU) mechanism, which enables stable convergence across the entire real-valued adaptivity spectrum:

\[\begin{aligned} m_t &= \beta_1 m_{t-1} + (1-\beta_1) g_t \quad\text{(first moment)}\\ v_t &= \beta_2 v_{t-1} + (1-\beta_2) g_t^2 \quad\text{(second moment)}\\ \bar{v}_t &= \frac{v_t}{1-\beta_2^{\max(t,1)}} + \epsilon \quad\text{(bias-corrected second moment)}\\ \text{fixed\_\allowbreak lr}_k &= \begin{cases} \bar{v}_k^{-\gamma/2} & \text{if } k = 0 \\ \sqrt{2 / \left(\frac{1}{\text{fixed\_\allowbreak lr}_{k-1}^2} + \bar{v}_k^\gamma\right)} & \text{if } k > 0 \end{cases}\\ \theta_{t+1} &= \theta_t - \frac{\eta}{1-\beta_1^t} \cdot m_t \odot \text{fixed\_\allowbreak lr}_{\lfloor\log_2 t\rfloor} \end{aligned}\]

Unlike standard adaptive optimizers, IDU updates the preconditioner only at powers of two (\(t = 2^k\)), accumulating gradient statistics between updates and aggregating them recursively. This soft, multi-scale accumulator avoids the hard max-tracking of AMSGrad while provably stabilizing convergence for any \(\gamma \in \mathbb{R}\). The parameter \(\gamma\) directly controls adaptivity: \(\gamma > 1\) accelerates escape from saddle points (favorable for complex architectures), \(\gamma = 1\) recovers Adam-like behavior, \(\gamma = 0\) yields SGD-like scaling, and \(\gamma < 0\) biases toward flatter minima (beneficial for classical architectures like CNNs).

Cons.

• Maintains three state buffers (\(m\), \(v\), fixed_lr) and refreshes the preconditioner on a logarithmic schedule.

• The adaptivity parameter \(\gamma\) must be selected per architecture.

Pros.

• Same \(\mathcal{O}(n)\) per-step complexity class as AdamW, with continuously tunable adaptivity.

• IDU provably attains \(\mathcal{O}(\sqrt{T})\) regret for convex problems and \(\mathcal{O}(\ln T / \sqrt{T})\) for non-convex settings.

Prodigy (Approximating the Distance Estimate).

Prodigy adapts the effective step scale through a running distance-like statistic, eliminating the need for explicit learning-rate tuning. The algorithm maintains a cumulative distance estimate: \[\begin{aligned} u_t &= \frac{g_t}{\sqrt{v_t}+\epsilon} \quad\text{(unnormalized adaptive step)} \\ s_t &= s_{t-1} + \langle u_t, \theta_t - \theta_0 \rangle \quad\text{(cumulative signed distance)} \\ d_t &= \max(d_{t-1}, d_0 + d_{\text{coef}} \cdot s_t) \quad\text{(distance estimate with lower bound)} \\ \theta_{t+1} &= \theta_t - \eta \cdot d_t \cdot u_t \end{aligned}\] where \(d_0\) is an initial bound and \(d_{\text{coef}}\) controls how aggressively the estimate adapts. The distance estimate \(d_t\) captures the combined magnitude of past gradients weighted by actual parameter displacement, implementing an implicit effective learning rate.

Cons.

• Introduces \(d_0\) and \(d_{\text{coef}}\) hyperparameters for the distance estimate.

• Adds an inner-product bookkeeping step per iteration.

Pros.

• Removes explicit learning-rate tuning and scheduling.

• Robust convergence across problem scales and batch sizes.

Second-Order, Geometric, and Orthogonality Optimizers

Shampoo (Matrix Preconditioning via Kronecker Products).

Shampoo is a structured second-order method that computes matrix preconditioners from Kronecker-structured outer-product statistics. For a 2D parameter matrix \(W \in \mathbb{R}^{m \times n}\): \[\begin{aligned} L_t &= L_{t-1} + G_t G_t^\top \quad\text{(left/row Gram matrix, $m \times m$)} \\ R_t &= R_{t-1} + G_t^\top G_t \quad\text{(right/column Gram matrix, $n \times n$)} \\ L_t^{-1/4} &= Q_L (\Lambda_L)^{-1/4} Q_L^\top \quad\text{(via eigendecomposition)} \\ R_t^{-1/4} &= Q_R (\Lambda_R)^{-1/4} Q_R^\top \\ \Delta W_t &= L_t^{-1/4} G_t R_t^{-1/4} \quad\text{(preconditioned gradient)} \end{aligned}\] Shampoo approximates the full-matrix Adagrad preconditioner \(H^{-1/2}\) (where \(H\) is the Hessian) using Kronecker-factored structure. Instead of storing a \((mn) \times (mn)\) preconditioner, Shampoo maintains two smaller matrices (\(m \times m\) and \(n \times n\)), capturing cross-parameter correlations within and across groups. Eigendecompositions are performed periodically (e.g., every \(K=10\) steps) to amortize cost.

Cons.

• Periodic \(O(m^3 + n^3)\) eigendecompositions and higher per-step compute.

• Assumes 2D parameter structure; embeddings and biases need separate handling.

Pros.

• Kronecker-factored second-order preconditioning without a dense \((mn) \times (mn)\) matrix.

• Proven effective at production scale and well suited to high-rank transformer matrices, with improved conditioning and accelerated convergence.

SOAP (Shampoo with Adam in eigenbasis).

SOAP is a simplified variant of Shampoo that decouples preconditioning from momentum tracking. Instead of complex matrix algebra on preconditioned gradients, SOAP runs standard Adam in the eigenbasis of Shampoo’s preconditioners: \[\begin{aligned} L_t &= L_{t-1} + G_t G_t^\top, \quad R_t = R_{t-1} + G_t^\top G_t \quad\text{(Gram accumulation)} \\ Q_L, \Lambda_L &= \text{eigh}(L_t), \quad Q_R, \Lambda_R = \text{eigh}(R_t) \quad\text{(periodic eigendecomposition)} \\ G_t^{\text{rot}} &= Q_L^\top G_t Q_R \quad\text{(rotate gradient into eigenbasis)} \\ m_t^{\text{rot}} &= \beta_1 m_{t-1}^{\text{rot}} + (1-\beta_1) G_t^{\text{rot}} \\ v_t^{\text{rot}} &= \beta_2 v_{t-1}^{\text{rot}} + (1-\beta_2)(G_t^{\text{rot}})^2 \quad\text{(Adam in rotated space)} \\ U_t^{\text{rot}} &= \frac{m_t^{\text{rot}}}{\sqrt{v_t^{\text{rot}}}+\epsilon}, \quad U_t = Q_L U_t^{\text{rot}} Q_R^\top \quad\text{(rotate back)} \end{aligned}\] The key insight: all momentum tracking occurs in the well-conditioned eigenbasis, simplifying numerical stability and theoretical analysis.

Cons.

• Inherits Shampoo’s periodic eigendecomposition cost.

• Rotation and inverse rotation add matrix multiplications per step.

Pros.

• Combines Shampoo’s explicit curvature structure with Adam’s proven momentum mechanics.

• Often cleaner and more stable than full Shampoo, with eigendecompositions amortized every \(K\) steps.

Lion (EvoLved Sign Momentum).

Lion achieves minimal memory overhead by using sign-based updates instead of full adaptive scaling: \[\begin{aligned} c_t &= \beta_1 m_t + (1-\beta_1) g_t \quad\text{(momentum input)} \\ \theta_{t+1} &= \theta_t - \eta \left(\text{sign}(c_t) + \lambda\theta_t\right) \quad\text{(sign operator update)} \\ m_{t+1} &= \beta_2 m_t + (1-\beta_2) g_t \quad\text{(momentum accumulation)} \end{aligned}\] All element-wise scaling operations are replaced with sign(), which outputs \(\{-1, 0, +1\}\). This dramatically reduces memory compared to Adam-like methods: Lion stores only the momentum buffer, no second moment variance.

Cons.

• Sign-based updates sacrifice the element-wise learning-rate adaptation that makes Adam effective.

• Can be less robust than Adam-family methods on tasks that need adaptive scaling.

Pros.

• Stores only the momentum buffer, with no second moment variance.

• High-throughput, low-memory option for large-batch regimes where activation memory dominates.

Sophia (Second-Order Hessian Information with Optimized Approximation).

Sophia uses diagonal Hessian estimates for curvature-aware scaling without the memory and compute cost of dense second-order methods. \[\begin{aligned} m_t &= \beta_1 m_{t-1} + (1-\beta_1) g_t \quad\text{(first moment)} \\ h_t &= \beta_2 h_{t-(k-1)} + (1-\beta_2) \hat{h}_t \quad\text{(diagonal Hessian, updated every $k$ steps)} \\ \text{clip}(x, C) &= \min(\max(x, -C), C) \quad\text{(element-wise clipping)} \\ \theta_{t+1} &= \theta_t - \eta \cdot \text{clip}\left(\frac{m_t}{\max(\gamma h_t, \epsilon)}, 1\right) \end{aligned}\] where \(\hat{h}_t\) is a diagonal estimate (e.g., Hutchinson-trace estimator) and \(\gamma\) is a scaling factor. The diagonal Hessian \(h_t\) captures local curvature, allowing the optimizer to take smaller steps in sharp directions and larger steps in flat directions. The clipping operation \(\text{clip}(\cdot, 1)\) prevents adaptive steps from exploding.

Cons.

• Hutchinson estimates add periodic compute and noise to the curvature signal.

• The Hessian estimate can become stale between refreshes.

Pros.

• Curvature-aware scaling without the \(O(n^2)\) or \(O(n^3)\) cost of full second-order methods.

• Clipped steps prevent explosions and work particularly well in second-pass fine-tuning.

Muon (Newton-Schulz Orthogonalization).

Muon is an orthogonality-oriented optimizer that reshapes update geometry using Newton-Schulz polynomial iterations for orthogonalization. For a 2D parameter matrix \(W\) with gradient \(G_t\): \[\begin{aligned} X_0 &= \frac{G_t}{\|G_t\|_F + \epsilon} \quad\text{(normalized gradient)} \\ A_k &= X_k X_k^\top \quad\text{(Gramian)} \\ B_k &= b A_k + c A_k^2 \quad\text{(polynomial step, coefficients $b, c$ from Newton-Schulz)} \\ X_{k+1} &= a X_k + B_k X_k \quad\text{(iteration $k=0,1,\ldots,4$)} \\ W_{t+1} &= W_t - \eta X_K \quad\text{(update with orthogonalized direction)} \end{aligned}\] Muon performs 5 Newton-Schulz iterations to produce an approximately orthogonal direction, ensuring updates respect geometric constraints. The rationale: orthogonal updates preserve norms during training, avoiding the norm creep observed in momentum-based methods.

Following the hybrid recipe of , the orthogonalized update is reserved for hidden-layer weight matrices (\(\mathtt{attention}\), \(\mathtt{ffn}\) and other hidden-block groups). Embedding tables, learned positional encodings, output heads and every non-matrix parameter (biases, normalization scales, convolution kernels) are trained with decoupled AdamW instead of being skipped, using the group-level learning rate, weight decay, \(\beta_1, \beta_2\) and \(\epsilon\) settings exposed by the per-group configuration.

Cons.

• Several matrix multiplications per Newton-Schulz iteration; efficient custom kernels are required.

• The orthogonalized construction is defined only for 2D matrices; the remaining parameter classes depend on the AdamW fallback and its per-group learning-rate calibration.

Pros.

• Orthogonalized updates preserve norms and avoid the norm creep of momentum-based methods.

• Promising results on large-scale training with fundamentally better-conditioned update geometry.

Turbo-Muon (AOL-Preconditioned Orthogonalization).

Turbo-Muon adds an almost-orthogonal preconditioning (AOL) step before the Newton-Schulz iterations. The preconditioner \(P = 1.5I - 0.5X_0X_0^\top\) is applied to the normalized gradient, reducing the number of required orthogonalization steps to four while preserving the orthogonalized update geometry. It inherits Muon’s hybrid routing: only hidden-layer 2D matrices are orthogonalized, while embedding tables, learned positional encodings, output heads and non-matrix parameters take decoupled AdamW updates.

Cons.

• Still compute-heavy; AOL adds a preconditioning matrix multiplication.

• Requires the same custom-kernel support as Muon, and shares its dependence on the AdamW fallback for non-hidden parameter classes.

Pros.

• AOL reduces the Newton-Schulz iteration count (typically five to four).

• Retains orthogonalized update geometry at lower iteration cost.

AF-Muon (AdamW-Free Tied-Embedding Muon).

AF-Muon is an AdamW-free extension of Muon that removes the auxiliary AdamW fallback of the Hybrid Muon recipe by giving every parameter class its own norm-constrained geometry. Hidden matrices keep the spectral update with a Nesterov signal; the vocabulary table, which simultaneously receives sparse input-lookup and dense output-softmax gradients, is updated row by row with a support-aware finite-cap linear minimization oracle; one-dimensional auxiliary tensors receive an RMS-normalized update. No parameter class carries second-moment state, so a single first-moment buffer per parameter replaces the AdamW state of the hybrid recipe, which the authors report as roughly 20% less optimizer-state memory at matched training.

For the tied table \(W \in \mathbb{R}^{V \times d}\) with momentum row \(b_i\), the oracle solves the row problem \[\begin{aligned} u_i &= \operatorname{sign}(b_i) \odot \min\left(c\,\mathbf{1},\ \tau_i\,|b_i|\right), \\ \|u_i\|_2^2 &= \min\left(d,\ c^2\,\bigl|\operatorname{supp}(b_i)\bigr|\right), \end{aligned}\] where the per-row scale \(\tau_i\) is obtained by bisection in FP32 over complete-row chunks. Rows whose support cannot fill the row-\(\ell_2\) radius saturate every useful coordinate at the cap \(c\), so sparse lookup rows are handled without inflating absent coordinates. The cap interpolates between two known endpoints: \(c = 1\) recovers the SCION-style Sign update and \(c \geq \sqrt{d}\) reduces to a row-RMS update. The tied step size is \(\eta_{tie} = \eta_M (\rho_{out}/\rho_{hid})\, s / d\); the paper’s main setting is \(c = 3\) and \(s = 0.5\) with \(\rho_{out}/\rho_{hid} = 60\), and its diagnostics report the finite cap active on under 0.3% of coordinates while affecting most rows.

Cons.

• Every parameter must be assigned an explicit role (matrix, tied table, or vector); an untied output head or an extra embedding table is not a tied-table block.

• The tied-table oracle requires a per-row bisection with complete-row chunking; the paper reports roughly 1% median step-time overhead at matched training.

• The tied-table claim is established for tied-vocabulary models; frontier-scale pretraining and universal endpoint superiority are not claimed.

Pros.

• AdamW-free: one first-moment buffer per parameter and no second-moment state, reported as roughly 20% less optimizer-state memory than Hybrid Muon.

• The finite cap preserves within-row magnitudes while bounding coordinate concentration, and recovers Sign (\(c=1\)) and row-RMS (\(c \geq \sqrt{d}\)) as limit cases.

• Improves mean validation loss and perplexity over both Hybrid Muon and a SCION-style Sign endpoint across nine tied-token settings spanning text, image-token, protein, MoE, and encoder-decoder models.

AdaMuon (Adaptive Muon).

AdaMuon is an adaptive extension of Muon that keeps the spectral (polar) update while adding coordinate-wise variance normalization. Muon orthogonalizes the momentum matrix into a direction-only update, which removes anisotropic scaling but assigns the same magnitude to every coordinate; AdaMuon addresses the resulting mismatch between noisy and informative coordinates by tracking the element-wise variance of the orthogonalized direction rather than of the raw gradient or of the momentum. Because the polar factor is globally scale-invariant, the sign transform applied before orthogonalization preserves the update direction while stabilizing its element-wise magnitudes; the paper proves that \(\operatorname{sign}\) is the unique element-wise map that is scale-free, sign-consistent, odd and bounded (Theorem 1). The second-moment estimate therefore sees a stable signal, and its variance-normalized version is rescaled to an RMS of \(0.2\), matching Adam’s empirical update scale so that existing learning-rate schedules are reused without tuning.

For a weight \(W_t \in \mathbb{R}^{n \times m}\), momentum buffer \(M_t\), momentum \(\beta\), Newton–Schulz steps \(T\) and small constant \(\varepsilon\): \[\begin{aligned} M_t &= \beta M_{t-1} + G_t, \\ O_t &= \operatorname{NewtonSchulz}\bigl(\operatorname{sign}(M_t), T\bigr), \\ V_t &= \beta V_{t-1} + (1-\beta)\, O_t \odot O_t, \\ \widehat{O}_t &= O_t \oslash \bigl(\sqrt{V_t} + \varepsilon \mathbf{1}\bigr), \\ \gamma_t &= \frac{0.2\sqrt{mn}}{\|\widehat{O}_t\|_F}, \\ W_{t+1} &= W_t - \eta\bigl(\gamma_t \widehat{O}_t + \lambda W_t\bigr). \end{aligned}\] The Newton–Schulz map uses \(X_0 = \operatorname{sign}(M_t)/\|\operatorname{sign}(M_t)\|_F\) and the cubic iteration \(X_k = aX_{k-1} + b(X_{k-1}X_{k-1}^{\mathsf T})X_{k-1} + c(X_{k-1}X_{k-1}^{\mathsf T})^2X_{k-1}\) with \((a,b,c) = (3.4445, -4.7750, 2.0315)\) and \(T = 5\). No bias correction is applied to either moment: a constant factor in \(V_t\) is cancelled exactly by \(\gamma_t \propto \sqrt{c}\), and a constant factor in \(M_t\) is removed by the polar map, \(\operatorname{polar}(cM) = \operatorname{polar}(M)\).

Cons.

• One additional second-moment buffer per matrix parameter (two buffers versus Muon’s single first-moment buffer), plus per-step element-wise EMA, square-root and normalization work on the update.

• The update is defined for two-dimensional parameters; one-dimensional parameters (biases, norm scales) require an auxiliary optimizer in the paper’s recipe.

• The RMS target (\(0.2\)) is a fixed constant rather than a learned quantity, so the method assumes Adam-like update scales when reusing existing schedules.

Pros.

• No new hyperparameters: the second-moment EMA reuses Muon’s momentum coefficient \(\beta\), and the paper reports low sensitivity across \(\beta \in [0.8, 0.99]\).

• The sign transform makes the variance estimate stable enough to normalize, and the RMS alignment cancels any constant bias in \(V_t\), so no bias correction is needed for either moment.

• RMS alignment matches Adam’s update scale, allowing direct reuse of learning-rate schedules; the authors report up to roughly 40% training-efficiency gains over AdamW and consistent wins over Muon on GPT-2 (125M–1.5B) and Qwen2.5 (1.5B, 7B).

Dion3 (Fractional Gram-Newton-Schulz Orthogonalization).

Dion3  is a revision of Muon that reduces the cost of the orthogonalization step at the algorithmic level. It replaces standard Newton–Schulz with Gram Newton–Schulz, a mathematically equivalent iteration that runs on the small symmetric Gram matrix \(XX^{\mathsf T}\), and it orthogonalizes only a fraction \(f \in (0,1]\) of the momentum rows at each step, updating only those rows of the weight matrix. Error feedback decays only the selected rows of the momentum buffer, so rows that were not selected keep accumulating and are eventually chosen. At \(f = 1\) the update reduces to Muon in exact arithmetic (or to NorMuon with the optional row-wise normalization). For a weight \(W \in \mathbb{R}^{n \times m}\) (operating on the smaller dimension, \(n \le m\)), fraction \(f\), \(k = \lceil f n \rceil\), momentum \(\mu\), learning rate \(\eta\) and weight decay \(\lambda\): \[\begin{aligned} M_t &= M_{t-1} + G_t \quad\text{(momentum accumulation; damping deferred)} \\ S_t &= \operatorname{topk}_{i}\bigl(\lVert M_{t,i,:}\rVert_1,\ k\bigr) \quad\text{(select the } k \text{ largest rows)} \\ O_t &= \operatorname{GramNS}\bigl(M_t[S_t,:]\bigr) \quad\text{(orthogonalize the } k \times m \text{ submatrix)} \\ W_t &= (1 - \eta\lambda)\, W_{t-1}, \qquad W_t[S_t,:] \gets W_t[S_t,:] - \eta\, O_t \\ M_t[S_t,:] &\gets \mu\, M_t[S_t,:] \quad\text{(error feedback: decay only selected rows)} \end{aligned}\] Gram Newton–Schulz normalizes \(X \gets M_t[S_t,:]/(\lVert M_t[S_t,:]\rVert_F + \epsilon)\) with \(\epsilon = 10^{-7}\), casts to float16, forms \(R_0 = X X^{\mathsf T}\) and \(Q_0 = I\), and runs \(T\) iterations \[\begin{aligned} Z_t &= b_t R_{t-1} + c_t R_{t-1}^2, & Q_t &= Q_{t-1} Z_t + a_t Q_{t-1}, \\ (RZ)_t &= R_{t-1} Z_t + a_t R_{t-1}, & R_t &= Z_t (RZ)_t + a_t (RZ)_t, \end{aligned}\] returning \(Q_T X\). A restart after iteration 2 (\(X \gets Q_2 X\), \(R_2 \gets X X^{\mathsf T}\), \(Q_2 \gets I\)) removes the spurious negative eigenvalues that half precision introduces in the Gram matrix; the coefficients \((a_t,b_t,c_t)\) are the Polar Express five-step triples . The cost falls from \(T(3\alpha+1)n^3\) FLOPs for standard Newton–Schulz to \((4T+3\alpha-3)n^3\) for Gram Newton–Schulz (\(\alpha = m/n\)), and selecting \(k = \lceil f n \rceil\) rows both shrinks the iteration and raises the aspect ratio by \(1/f\). The optional NorMuon flavor adds a row-wise second moment \(v_i \gets \beta_2 v_i + (1-\beta_2)\,\operatorname{mean}_j(O_{ij}^2)\) for \(i \in S\), rescales \(\tilde O_{ij} = O_{ij}/(\sqrt{v_i}+\epsilon)\), and renormalizes \(\tilde O \gets \tilde O\,\lVert O\rVert_F/\lVert\tilde O\rVert_F\) in float32 before the update.

Cons.

• Each step updates only a fraction of the rows, so the optimization trajectory differs from Muon; the paper reports the change is benign (Section 8.1 of the Dion3 report).

• Row selection and error feedback add top-\(k\) bookkeeping and a row-wise momentum decay per step.

• The largest wall-clock gains in the paper come from symmetric GEMM kernels and megabatched communication, which are hardware- and cluster-specific; only the algorithmic reduction and the single first-moment buffer apply here.

Pros.

• Gram Newton–Schulz is mathematically equivalent to Newton–Schulz and cuts orthogonalization FLOPs by about 55% at typical transformer aspect ratios with five iterations.

• Fractional row selection reduces the orthogonalized submatrix to \(k \times m\) and raises its aspect ratio by \(1/f\), compounding the Gram savings; the paper reports up to \(6\times\) faster optimizer steps at matched loss.

• At \(f = 1\) the update reduces to Muon (or NorMuon), and a single momentum buffer per matrix parameter suffices; error feedback lets ignored rows participate eventually.

ROOT (Robust Orthogonalized Optimizer).

ROOT  is a robustness-oriented revision of Muon that attacks two fragilities of the fixed-coefficient Newton–Schulz pipeline. First, the universal coefficients \((a,b,c)=(3.4445,-4.7750,2.0315)\) are calibrated for an “average” matrix shape and lose precision exactly on square and near-square layers: the orthogonalization MSE ranges from \(0.0499\) at \(2048\times2048\) to \(0.0003\) at \(2048\times16384\) in Table 1 of the paper. ROOT replaces them with per-shape coefficients \(\{a^{(m,n)}, b^{(m,n)}, c^{(m,n)}\}\) calibrated offline over momentum singular-value distributions (Mixed 1:3 real-to-random calibration), a scheme the authors call AdaNewton. Second, heavy-tailed gradient components distort the Frobenius normalization required by Newton–Schulz and are amplified by the polynomial iteration; ROOT models the momentum as \(M_t = B_t + O_t\) and solves a proximal L1 problem whose closed-form solution is soft-thresholding. Only the robust base component \(B_t\) is orthogonalized; the sparse outlier component \(O_t\) is discarded. For a hidden-layer weight matrix \(W \in \mathbb{R}^{m\times n}\), momentum \(\mu\), learning rate \(\eta\), percentile \(p\) and iteration count \(K\): \[\begin{aligned} M_t &= \mu M_{t-1} + \nabla \mathcal{L}(\theta_{t-1}) \quad\text{(momentum, Eq. 1)} \\ \varepsilon_t &= \operatorname{Quantile}(|M_t|, p) \quad\text{(dynamic threshold, } p=0.90 \text{ default)} \\ O_t &= \mathcal{T}_{\varepsilon_t}[M_t], \quad \mathcal{T}_\varepsilon[x]_i = \operatorname{sign}(x_i)\max(|x_i|-\varepsilon, 0) \quad\text{(soft-threshold, Eqs. 7--10)} \\ B_t &= M_t - O_t \quad\text{(clipped robust component, Eq. 9)} \\ X_0 &= \frac{B_t}{\lVert B_t\rVert_F}, \qquad X_k = a^{(m,n)}X_{k-1} + b^{(m,n)}X_{k-1}(X_{k-1}^\top X_{k-1}) + c^{(m,n)}X_{k-1}(X_{k-1}^\top X_{k-1})^2 \\ \theta_t &= \theta_{t-1} - \eta X_K \quad\text{(update, } K=5\text{)} \end{aligned}\]

Following the hybrid recipe of Muon, only hidden-layer 2-D matrices take the ROOT path; embedding tables, learned positional encodings, output heads and every non-matrix parameter are trained with decoupled AdamW using the group-level learning rate, weight decay, \(\beta_1, \beta_2\) and \(\epsilon\) settings. The paper’s LLM runs apply a \(0.2\) update-RMS scaling factor following Moonlight; this implementation lets the per-group learning rates absorb that calibration, so no scale constant is applied. Shapes absent from Table 9 fall back to Muon’s fixed coefficients, and the dynamic quantile costs an additional sort of \(|M_t|\) per matrix per step.

Cons.

• The percentile threshold requires an \(O(N\log N)\) sort of the momentum matrix at every step, on top of the Newton–Schulz matmuls.

• The calibrated coefficient table covers seven shapes only; every other shape falls back to Muon’s fixed coefficients.

• Outliers are discarded rather than shrunk, so an extremely noisy gradient can also lose informative signal; the paper’s ablation shows the trade-off is governed by the percentile (0.85 truncates too much, 0.99 too little).

Pros.

• Dimension-invariant orthogonalization precision: shape-calibrated coefficients remove the precision cliff that fixed coefficients show on square and near-square layers.

• Outlier-robust without retuning the Muon recipe; the paper reports lower loss and a better zero-shot average than Muon and AdamW at 1B scale (60.12 vs 59.59 vs 59.05) and gains on ViT/CIFAR-10.

• Same single momentum buffer per matrix parameter as Muon, and the softthresh/ada_newton toggles expose the paper’s ablations directly.

Muon² (Second-Moment Preconditioned Orthogonalization).

Muon\(^2\)  generalizes Muon by preconditioning the momentum with an Adam-style second-moment EMA before the Newton–Schulz orthogonalization. The paper shows that the practical bottleneck of Muon’s polar approximation is the ill-conditioned momentum matrix (early-training spectra span \(10^{-4}\)–\(1\) with roughly half the singular values in the NS “dead zone”); dividing by \(\sqrt{V_t}\) tightens the spectrum, so every NS step aligns better with the true orthogonalized update (cosine similarity \(0.916\) at \(K{=}3\) vs Muon’s \(0.931\) at \(K{=}5\), rising to \(0.975\) at \(K{=}5\)). For a hidden-layer weight matrix \(W \in \mathbb{R}^{n \times m}\) with learning rate \(\eta\): \[\begin{aligned} M_t &= \beta_1 M_{t-1} + (1-\beta_1) G_t && \text{(momentum EMA, Eq.~2)} \\ V_t &= \beta_2 V_{t-1} + (1-\beta_2)\, (G_t \odot G_t) && \text{(second-moment EMA, Eq.~3)} \\ \widetilde{M}_t &= M_t \oslash \big(\sqrt{V_t} + \epsilon \mathbf{1}\big) && \text{(preconditioning, Eq.~4)} \\ X_0 &= \frac{\widetilde{M}_t}{\|\widetilde{M}_t\|_F + \epsilon}, \qquad X_{k+1} = aX_k + b\, X_k X_k^\top X_k + c\, X_k (X_k X_k^\top)^2 X_k && \text{(NS, Eq.~5)} \\ W_{t+1} &= W_t - \eta \sqrt{m/n}\, X_K && \text{(update, Eq.~6)} \end{aligned}\] with the same quintic coefficients \((a,b,c)=(3.4445,-4.7750,2.0315)\) as Muon and the \(\sqrt{m/n}\) learning-rate factor proposed by Bernstein (2025). The memory-efficient variant Muon\(^2\)-F replaces \(V_t\) with Adafactor-style row/column statistics: \[\begin{aligned} r_t &= \beta_2 r_{t-1} + (1-\beta_2)\,\mathrm{sum}_{\mathrm{col}}(G_t \odot G_t), \qquad c_t = \beta_2 c_{t-1} + (1-\beta_2)\,\mathrm{sum}_{\mathrm{row}}(G_t \odot G_t), \\ \widehat{V}_t &= \frac{r_t\, c_t^\top}{\mathrm{sum}(r_t)}, \end{aligned}\] keeping memory within \(\sim\)0.3% of Muon at 7B scale (47.53 vs 47.29 GB) while retaining most of Muon\(^2\)’s benefit. Across GPT, LLaMA and Mixture-of-Experts pre-training up to 13B parameters, Muon\(^2\) outperforms Muon at equal NS iterations and matches it with \(K{=}3\) instead of \(5\) (a 40% reduction), saving up to a quarter of training time at equal loss.

Following the hybrid recipe of Muon, only hidden-layer 2-D matrices take the Muon\(^2\) path; embedding tables, learned positional encodings, output heads and every non-matrix parameter are trained with decoupled AdamW using the group-level learning rate, weight decay, \(\beta_1, \beta_2\) and \(\epsilon\) settings — the same \(\beta_1, \beta_2, \epsilon\) also parameterize Eqs. 2–4, so one per-group knob set configures both paths.

Cons.

• One extra full-size state buffer over Muon (two moment EMAs per matrix), doubling optimizer memory on the orthogonalized path unless the factored variant is enabled.

• The \(\sqrt{m/n}\) update scale differs from this repository’s plain muon path, so learning rates do not transfer one-to-one between the two registry entries.

• \(\beta_1, \beta_2\) defaults are not fixed by the paper; this implementation inherits the shared per-group defaults \((0.9, 0.95)\).

Pros.

• Better orthogonalization at equal NS iterations and equal quality with 40% fewer iterations (3 instead of 5), cutting NS compute and communication.

• Consistent gains over Muon across GPT, LLaMA and Mixture-of-Experts pre-training up to 13B parameters, with up to a quarter of training time saved at matched loss.

• The factored toggle exposes Muon\(^2\)-F, which keeps memory within \(\sim\)0.3% of Muon while preserving most of the benefit.

NorMuon (Neuron-wise Normalized Muon).

NorMuon  keeps Muon’s orthogonalized update and adds a neuron-wise (row-wise) adaptive learning rate on top of it. Muon’s Newton–Schulz iteration equalizes the singular values of the update, but the \(L_2\) norm of the individual neuron rows remains highly non-uniform: some neurons receive disproportionately large updates while others are barely trained. NorMuon tracks the mean squared magnitude of each row of the orthogonalized direction in a second-moment EMA, divides every neuron row by its RMS, and rescales the whole update to Adam’s empirical update scale so existing learning-rate schedules transfer unchanged. For a hidden-layer weight matrix \(W \in \mathbb{R}^{m \times n}\) with learning rate \(\eta\): \[\begin{aligned} M_t &= \beta_1 M_{t-1} + (1-\beta_1) G_t && \text{(momentum, line 5)} \\ O_t &= \mathrm{NS}_5(M_t) && \text{(orthogonalization, line 6)} \\ v_t &= \beta_2 v_{t-1} + (1-\beta_2)\,\mathrm{mean}_{\mathrm{cols}}(O_t \odot O_t) && \text{(neuron EMA, line 7)} \\ V_t &= \mathrm{ExpandRows}(v_t), \qquad \widehat{O}_t = O_t \oslash \big(\sqrt{V_t} + \epsilon\big) && \text{(row normalization, lines 8--9)} \\ \hat\eta &= \frac{0.2\,\eta\,\sqrt{mn}}{\lVert \widehat{O}_t \rVert_F}, \qquad W_{t+1} = W_t - \eta\lambda W_t - \hat\eta\,\widehat{O}_t && \text{(update, lines 10--11)} \end{aligned}\] with \(v_t \in \mathbb{R}^m\) (one scalar per neuron) and the same quintic Newton–Schulz coefficients \((a,b,c)=(3.4445,-4.7750,2.0315)\) as Muon. The per-neuron statistic costs \(m\) scalars against the \(m \times n\) first moment, so optimizer state is \(m(n+1)\) — within a \(1/n\) factor of Muon and roughly half of Adam’s \(2mn\). On 1.1B and 5.4B pretraining NorMuon needs 21.74% and 13.91% fewer steps than AdamW to reach the same validation loss; on Modded-NanoGPT (124M/350M) it reaches Muon’s loss with 6% and 15% fewer iterations, at a 2.9% training-step-time overhead over AdamW and 21.19 GB of optimizer state versus 40.56 GB for AdamW at 5.4B.

Following the hybrid recipe of Muon, only hidden-layer 2-D matrices take the NorMuon path; embedding tables, learned positional encodings, output heads and every non-matrix parameter are trained with decoupled AdamW using the group-level learning rate, weight decay, \(\beta_1, \beta_2\) and \(\epsilon\) settings. This implementation maps the paper’s parameters to the registry as follows: the momentum suffix is \(\beta_1\), the shared per-group betas[1] is \(\beta_2\), the shared per-group eps is \(\epsilon\), and ns_steps/ns_eps drive the Newton–Schulz iteration. Because NS5 starts with a Frobenius normalization, the paper’s \((1-\beta_1)\) momentum scaling and this repository’s standard momentum accumulation produce identical updates. The paper’s second contribution, Algorithm 2, distributes the orthogonalization across FSDP2 ranks (round-robin parameter assignment, gather/NS5/scatter, shard-local row normalization) to avoid replicated work; that deployment strategy is orthogonal to the update rule and is not implemented here.

Cons.

• One extra per-row state vector per matrix parameter (stored in float32); small next to the \(m \times n\) momentum, but non-zero.

• The normalization must follow orthogonalization: the paper’s “NorMuon (Front)” ablation, which applies it before NS5, is slightly worse.

• \(\epsilon\) is placed outside the root (\(\sqrt{V_t}+\epsilon\)), following Algorithm 2; Algorithm 1’s rendering of line 9 is typographically ambiguous.

• The neuron-wise statistics do not replace the AdamW fallback: embeddings, pos, heads and non-2-D tensors still carry full first/second moments.

Pros.

• Combines Muon’s low condition number with AdamW-like uniform per-neuron update norms, consistently outperforming both on 124M–5.4B pretraining.

• The \(0.2\eta\sqrt{mn}/\lVert\widehat{O}\rVert_F\) rescaling targets RMS \(0.2\eta\), so Adam learning-rate schedules apply without retuning.

• Memory stays close to Muon (\(m(n+1)\) versus \(mn\)), roughly halving Adam’s optimizer state, and the per-step overhead is a few percent.

Conda (Column-Normalized Adam).

Conda  keeps Adam’s coordinate-wise adaptivity and adds a milder, column-specific spectral normalization. The authors first reformulate Muon as a first/second-moment method whose second moment is the row-constant matrix \(\mathrm{diag}(\mathbf{\Sigma}_t)\mathbf{1}^\top\): Muon normalizes every column of the update by the same \(1/\sigma_i\) factors, which over-amplifies directions with tiny singular values. Conda instead constrains the second moment to the same subspace as the first moment by projecting the stochastic gradient, tracking an element-wise EMA of the squared projected gradients, and dividing the projected first moment by its square root. For a weight matrix \(W \in \mathbb{R}^{m \times n}\) with \(m \le n\), learning rate \(\eta\), decay rates \(\beta_1,\beta_2\), subspace refresh period \(T\), scale factor \(\alpha\) and stability constant \(\epsilon\): \[\begin{aligned} M_t &= \beta_1 M_{t-1} + (1-\beta_1) G_t \\ \bar U_t &= \begin{cases} \mathrm{SVD}(M_t), & t \bmod T = 0, \\ \bar U_{t-1}, & \text{otherwise}, \end{cases} \\ M'_t &= \bar U_t^\top M_t, \qquad N_t = \beta_2 N_{t-1} + (1-\beta_2)\big(\bar U_t^\top G_t\big)^2, \\ \widehat M_t &= \frac{M'_t}{1-\beta_1^t}, \qquad \widehat N_t = \frac{N_t}{1-\beta_2^t}, \\ W_t &= W_{t-1} - \eta\,\alpha\,\bar U_t\, \frac{\widehat M_t}{\sqrt{\widehat N_t} + \epsilon}. \end{aligned}\] For \(m > n\) the right singular basis \(V_t\) replaces \(\bar U_t\) (\(M'_t = M_t V_t\), \(N_t\) tracks \((G_t V_t)^2\), and the update is \((\widehat M_t/\sqrt{\widehat N_t})\,V_t^\top\)), exactly as in the released reference implementation. The SVD is the only super-linear step and runs once every \(T\) steps (\(T=2000\) in the paper’s pretraining runs, \(T=200\) in its LoRA runs); the wall-clock overhead is negligible for \(T \ge 100\).

On LLaMA 60M–1B pretraining Conda reaches validation perplexities of 28.32, 21.38, 16.44 and 13.59, against 34.06/25.08/18.80/15.56 for AdamW and 29.89/22.15/16.51/14.17 for Muon, i.e. 2–2.5\(\times\) AdamW’s convergence speed in steps and wall-clock time; it also leads on GPT-2 (125M/355M: 17.40/13.92) and on LoRA commonsense fine-tuning. Unlike AdamW and Muon it remains stable at learning rates up to \(2\times10^{-2}\), and the ablations show robustness to \(T \in [2, 20000]\), \(\beta_1 \in \{0.9, 0.95\}\) and \(\beta_2 \in \{0.95,\ldots,0.999\}\). Optimizer memory stays within 0.03–1.3 GB of AdamW from 60M to 1B. In this implementation the update_proj_gap suffix is \(T\) (default 2000), scale is \(\alpha\) (default 1.0; the paper’s pretraining uses 0.25), correct_bias toggles the \(1-\beta^t\) corrections, and the shared per-group betas/eps are \((\beta_1,\beta_2)\) (the paper uses 0.9/0.99) and \(\epsilon\); the shared wd is decoupled. Every 2-D parameter takes the Conda path regardless of group — there is no embedding/pos/head exclusion as in Muon — while every non-2-D parameter is trained with decoupled AdamW.

Cons.

• The SVD is a per-matrix cost spike every \(T\) steps; at \(T=2\) it raises LLaMA-60M wall-clock training time from 2h49min to 4h37min.

• The second moment must live in the same subspace as the first: the unprojected Adam estimator is unstable at scale and fails to converge on LLaMA-1B.

• Optimizer state is Adam-sized (full \(m\times n\) first moment) plus the projection basis; the memory advantage over AdamW is small, not the \(O(nr)\) of GaLore.

• Evaluated only up to 1B parameters and without Mixture-of-Experts; the paper flags both as future work.

Pros.

• Combines Adam’s coordinate-wise adaptivity with improved spectral conditioning in a single update rule, consistently beating AdamW, Muon, SOAP and Adafactor.

• The column-wise normalizer is milder than Muon’s uniform \(1/\sigma_i\) scaling and preserves the relative order of singular values, avoiding overshoot in small-singular directions.

• A single hyperparameter set transfers across 60M–1B, and training is stable at learning rates 10–20\(\times\) larger than AdamW’s.

• Memory and per-step cost stay close to AdamW for \(T \gg 1\); the lazy SVD makes the subspace refresh nearly free in wall-clock terms.

TUON (Tensor-Aware Optimization via T-SVD and Polarization).

TUON  lifts Muon’s layer-wise orthogonalization to cross-layer tensors using a transform-based tensor algebra derived from the Tensor-SVD. Layer-wise polar maps discard cross-layer correlations, while unfolding-based tensor optimizers (TEON) suffer dimension imbalance and cross-layer misalignment and degrade beyond two stacked layers or QKV-only modules. TUON instead mixes then separates: for one compatible group of \(K\) same-shape matrices with stacked momentum \(\mathcal{M}_t \in \mathbb{R}^{m \times n \times K}\): \[\begin{aligned} \mathcal{M}_t &= \beta \mathcal{M}_{t-1} + (1-\beta)\, \mathcal{G}_t && \text{(momentum tensor, Alg.~1 line 2)} \\ C_t &= \mathsf{M}_t \mathsf{M}_t^{\top} \in \mathbb{R}^{K \times K}, \qquad C_t = U_t \Lambda_t U_t^{\top}, \qquad Q_t = U_t^{\top} && \text{(layer Gram + eigendecomposition, lines 3--5)} \\ \widetilde{\mathcal{M}}_t &= \mathcal{M}_t \times_3 Q_t, \qquad \widetilde{P}_{t,j} = \operatorname{polar}\!\big(\widetilde{\mathcal{M}}_{t,j}\big) && \text{(transform + slice polarization, lines 6--9)} \\ \mathcal{P}_t &= \widetilde{\mathcal{P}}_t \times_3 Q_t^{H} && \text{(inverse transform, line 10)} \\ W_{t+1} &= (1 - \eta_t \lambda)\, W_t - \eta_t \gamma_t \mathcal{P}_t, \qquad \gamma_t = \frac{s}{\operatorname{RMS}(\mathcal{P}_t) + \epsilon} && \text{(RMS-scaled update, line 11)} \end{aligned}\] where \(\mathsf{M}_t\) is the mode-3 unfolding (\(k\)-th row \(= \operatorname{vec}(\mathcal{M}_t^{(k)})^{\top}\)), the mode-3 product mixes slices as \(\widetilde{\mathcal{M}}_{t,j} = \sum_k (Q_t)_{jk} \mathcal{M}_t^{(k)}\), \(\operatorname{RMS}(\mathcal{P}_t) = \sqrt{(mnK)^{-1} \sum_{i,j,k} |\mathcal{P}_t(i,j,k)|^2}\), and the target stack RMS is \(s = 1/\sqrt{\max(m,n)}\) (the rank–energy identity: full-rank slices give \(\operatorname{RMS} = 1/\sqrt{\max(m,n)}\), so \(\gamma_t \approx 1\) and per-matrix update magnitudes match Muon). The default transform DecorQ takes the eigenvectors of the layer Gram matrix (\(Q_t = U_t^{\top}\)) and is refreshed every \(T\) steps (default 200); it diagonalizes the transformed Gram, so the transformed slices are mutually Frobenius-orthogonal and each polar map acts in its own subspace — the coupled \(K\)-layer problem decouples into \(K\) independent single-channel problems. Fixed unitary transforms (DCT, Haar, DFT) are valid but impose a prior mismatched to layer statistics; in the paper’s ablation DecorQ (2.7598) beats DFT (2.7631), Haar (2.7635), static DecorQ (2.7658) and DCT (2.8100) on Qwen3-4B. Slice polarization uses the same quintic Newton–Schulz iteration as Muon, \((a,b,c) = (3.4445, -4.7750, 2.0315)\) with \(J = 5\) steps after a Frobenius normalization and a transpose-when-tall.

This implementation maps the paper’s “one compatible group” to each (parameter group, matrix shape) bucket: same-shape hidden 2-D matrices of a group are stacked along the layer mode in registration order (stack: full); stack: pairs restricts the stack to consecutive pairs (\(K=2\)). The transform suffix selects \(Q\) (decorq default; dct/haar/dft fixed transforms), momentum is \(\beta\), refresh_interval is \(T\), ns_steps/ns_eps drive the Newton–Schulz iteration, and rms_eps is the \(\epsilon\) of \(\gamma_t\). For \(K=1\) the transform cancels exactly and TUON reduces to Muon (paper Prop. 2); following the hybrid recipe of Muon, only hidden-layer 2-D matrices take the TUON path — embedding tables, learned positional encodings, output heads and every non-matrix parameter are trained with decoupled AdamW using the group-level learning rate, weight decay, \(\beta_1, \beta_2\) and \(\epsilon\) settings. Theory: TUON is steepest descent under the \(Q\)-transformed tensor spectral norm (dual norm pair \(\|\cdot\|_{Q,Ts}\) / \(\|\cdot\|_{Q,Tn}\)), its update equals the tensor polar factor of a T-product T-SVD, no unitary layer transform distorts the Muon geometry by more than \(\sqrt{K}\), and at equal operator radius the TUON oracle is never weaker than TEON’s and up to \(\sqrt{K}\) stronger, with an \(O(T^{-1/2})\) convergence certificate. Empirically TUON beats Muon and TEON on all of GPT-124M/350M and Qwen-0.6B/4B (validation loss 2.9357/2.6925/2.8887/2.7598), is the only method that improves from \(K{=}2\) to \(K{=}\mathrm{full}\) stacking on all four models, transfers from QKV to QKVO+MLP, and costs 0.2–1.1% step-time over TEON and at most 12.5% over AdamW; per-iteration complexity is \(\mathcal{O}(Kmn\min(m,n) + K^2 mn)\), the same dominant term as Muon.

Cons.

• Cross-layer mixing assumes the stacked matrices are genuinely compatible; this implementation stacks by exact shape within a parameter group, which is the closest faithful realization available from the unnamed factory param lists.

• The Gram/eigendecomposition adds \(\mathcal{O}(K^2 mn)\) per step and a \(K \times K\) cached basis per bucket; negligible for \(K \ll \min(m, n)\), but a per-bucket cost that Muon does not pay.

• \(\beta\) is required by Algorithm 1 but its default value is not stated in the paper; this implementation inherits the repo Muon-family convention (\(0.95\)).

• The dft transform runs the Newton–Schulz iteration over complex slices (Hermitian Gram); the real part is taken after the inverse transform, justified by the conjugate symmetry of the DFT of a real stack.

Pros.

• Exploits cross-layer structure that layer-wise Muon discards, with a convergence guarantee and exact steepest-descent interpretation under the induced tensor norm.

• DecorQ is data-dependent and refreshed every \(T\) steps; stacking \(K{=}\mathrm{full}\) improves with depth instead of degrading as unfolding methods do.

• Memory stays at Muon’s level (one momentum buffer per matrix plus a \(K \times K\) basis), and the RMS stack scaling keeps updates at Muon’s magnitude so existing learning rates transfer.

• Applies uniformly beyond QKV (attention, MLP and output projections), where prior tensor-aware optimizers fail to transfer.

NAMO (Norm-Based Adaptive Moment Estimation with Orthogonalized Momentum).

NAMO  keeps Muon’s orthogonalized update direction and adds an Adam-type norm-based noise adaptation on top of it. Matrix orthogonalization is an unbounded operation that can amplify stochastic noise, so Muon’s effective stepsize does not shrink when gradients become noisy. NAMO tracks a single scalar EMA of the squared Frobenius norm of the gradient and rescales the whole orthogonalized update by the bias-corrected norm ratio \(\lVert\widehat M_t\rVert_F/(\sqrt{\widehat v_t}+\epsilon_t)\) — an aggregated signal-to-noise statistic analogous to Adam’s coordinate-wise ratio. Because the scaling is a single scalar, the update direction stays exactly orthogonal, and the optimizer state grows by only one scalar per matrix over Muon. For a hidden-layer weight matrix \(W \in \mathbb{R}^{m\times n}\) with learning rate \(\eta\), momenta \(\mu_1,\mu_2 \in [0,1)\), stability constant \(\epsilon\) and decoupled weight decay \(\lambda\): \[\begin{aligned} M_t &= \mu_1 M_{t-1} + (1-\mu_1) G_t && \text{(momentum)} \\ v_t &= \mu_2 v_{t-1} + (1-\mu_2)\,\lVert G_t\rVert_F^2 && \text{(scalar norm EMA)} \\ O_t &= \mathrm{NS}_K(M_t) && \text{(orthogonalization)} \\ \alpha_t &= \frac{\sqrt{1-\mu_2^t}}{1-\mu_1^t}\cdot\frac{\lVert M_t\rVert_F}{\sqrt{v_t}+\epsilon} = \frac{\lVert\widehat M_t\rVert_F}{\sqrt{\widehat v_t}+\epsilon_t}, \qquad \epsilon_t := \frac{\epsilon}{\sqrt{1-\mu_2^t}} && \text{(adaptive stepsize)} \\ W_t &= W_{t-1} - \eta\,\alpha_t\,\bigl(O_t + \lambda W_{t-1}\bigr) && \text{(update)} \end{aligned}\] with the same quintic Newton–Schulz coefficients \((a,b,c)=(3.4445,-4.7750,2.0315)\) as Muon. Under nuclear-norm smoothness and bounded-variance noise, NAMO attains the optimal deterministic rate \(\mathcal{O}(T^{-1/2})\), and its stochastic bound decomposes into an optimization term \(\mathcal{O}(T^{-1/4})\) plus a variance term \(\mathcal{O}(\sqrt{\sigma}\,b^{-1/4}T^{-1/8})\), recovering the optimal \(\mathcal{O}(T^{-1/4})\) once the batch size satisfies \(b=\Omega(\sigma^2 T^{1/2})\).

On GPT-2 pretraining over OpenWebText (124M/355M, 2000-step linear warmup then constant LR), NAMO reaches validation losses of 3.0351/2.9516 versus 3.0643/2.9914 for AdamW and 3.0435/2.9684 for Muon, with optimal learning rates \(\eta=0.012/0.007\) (roughly \(9\times\) AdamW’s) at \(\mu_1=0.95\), \(\mu_2=0.99\), \(\lambda=0.01\), and it remains stable over a wider learning-rate grid than both baselines. In this implementation the global namo-betas suffix is \((\mu_1,\mu_2)\) (default \([0.95,0.99]\)), ns_steps/ns_eps configure Newton–Schulz, the shared per-group eps is \(\alpha_t\)’s \(\epsilon\), the shared wd is decoupled and scaled by \(\alpha_t\) as in the paper’s update, and the shared betas/eps configure the AdamW fallback for the embeddings/pos/head groups and every non-2-D parameter. The reference release additionally exposes Nesterov momentum, an effective-stepsize clamp and a \(0.2\sqrt{\max(m,n)}\) shape scaling; Algorithm 1 of the paper defines none of these, so this module follows the paper.

Cons.

• A single scalar stepsize carries no per-neuron information; the paper itself addresses this with the diagonal extension NAMO-D.

• The update is not RMS-aligned, so Adam learning-rate schedules do not transfer unchanged: the paper’s optima are roughly \(9\times\) AdamW’s.

• The theory assumes exact orthogonalization while practice uses \(K{=}5\) Newton–Schulz iterations, and the empirical evidence is limited to GPT-2-scale pretraining.

Pros.

• First principled integration of orthogonalized momentum with Adam-type noise adaptation, with optimal deterministic and noise-adaptive stochastic convergence rates.

• The scalar scaling preserves orthogonality exactly, and the effective stepsize shrinks automatically under noise or near stationarity.

• Zero additional memory over Muon and negligible per-step overhead; outperforms AdamW and Muon on GPT-2 (124M/355M) over a wider stable learning-rate range.

NAMO-D (Diagonal Extension of NAMO).

NAMO-D  replaces NAMO’s single adaptive scalar with a diagonal matrix that right-multiplies the orthogonalized momentum, so every neuron receives its own adaptive stepsize. Column-wise scaling aligns with the near block-diagonal Hessian of neural networks and yields finer-grained noise adaptation, but it no longer preserves strict orthogonality of the update direction. With the column-norm operator \([N^c(M)]_j := \lVert M_{:,j}\rVert\) for \(j=1,\dots,n\): \[\begin{aligned} M_t &= \mu_1 M_{t-1} + (1-\mu_1) G_t && \text{(momentum)} \\ [v_t]_j &= \mu_2 [v_{t-1}]_j + (1-\mu_2)\,\lVert [G_t]_{:,j}\rVert^2 && \text{(column norm EMA)} \\ d_t &= \frac{\sqrt{1-\mu_2^t}}{1-\mu_1^t}\,N^c(M_t) \oslash \bigl(\sqrt{v_t}+\epsilon\bigr) && \text{(column stepsizes)} \\ \bar d_t &= \lVert d_t\rVert_1/n, \qquad \tilde d_t = \min\bigl\{\max\{d_t, c\,\bar d_t \mathbf 1\}, (\bar d_t/c)\,\mathbf 1\bigr\} && \text{(clamp)} \\ O_t &= \mathrm{NS}_K(M_t), \qquad W_t = W_{t-1} - \eta\,\bigl(O_t + \lambda W_{t-1}\bigr)\operatorname{diag}(\tilde d_t) && \text{(update)} \end{aligned}\] By Lemma 7 of the paper the entries satisfy \(\lVert d_t\rVert_\infty \le \sqrt{(1-\mu_1)/(1-\mu_2)}\), and the clamp bounds the diagonal’s condition number by \(1/c^2\). The additional state over NAMO is one vector of length \(n\) per matrix; NAMO-D keeps the same optimal deterministic \(\mathcal{O}(T^{-1/2})\) and noise-adaptive stochastic rates as NAMO (Theorems 5–6). The extra clamping constant \(c\) balances the two competing goals: small \(c\) clamps harder toward the average and keeps the direction close to orthogonal, while larger \(c\) allows more neuron-wise variation.

The paper fixes \(c=0.1\) for GPT-2 (124M) and sweeps \(c\in\{0.12,0.40,0.75,0.90\}\) for GPT-2 (355M), selecting \(c=0.9\); the reference implementation defaults to \(0.75\). NAMO-D reaches validation losses of 3.0246 and 2.9507 on GPT-2 (124M/355M), improving on NAMO’s 3.0351/2.9516, at optimal learning rates \(\eta=0.009\) with the same \(\mu_1=0.95\), \(\mu_2=0.99\) and \(\lambda=0.01\). In this implementation the global namo_d-betas suffix is \((\mu_1,\mu_2)\) (default \([0.95,0.99]\)), namo_d-clamp is \(c\) (default \(0.1\)), ns_steps/ns_eps configure Newton–Schulz, the shared per-group eps is the denominator \(\epsilon\), the shared wd is decoupled and right-multiplied by the diagonal, and the shared betas/eps configure the AdamW fallback. As with NAMO, the official release’s Nesterov momentum, effective-stepsize clamp and \(0.2\sqrt{\max(m,n)}\) shape scaling are not part of Algorithm 2 and are not implemented.

Cons.

• The update no longer preserves strict orthogonality: the diagonal scaling trades the structural guarantee for finer-grained noise adaptation.

• The clamp constant \(c\) is a new hyperparameter whose optimum varies with model scale (\(0.1\) at 124M, \(0.9\) at 355M), requiring a sweep.

• One extra \(\mathbb{R}^n\) state per matrix over NAMO, and the theory still assumes exact orthogonalization.

Pros.

• Neuron-wise adaptation matches the near block-diagonal Hessian of neural networks and improves on NAMO on both GPT-2 sizes.

• The clamp bounds the diagonal’s condition number by \(1/c^2\), keeping the scaled direction well-conditioned under heavy-tailed or noisy gradients.

• Optimal deterministic and noise-adaptive stochastic convergence rates carry over from NAMO, with the same Newton–Schulz cost and hybrid AdamW routing.

Evolutionary and Population-Based Hybrids

EvE (Evolutionary Explorer).

EvE  is a population-based alternate to Adam targeted at the search-heavy, budget-constrained regime one level above final-stage training: ranking or pruning candidate configurations inside hyperparameter and architecture search. Instead of refining a single point, EvE keeps a steady-state population of four candidates and challenges one round-robin slot per iteration with a rand-to-pbest/1 differential-evolution child; only when the child fails to improve the slot’s incumbent does EvE run a short burst of stock Adam on the same minibatch, keeping the refined point only if it beats the child. Selection is greedy, so the best-so-far value is monotone non-increasing on deterministic objectives, and because gradients are used only as a targeted rescue, per-iteration cost stays within a constant factor of a single Adam step regardless of dimension.

All internal search happens on the unit cube \(z \in [0,1]^n\) via the affine map \(\mathbf{x} = \mathbf{x}_l + \mathbf{z} \odot (\mathbf{x}_u - \mathbf{x}_l)\), so a fixed-size step never overshoots narrow-range coordinates or under-moves wide ones. The population is initialized by proximity — one anchor with three Gaussian jitters, \(\sigma = 0.1\) — anchored at the model’s own initial weights (the paper’s LoRA precedent), with the per-tensor search box centered on the anchor at the Kaiming half-width \(\sqrt{6/\mathrm{fan\_\allowbreak in}}\). Each iteration draws three distinct donors \(a, b, c \ne i\) and builds the child with a jittered difference weight (the jitter applies only to the difference term; the pbest pull uses the unjittered \(F\)): \[\begin{aligned} z_{\text{child}} &= \operatorname{clamp}\big(z_a + F_t (z_b - z_c) + F (z_{\text{pbest}} - z_a),\ 0,\ 1\big), & F_t &= F\big(1 + \gamma (u - \tfrac{1}{2})\big), \quad u \sim U(0,1), \end{aligned}\] with \(F = 0.5\) and \(\gamma = 10^{-4}\). An optional polynomial mutation (PM; off by default for neural-network training, since a single-coordinate nudge essentially never wins its accept/reject gate at network scale) perturbs exactly one gene at rate \(p = 1/n\) with distribution index \(\eta = 20\), the perturbation bounded and self-scaling to the distance to each boundary: \[\begin{aligned} z_d &\leftarrow z_d + \delta_q (z_u - z_l)_d, & \delta_q &= \begin{cases} \bigl[2r + (1 - 2r)(1 - \delta_1)^{\frac{1}{\eta+1}}\bigr]^{\frac{1}{\eta+1}} - 1, & r \le 0.5,\\[2pt] 1 - \bigl[2(1 - r) + 2(r - 0.5)(1 - \delta_2)^{\frac{1}{\eta+1}}\bigr]^{\frac{1}{\eta+1}}, & r > 0.5, \end{cases} \end{aligned}\] with \(\delta_1 = (z_d - z_l)/(z_u - z_l)\), \(\delta_2 = (z_u - z_d)/(z_u - z_l)\) and \(r \sim U(0,1)\). When the (possibly mutated) child fails to beat the slot’s value, EvE runs \(n_{\text{adam}}\) stock-Adam steps on it in real coordinates, with classic L2 weight decay folded into the gradient before the moment update (PyTorch’s stock Adam, not AdamW’s decoupled rule) and every iterate clamped back into the box: \[\begin{aligned} g_t &= \nabla_{\mathbf{x}} f(\mathbf{x}_{t-1}; B_t) + \lambda \mathbf{x}_{t-1}, & m_t &= \beta_1 m_{t-1} + (1 - \beta_1) g_t, & v_t &= \beta_2 v_{t-1} + (1 - \beta_2) g_t^2, \\ \widehat{m}_t &= \frac{m_t}{1 - \beta_1^t}, & \widehat{v}_t &= \frac{v_t}{1 - \beta_2^t}, & \mathbf{x}_t &= \operatorname{clamp}\Big(\mathbf{x}_{t-1} - \eta_{lr} \frac{\widehat{m}_t}{\sqrt{\widehat{v}_t} + \epsilon},\ \mathbf{x}_l,\ \mathbf{x}_u\Big), \end{aligned}\] with \(\beta_1 = 0.9\), \(\beta_2 = 0.999\), \(\epsilon = 10^{-8}\). Each population slot keeps its own persistent \((m, v, t)\) state across the whole run, so fallback bursts form one continuous Adam trajectory rather than repeated cold starts. Every evaluation of an iteration — the slot refresh, the child and the whole fallback burst — is scored on one freshly drawn minibatch \(B_t\) (common random numbers), and a candidate replaces the incumbent only if it is strictly better on that same batch. In this implementation the shared per-group lr_, wd_, betas_ and eps_ configure the Adam fallback per parameter group (wd_ is the L2 term \(\lambda\) inside \(g_t\)), the live model parameters always hold the best-so-far slot between steps, and the step is closure-based — the trainer re-evaluates the current batch on each candidate point inside optimizer.step(), which is why EvE requires gradient_accumulation_steps: 1 and is not available for SBERT.