Maba v1.5-exp Architecture

Reference PyTorch implementation of the Maba v1.5 experimental hybrid architecture (DGDA recurrence and MABA-SA sparse attention).

Note: For the production release, see Maba v2 Architecture.

Maba v1.5 is an interleaved sub-quadratic hybrid model combining:

• 75% Decoupled Gated Delta Attention (DGDA): Linear recurrence with independent key erase gate b_t, value write gate w_t, channel-wise negative-softplus decay alpha_t, and chunkwise parallel prefill (C=16) via Order-3 Neumann series inversion.
• 25% Global MABA-SA Attention: Low-rank MLA key-value compression (d_c=128), strict NoPE (no positional embeddings), Delta-Guided Centroid Indexing (Top-32 blocks), and 3-stream output superposition (Local window + Top-32 sparse blocks + 64:1 HCA).
• High Computation Core (95.21%): Factorized token embeddings (32,768 -> 128 -> 640) constrain vocabulary tax to 4.30%, leaving 95.21% of parameters for sequence modeling layers.
• Native Speculative Drafter: Built-in Multi-Token Prediction (MTP k=2) heads for parallel token verification without companion models.

Documentation Index

Detailed technical documentation is organized across dedicated files:

• SCALING.md: Multi-scale parameter derivations from 100M to 30B, KV cache scaling to 1M context, and audits against 2026 foundation architectures (Qwen3.5, Muse-30B, Gemma4).
• BENCHMARK_REPORT.md: Empirical test verification, state memory invariance, and runtime scaling benchmarks.
• MABA_SPARSE_SPEC_AND_ROADMAP.md: Full mathematical derivations, proofs, gate mechanics, and algorithm pseudo-code.

Architectural Comparison

[!NOTE] Pretrained Weights and Downstream Evaluation This repository contains the reference architectural specification and PyTorch engine. Downstream task evaluations (ARC-Easy, HellaSwag, Story Cloze) and trained Safetensors weights will be published in the dedicated model weights repository upon completing pretraining runs.

Exact Parameter & Memory Breakdown (101M Reference Model)

1. Parameter Accounting

2. State Memory Scaling Across Context Horizons

* Note on 1M Context: The DG-Indexer selects Top-32 blocks (2048 active tokens), bounding working attention memory at 2.50 MB regardless of sequence length.

Core Formulation Reference

Clean mathematical representations without complex LaTeX macros:

1. Recurrence Step (DGDA, 75% of layers):
   State Update : S_t = diag(alpha_t) * S_{t-1} - k_t * ((b_t * k_t)^T * diag(alpha_t) * S_{t-1}) + (w_t * v_t) * k_t^T
   Readout      : o_t = q_t * S_t
   Where        : b_t in [0, 1]^d_k (erase gate), w_t in [0, 1]^d_v (write gate)
                  alpha_t = exp(-softplus(x_t * W_alpha)) (channel-wise decay)

2. Dynamic Sparse Attention (MABA-SA, 25% of layers):
   KV Latent    : c_t = W_down * x_t  (compressed from dim 640 to rank 128)
   Centroids    : C_i = 0.5 * mean(Block_i) + 0.5 * max(Block_i)
   Block Score  : Score(q_t, C_i) = q_t * C_i^T - lambda * log(1 + |t/B - i|)
   Routing      : Dynamically queries the Top-32 most salient blocks of size B=64

3. 3-Stream Output Superposition:
   O_t = g_local * O_local + g_sparse * O_sparse + g_hca * O_hca
   Where {g_local, g_sparse, g_hca} = softmax(x_t * W_gate)
   - Stream 1: Local window W=128 + 4 attention sinks
   - Stream 2: Dynamic Top-32 sparse blocks
   - Stream 3: 64:1 macro-averaged compressed history (HCA)

Tests

The test suite verifies numerical stability, boundary sequence lengths, causal masking, and O(1) state memory invariance:

pytest -q

• tests/test_dgda.py: Chunkwise Neumann prefill parity against recurrent decode (error < 4.58e-5).
• tests/test_dgda_stress.py: Arbitrary sequence lengths (1, 17, 33, 65, 128, 256) and extreme inputs (+-100).
• tests/test_indexer.py: Hybrid centroid pooling, distance penalties, and causal block masking.
• tests/test_sparse_attention.py: MLA compression, 3-stream superposition, and KV cache updates.
• tests/test_model.py: 101M parameter accounting, tied embeddings, MTP auxiliary loss, and generation.
• tests/test_nope_order_sensitivity.py: Permutation sensitivity under zero RoPE (L2 divergence = 39.518).
• tests/test_challenger_empirical.py: Strict O(1) state memory invariance (exactly 2,457,600 bytes).
• tests/test_ablations.py: Layer ablations (pure DGDA, no HCA, dense transformer baseline).

Quick Start & Reference API

1. Installation

git clone https://huggingface.co/AndrewThompson1233/maba-v1.5-exp-architecture
cd maba-v1.5-exp-architecture
pip install -e .

2. Autoregressive Generation

import torch
from maba_sparse.config import MabaSparseConfig
from maba_sparse.model import MabaSparseForCausalLM, get_101m_config

# Initialize 101M reference model
cfg = get_101m_config()
model = MabaSparseForCausalLM(cfg)

# Autoregressive generation
prompt = torch.tensor([[101, 2045, 312]])
generated = model.generate(prompt, max_new_tokens=32, temperature=0.7)
print("Generated token sequence:", generated.tolist())

3. Training & Benchmarking

# Model training smoke test
python train.py --model maba_sparse --steps 5 --batch_size 2 --seq_len 64

# Benchmark latency and state memory scaling
python benchmark.py --contexts 512,1024,2048,4096 --warmup 1 --repeats 2