burn-gdn2-0.1.2 has been yanked.

burn-gdn2 — Gated DeltaNet 2

Linear-complexity recurrent token mixer for Burn. Channel-wise erase and write gates replace the scalar write-strength gate of the gated delta rule. Two forward modes: token-by-token recurrent (inference) and chunked WY decomposition (training).

O(T) complexity instead of O(T²) — trains on 4K+ sequences on a single GPU.

Requirements

Rust: ≥ 1.75
Burn: 0.21
GPU (optional): CUDA or Vulkan/Metal via WGPU

Install

cargo add burn-gdn2

Enable CUDA:

burn-gdn2 = { version = "0.1", features = ["cuda"] }

Quick start

use burn_ndarray::NdArray;
use burn_gdn2::{Gdn2Config, Gdn2Mode, GatedDeltaNet2};

let device = Default::default();
let cfg = Gdn2Config {
    hidden_size: 256,
    num_heads: 4,
    head_dim: 64,
    mode: Gdn2Mode::Chunk,
    ..Default::default()
};
let model = GatedDeltaNet2::<NdArray>::new(&cfg, &device);

// Training — chunked WY forward
let x = Tensor::random([1, 1024, 256], Distribution::Normal(0.0, 1.0), &device);
let output = model.forward_train(x);

// Inference — token-by-token with persistent state
let mut state = None;
let token = Tensor::random([1, 1, 256], Distribution::Normal(0.0, 1.0), &device);
let output = model.forward(token, &mut state, true);

Architecture

x ──→ Q/K/V/B/W/G projections ──→ ShortConv ──→ L2-Norm
                                       │
                    ┌──────────────────┤
                    ▼                  ▼
            Chunk WY (train)    Fused Recurrent (infer)
                    │                  │
                    ▼                  ▼
              Gated RMS Norm ←── SiLU Gate
                    │
                    ▼
               Output Proj

Per-token recurrence on matrix state S ∈ R^{d_k × d_v}:

S_t = (I - k_t (b_t ⊙ k_t)^T) diag(α_t) S_{t-1} + k_t (w_t ⊙ v_t)^T
o_t = q_t^T S

Where b_t is the channel-wise erase gate (key axis) and w_t is the channel-wise write gate (value axis), replacing the scalar gate of KDA.

Configuration

Field	Default	Description
`hidden_size`	2048	Model dimension
`num_heads`	16	Query/key heads
`head_dim`	128	Dimension per key/query head
`expand_v`	1.0	Value dimension expansion factor
`num_v_heads`	None	Value heads (None = same as num_heads)
`use_short_conv`	true	Depthwise conv before recurrence
`allow_neg_eigval`	false	Allow negative eigenvalues (b × 2)
`mode`	FusedRecurrent	Chunk (train) or FusedRecurrent (infer)
`chunk_size`	64	Chunk size for Chunk mode
`norm_eps`	1e-5	Output norm epsilon

Performance

Benchmarked on RTX 5060 Ti, Burn 0.21 vs equivalent PyTorch chunk implementation.

Chunk forward (training mode) is 2.2–5.5× faster across all configs — batched matmul for cumulative sum and pre-allocated slice_assign replace O(T) CPU-side loops.

Projections only (Q/K/V/B/W/G) are 0.6–0.8× vs PyTorch — Burn's cubecl matmul overhead for small matrix sizes. This is a fixed cost per layer and does not scale with sequence length.

Config	Burn CUDA	PyTorch CUDA	Speedup
seq=256, d=128	29 ms	63 ms	2.2×
seq=1024, d=256	117 ms	267 ms	2.3×
seq=2048, d=512	221 ms	645 ms	2.9×
seq=4096, d=1024	341 ms	1,875 ms	5.5×

License

AGPL-3.0. See LICENSE.

burn-gdn2 0.1.2