OSDN

OSDN introduces an online diagonal preconditioner into the DeltaNet fast-weight write. Instead of changing the recurrent state size or the read path, OSDN keeps the standard DeltaNet residual read and writes with a feature-scaled key:

S_t = S_{t-1} + beta_t (v_t - S_{t-1} k_t) (d_t * k_t)^T

The preconditioner d_t is updated online from the key sequence and scalar write gates. This gives each key dimension its own adaptive scale while preserving the chunkwise DeltaNet computation pattern. OSDN-APF further adds a data-dependent retention gate on the preconditioner state, improving calibration under non-stationary contexts.

The same storage-key scaling applies cleanly beyond DeltaNet:

• OSDN: online-scaled DeltaNet.
• OSDN-APF: OSDN with adaptive preconditioner forgetting.
• OSGDN: online-scaled Gated DeltaNet.
• OSGDN-APF: OSGDN with adaptive preconditioner forgetting.
• OSKDA: online-scaled KDA.
• OSKDA-APF: OSKDA with adaptive preconditioner forgetting.

• Mechanism-targeted improvement. OSDN improves associative recall where fast-weight writes are expected to matter most.
• Provable online update. The diagonal preconditioner follows a token-local online update derived from the Delta Rule regression objective.
• Simple write-key substitution. The method is implemented by replacing the storage-side key with d_t * k_t.
• Broad backbone coverage. The reference code includes DeltaNet, Gated DeltaNet, KDA, and their online-scaled variants.

Matched 340M models are trained on FineWeb-Edu for 10B tokens. The 1.3B rows come from the matched 100B-token scaling run. LM PPL is the WikiText/LAMBADA geometric mean. Recall is JRT-style contains accuracy at 2K context. Only paper-reported metrics are included here.

The paper also measures the residual contraction ratio q_geo = f_t(S_t) / f_t(S_{t-1}) on repeated-recall prompts. Lower values mean stronger per-token residual contraction.

pip install -e ".[test]"

To run the optional Triton phase-1 parity test:

pip install -e ".[triton]"

import torch
from osdn import osdn_recurrent

B, T, H, K, V = 2, 16, 4, 32, 64
q = torch.randn(B, T, H, K)
k = torch.randn(B, T, H, K)
v = torch.randn(B, T, H, V)
beta = torch.rand(B, T, H).sigmoid()

result = osdn_recurrent(
    q, k, v, beta,
    d_min=0.5,
    d_max=2.0,
    initial_d=torch.ones(B, H, K),
)

out = result.output
d_traj = result.d_trajectory

All tensors use [batch, time, heads, dim] layout. Recurrent states are kept in fp32 in the reference implementation.

python examples/osdn_smoke.py
pytest -q

The test suite covers recurrence shapes, baseline equivalence when d_t = 1, autograd on the PyTorch reference paths, and optional Triton parity.

The production chunk and fused recurrent kernels are provided as an FLA overlay under fla/ops/. Copy these directories into a flash-linear-attention checkout or installed FLA source tree before using the high-performance paths:

rsync -av fla/ops/os_delta_rule/ /path/to/flash-linear-attention/fla/ops/os_delta_rule/
rsync -av fla/ops/os_gated_delta_rule/ /path/to/flash-linear-attention/fla/ops/os_gated_delta_rule/
rsync -av fla/ops/os_kda/ /path/to/flash-linear-attention/fla/ops/os_kda/

Included paper kernel groups:

• fla/ops/os_delta_rule/: OSDN and OSDN-APF DeltaNet chunk paths, fused recurrent paths, phase-1 sweeps, phase-2/WY helpers, and shared scaled-dot helpers used by OSGDN and OSKDA.
• fla/ops/os_gated_delta_rule/: OSGDN chunk and fused recurrent kernels, including the post-gate-regret path.
• fla/ops/os_kda/: OSKDA chunk, fused recurrent, intra-chunk, gate, and backward kernels.

The overlay is trimmed to these paper paths and their direct dependencies. Some internal file and function names keep the original FLA-compatible osgm prefix so the kernels can be copied into an FLA tree without additional import rewrites.

Public paper configs are stored under configs/ using the clean model names:

configs/
  osdn_340m/
    osdn_340m.json
    osdn_apf_340m.json
    osgdn_340m.json
    osgdn_apf_340m.json
    oskda_340m.json
    oskda_apf_340m.json
  osdn_1p3b/
    deltanet_1p3b_100b.json
    osdn_apf_1p3b_100b.json

The 340M OSDN entries use the paper hyperparameters:

{
  "d_min": 0.5,
  "d_max": 2.0,
  "initial_d": "ones"
}

src/osdn/reference.py       PyTorch reference recurrences
src/osdn/triton_phase1.py   optional Triton preconditioner sweep
fla/ops/                    FLA overlay for chunk/fused kernels
configs/                    public paper configuration records
examples/osdn_smoke.py      runnable smoke example
tests/                      correctness and parity checks