| """Diagonal Kalman Delta Network token mixer (arXiv:2609.07816).""" |
|
|
| from __future__ import annotations |
|
|
| import math |
| from typing import Any |
|
|
| import torch |
| from einops import rearrange |
| from fla.layers.utils import get_layer_cache, update_layer_cache |
| from fla.modules import FusedRMSNormGated, ShortConvolution |
| from torch import nn |
|
|
| from tiny_gdn.kdn import recurrence |
| from tiny_gdn.kdn.kernels import ( |
| _diag_kdn_frontend_precumsum_paired, |
| _diag_kdn_gain, |
| _diag_kdn_normalize, |
| chunk_kalman, |
| ) |
|
|
| SUPPORTED_HEAD_DIMS = (64, 128) |
|
|
|
|
| class DiagonalKalmanDeltaNet(nn.Module): |
| """Delta-rule memory whose write strength is a per-channel Kalman gain. |
| |
| Each head keeps a diagonal posterior covariance over its key channels: |
| channels that have absorbed little evidence take large corrections and |
| well-determined ones small corrections, replacing KDA's scalar ``beta``. |
| Stateless CUDA BF16 calls run the fused Triton kernels. Cached prefill, |
| decode, padding masks, and other dtypes use the exact PyTorch recurrences. |
| |
| The cache holds ``recurrent_state = (memory [B,H,K,K], precision [B,H,K])`` |
| in float32 plus the three short-convolution states. |
| """ |
|
|
| def __init__( |
| self, |
| hidden_size: int, |
| head_dim: int, |
| num_heads: int, |
| *, |
| conv_size: int = 4, |
| norm_eps: float = 1e-5, |
| layer_idx: int | None = None, |
| ) -> None: |
| super().__init__() |
| if head_dim not in SUPPORTED_HEAD_DIMS: |
| raise ValueError( |
| f"Diagonal KDN head_dim must be one of {SUPPORTED_HEAD_DIMS}, got {head_dim}" |
| ) |
| if num_heads <= 0 or hidden_size <= 0: |
| raise ValueError("hidden_size and num_heads must be positive") |
| if conv_size <= 0: |
| raise ValueError("conv_size must be positive") |
| self.hidden_size = hidden_size |
| self.head_dim = head_dim |
| self.num_heads = num_heads |
| self.key_dim = num_heads * head_dim |
| self.layer_idx = layer_idx |
| self.scale = head_dim**-0.5 |
| self.info_scale = float(head_dim) |
|
|
| self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False) |
| self.k_proj = nn.Linear(hidden_size, self.key_dim, bias=False) |
| self.v_proj = nn.Linear(hidden_size, self.key_dim, bias=False) |
| self.q_conv1d, self.k_conv1d, self.v_conv1d = ( |
| ShortConvolution( |
| hidden_size=self.key_dim, |
| kernel_size=conv_size, |
| bias=False, |
| activation="silu", |
| ) |
| for _ in range(3) |
| ) |
|
|
| self.f_proj = nn.Sequential( |
| nn.Linear(hidden_size, head_dim, bias=False), |
| nn.Linear(head_dim, self.key_dim, bias=False), |
| ) |
| self.A_log = nn.Parameter( |
| torch.log(torch.empty(num_heads, dtype=torch.float32).uniform_(1, 16)) |
| ) |
| self.A_log._no_weight_decay = True |
| dt = torch.exp( |
| torch.rand(self.key_dim, dtype=torch.float32) |
| * (math.log(0.1) - math.log(0.001)) |
| + math.log(0.001) |
| ).clamp(min=1e-4) |
| self.dt_bias = nn.Parameter(dt + torch.log(-torch.expm1(-dt))) |
| self.dt_bias._no_weight_decay = True |
|
|
| self.omega_proj = nn.Sequential( |
| nn.Linear(hidden_size, head_dim, bias=False), |
| nn.Linear(head_dim, self.key_dim, bias=True), |
| ) |
| self.r_proj = nn.Linear(hidden_size, num_heads, bias=True) |
| self.initial_precision_param = nn.Parameter( |
| torch.full( |
| (num_heads,), |
| recurrence.initial_precision_logit(), |
| dtype=torch.float32, |
| ) |
| ) |
| self.initial_precision_param._no_weight_decay = True |
|
|
| self.g_proj = nn.Sequential( |
| nn.Linear(hidden_size, head_dim, bias=False), |
| nn.Linear(head_dim, self.key_dim, bias=True), |
| ) |
| self.o_norm = FusedRMSNormGated(head_dim, activation="sigmoid", eps=norm_eps) |
| self.o_proj = nn.Linear(self.key_dim, hidden_size, bias=False) |
| self.reset_noise_parameters() |
|
|
| def reset_noise_parameters(self) -> None: |
| """Isotropic process noise and a neutral observation-noise bias at init. |
| |
| Blanket initializers overwrite these, so callers re-run this afterwards. |
| """ |
|
|
| nn.init.zeros_(self.omega_proj[1].weight) |
| nn.init.zeros_(self.omega_proj[1].bias) |
| nn.init.zeros_(self.r_proj.bias) |
|
|
| def _fused_supported(self, hidden_states: torch.Tensor) -> bool: |
| return hidden_states.is_cuda and hidden_states.dtype == torch.bfloat16 |
|
|
| def _split_heads(self, tensor: torch.Tensor) -> torch.Tensor: |
| return rearrange(tensor, "... (h d) -> ... h d", d=self.head_dim).contiguous() |
|
|
| def _fused( |
| self, |
| q: torch.Tensor, |
| k: torch.Tensor, |
| v: torch.Tensor, |
| retention_raw: torch.Tensor, |
| omega_raw: torch.Tensor, |
| r_raw: torch.Tensor, |
| ) -> torch.Tensor: |
| q_memory, k_gain, k_memory = _diag_kdn_normalize(q, k) |
| g_cumsum, alpha, omega, r, initial_precision = ( |
| _diag_kdn_frontend_precumsum_paired( |
| retention_raw, |
| omega_raw, |
| r_raw, |
| self.A_log.float().contiguous(), |
| self.dt_bias.float().view(self.num_heads, self.head_dim).contiguous(), |
| self.initial_precision_param.float().contiguous(), |
| ) |
| ) |
| kappa = _diag_kdn_gain( |
| k_gain, alpha, omega, r, initial_precision, omega_raw, r_raw |
| ) |
| return chunk_kalman(q_memory, k_memory, kappa, v, g_cumsum, self.scale) |
|
|
| def _stateful( |
| self, |
| q: torch.Tensor, |
| k: torch.Tensor, |
| v: torch.Tensor, |
| retention_raw: torch.Tensor, |
| omega_raw: torch.Tensor, |
| r_raw: torch.Tensor, |
| state: tuple[torch.Tensor, torch.Tensor] | None, |
| token_mask: torch.Tensor | None, |
| ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: |
| log_alpha, omega, r, initial_precision = recurrence.kalman_noise( |
| retention_raw, |
| omega_raw, |
| r_raw, |
| self.A_log, |
| self.dt_bias, |
| self.initial_precision_param, |
| ) |
| q, k = recurrence.normalize_query_key(q, k) |
| v = v.to(q.dtype) |
| if token_mask is not None: |
| k, v, log_alpha, omega = recurrence.mask_padding( |
| token_mask, k, v, log_alpha, omega |
| ) |
| if state is None: |
| memory, precision = recurrence.initial_state( |
| initial_precision, |
| batch_size=q.shape[0], |
| key_dim=self.head_dim, |
| value_dim=self.head_dim, |
| ) |
| else: |
| memory, precision = state |
| if q.shape[1] == 1: |
| output, memory, precision = recurrence.diag_kdn_step( |
| q[:, 0], |
| k[:, 0], |
| v[:, 0], |
| log_alpha[:, 0], |
| omega[:, 0], |
| r[:, 0], |
| memory, |
| precision, |
| scale=self.scale, |
| info_scale=self.info_scale, |
| ) |
| output = output[:, None] |
| else: |
| output, memory, precision = recurrence.diag_kdn_chunked( |
| q, |
| k, |
| v, |
| log_alpha, |
| omega, |
| r, |
| memory, |
| precision, |
| scale=self.scale, |
| info_scale=self.info_scale, |
| ) |
| return output, (memory, precision) |
|
|
| def forward( |
| self, |
| hidden_states: torch.Tensor, |
| attention_mask: torch.Tensor | None = None, |
| past_key_values: Any | None = None, |
| use_cache: bool = False, |
| output_attentions: bool = False, |
| **kwargs: Any, |
| ) -> tuple[torch.Tensor, None, Any | None]: |
| if hidden_states.ndim != 3 or hidden_states.shape[-1] != self.hidden_size: |
| raise ValueError( |
| f"hidden_states must have shape [B, T, {self.hidden_size}]" |
| ) |
| if attention_mask is not None and attention_mask.ndim != 2: |
| raise ValueError("attention_mask must be a [batch, sequence] padding mask") |
| if output_attentions: |
| raise ValueError("Diagonal KDN has no attention weights to return") |
| length = hidden_states.shape[1] |
| last_state = get_layer_cache(self, past_key_values) |
| token_mask = ( |
| None |
| if attention_mask is None |
| else attention_mask[:, -length:].to(device=hidden_states.device, dtype=torch.bool) |
| ) |
| conv_mask = ( |
| None if token_mask is None else token_mask.to(hidden_states.dtype) |
| ) |
| conv_q, conv_k, conv_v = ( |
| (None, None, None) if last_state is None else last_state["conv_state"] |
| ) |
| q, conv_q = self.q_conv1d( |
| x=self.q_proj(hidden_states), |
| mask=conv_mask, |
| cache=conv_q, |
| output_final_state=use_cache, |
| ) |
| k, conv_k = self.k_conv1d( |
| x=self.k_proj(hidden_states), |
| mask=conv_mask, |
| cache=conv_k, |
| output_final_state=use_cache, |
| ) |
| v, conv_v = self.v_conv1d( |
| x=self.v_proj(hidden_states), |
| mask=conv_mask, |
| cache=conv_v, |
| output_final_state=use_cache, |
| ) |
| q, k, v = (self._split_heads(tensor) for tensor in (q, k, v)) |
| retention_raw = self._split_heads(self.f_proj(hidden_states)) |
| omega_raw = self._split_heads(self.omega_proj(hidden_states)) |
| r_raw = self.r_proj(hidden_states).contiguous() |
|
|
| recurrent_state = None if last_state is None else last_state["recurrent_state"] |
| stateless = recurrent_state is None and token_mask is None and not use_cache |
| if stateless and self._fused_supported(hidden_states): |
| output = self._fused(q, k, v, retention_raw, omega_raw, r_raw) |
| else: |
| output, recurrent_state = self._stateful( |
| q, k, v, retention_raw, omega_raw, r_raw, recurrent_state, token_mask |
| ) |
| update_layer_cache( |
| self, |
| past_key_values, |
| recurrent_state=recurrent_state, |
| conv_state=(conv_q, conv_k, conv_v), |
| offset=length, |
| ) |
| output = self.o_norm( |
| output.to(v.dtype), |
| self._split_heads(self.g_proj(hidden_states)), |
| ) |
| output = self.o_proj(rearrange(output, "b t h d -> b t (h d)")) |
| return output, None, past_key_values |
|
|
