import dataclasses import math from typing import Optional, Union import torch import torch.nn.functional as F from torch import nn @dataclasses.dataclass class BDHConfig: n_layer: int = 2 n_embd: int = 128 n_head: int = 4 n_neuron: int = 2048 dropout: float = 0.01 vocab_size: int = 256 rope_theta: float = 2**16 state_decay: float = 1.0 state_clip: float = 0.0 attention_chunk_size: int = 64 def get_freqs(n: int, theta: float, dtype: torch.dtype) -> torch.Tensor: def quantize(t: torch.Tensor, q: int = 2) -> torch.Tensor: return (t / q).floor() * q return 1.0 / (theta ** (quantize(torch.arange(0, n, dtype=dtype)) / n)) / ( 2 * math.pi ) class BDHLinearAttention(nn.Module): """Paper-style BDH-GPU linear attention as explicit plastic state. For each layer/head we maintain rho in R^(n_head_neurons x d). At token t: a_t = rope(x_t) @ rho_{t-1} rho_t = decay * rho_{t-1} + outer(rope(x_t), v_t) This is the tensor form of the Hebbian synaptic update: co-active sparse neuron coordinates x_t and value/address vector v_t potentiate rho. """ def __init__(self, config: BDHConfig): super().__init__() if config.n_neuron % config.n_head != 0: raise ValueError("n_neuron must be divisible by n_head") self.n_head = config.n_head self.n_head_neuron = config.n_neuron // config.n_head self.n_embd = config.n_embd self.state_decay = config.state_decay self.state_clip = config.state_clip self.attention_chunk_size = config.attention_chunk_size self.register_buffer( "freqs", get_freqs( self.n_head_neuron, theta=config.rope_theta, dtype=torch.float32 ).view(1, 1, 1, self.n_head_neuron), persistent=False, ) @staticmethod def _phases_cos_sin(phases: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: phases = (phases % 1) * (2 * math.pi) return torch.cos(phases), torch.sin(phases) @staticmethod def _rope(phases: torch.Tensor, value: torch.Tensor) -> torch.Tensor: rotated = torch.stack((-value[..., 1::2], value[..., ::2]), dim=-1).view( *value.size() ) phases_cos, phases_sin = BDHLinearAttention._phases_cos_sin(phases) return (value * phases_cos).to(value.dtype) + (rotated * phases_sin).to( value.dtype ) def init_state( self, batch_size: int, dtype: torch.dtype, device: torch.device ) -> torch.Tensor: return torch.zeros( batch_size, self.n_head, self.n_head_neuron, self.n_embd, dtype=dtype, device=device, ) def apply_rope( self, x: torch.Tensor, offset: Union[int, torch.Tensor] = 0 ) -> torch.Tensor: seq_len = x.size(2) rel_pos = torch.arange(0, seq_len, device=x.device, dtype=self.freqs.dtype) if torch.is_tensor(offset): abs_pos = offset.to(device=x.device, dtype=self.freqs.dtype).view( x.size(0), 1, 1, 1 ) + rel_pos.view(1, 1, seq_len, 1) else: abs_pos = (rel_pos + offset).view(1, 1, seq_len, 1) phases = abs_pos * self.freqs return self._rope(phases, x) def forward( self, x: torch.Tensor, v: torch.Tensor, state: Optional[torch.Tensor] = None, offset: Union[int, torch.Tensor] = 0, ) -> tuple[torch.Tensor, torch.Tensor]: if state is None: state = self.init_state(x.size(0), v.dtype, x.device) if v.size(1) == 1: v = v.expand(-1, self.n_head, -1, -1) x_rope = self.apply_rope(x, offset=offset) chunk_size = max(1, self.attention_chunk_size) if chunk_size == 1: return self._forward_recurrent(x_rope, v, state) return self._forward_chunked(x_rope, v, state, chunk_size) def _forward_recurrent( self, x_rope: torch.Tensor, v: torch.Tensor, state: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: outputs = [] next_state = state for t in range(x_rope.size(2)): x_t = x_rope[:, :, t, :] v_t = v[:, :, t, :] a_t = torch.einsum("bhn,bhnd->bhd", x_t, next_state).unsqueeze(2) outputs.append(a_t) next_state = self.state_decay * next_state + torch.einsum( "bhn,bhd->bhnd", x_t, v_t ) return torch.cat(outputs, dim=2), next_state def _forward_chunked( self, x_rope: torch.Tensor, v: torch.Tensor, state: torch.Tensor, chunk_size: int, ) -> tuple[torch.Tensor, torch.Tensor]: outputs = [] next_state = state seq_len = x_rope.size(2) decay = float(self.state_decay) for start in range(0, seq_len, chunk_size): stop = min(start + chunk_size, seq_len) x_chunk = x_rope[:, :, start:stop, :] v_chunk = v[:, :, start:stop, :] cur_len = stop - start from_state = torch.einsum("bhtn,bhnd->bhtd", x_chunk, next_state) scores = x_chunk @ x_chunk.transpose(-1, -2) if decay == 1.0: causal = torch.ones( cur_len, cur_len, dtype=torch.bool, device=x_rope.device, ).tril(diagonal=-1) within = scores.masked_fill(~causal, 0) @ v_chunk state_update = torch.einsum("bhtn,bhtd->bhnd", x_chunk, v_chunk) next_state = next_state + state_update outputs.append(from_state + within) continue time = torch.arange(cur_len, device=x_rope.device, dtype=torch.float32) decay_base = torch.full_like(time, decay) state_weight = torch.pow(decay, time).to(dtype=from_state.dtype).view( 1, 1, cur_len, 1 ) from_state = from_state * state_weight i = torch.arange(cur_len, device=x_rope.device).view(cur_len, 1) j = torch.arange(cur_len, device=x_rope.device).view(1, cur_len) lower = i > j exponents = (i - 1 - j).clamp_min(0).to(torch.float32) weights = torch.pow(torch.full_like(exponents, decay), exponents).to( dtype=scores.dtype ) weights = weights.masked_fill(~lower, 0) within = (scores * weights.view(1, 1, cur_len, cur_len)) @ v_chunk update_weight = torch.pow(decay_base, cur_len - 1 - time).to( dtype=x_chunk.dtype ).view(1, 1, cur_len, 1) state_update = torch.einsum( "bhtn,bhtd->bhnd", x_chunk * update_weight, v_chunk ) next_state = (decay**cur_len) * next_state + state_update outputs.append(from_state + within) return torch.cat(outputs, dim=2), next_state class BabyDragonHatchling(nn.Module): """BDH-GPU state-space language model from the paper equations. Trainable offline parameters: E / encoder: R^(n x d) D_x / decoder_x: per-head R^(d x n/h) D_y / decoder_y: per-head R^(d x n/h) token encoder/decoder Online state: rho_l: per-layer recurrent Hebbian state, updated while tokens stream. """ def __init__(self, config: BDHConfig): super().__init__() if config.n_neuron % config.n_head != 0: raise ValueError("n_neuron must be divisible by n_head") if config.n_embd <= 0 or config.n_neuron <= 0: raise ValueError("n_embd and n_neuron must be positive") self.config = config n_per_head = config.n_neuron // config.n_head self.ln = nn.LayerNorm(config.n_embd, elementwise_affine=False, bias=False) self.wte = nn.Embedding(config.vocab_size, config.n_embd) self.drop = nn.Dropout(config.dropout) self.encoder = nn.Parameter(torch.empty(config.n_neuron, config.n_embd)) self.decoder_x = nn.Parameter( torch.empty(config.n_head, config.n_embd, n_per_head) ) self.decoder_y = nn.Parameter( torch.empty(config.n_head, config.n_embd, n_per_head) ) self.readout = nn.Parameter(torch.empty(config.n_embd, config.vocab_size)) self.attn = BDHLinearAttention(config) self.last_stats: dict[str, torch.Tensor] = {} self.apply(self._init_module) nn.init.normal_(self.encoder, mean=0.0, std=0.02) nn.init.normal_(self.decoder_x, mean=0.0, std=0.02) nn.init.normal_(self.decoder_y, mean=0.0, std=0.02) nn.init.normal_(self.readout, mean=0.0, std=0.02) @staticmethod def _init_module(module: nn.Module) -> None: if isinstance(module, nn.Embedding): nn.init.normal_(module.weight, mean=0.0, std=0.02) def init_plastic_state( self, batch_size: int, device: torch.device ) -> list[torch.Tensor]: return [ self.attn.init_state(batch_size, self.wte.weight.dtype, device) for _ in range(self.config.n_layer) ] def reset_plastic_state_rows( self, state: list[torch.Tensor], reset_mask: torch.Tensor ) -> list[torch.Tensor]: if not bool(reset_mask.any()): return state for s in state: s[reset_mask] = 0 return state def clamp_plastic_state(self, state: list[torch.Tensor]) -> list[torch.Tensor]: state_clip = float(self.config.state_clip) if state_clip <= 0: return state for s in state: s.clamp_(min=-state_clip, max=state_clip) return state def forward_features( self, idx: torch.Tensor, plastic_state: Optional[list[torch.Tensor]] = None, state_offset: Union[int, torch.Tensor] = 0, collect_stats: bool = False, ) -> tuple[torch.Tensor, list[torch.Tensor]]: batch_size, _ = idx.size() cfg = self.config if plastic_state is None: plastic_state = self.init_plastic_state(batch_size, idx.device) v = self.ln(self.wte(idx).unsqueeze(1)) x_prev = torch.zeros( batch_size, cfg.n_head, idx.size(1), cfg.n_neuron // cfg.n_head, dtype=v.dtype, device=idx.device, ) next_states = [] x_sparsity = [] y_sparsity = [] x_mean = [] y_mean = [] for layer in range(cfg.n_layer): x = x_prev + F.relu(v @ self.decoder_x) if collect_stats: x_sparsity.append((x > 0).float().mean().detach()) x_mean.append(x.detach().abs().mean()) a, next_state = self.attn( x=x, v=v, state=plastic_state[layer], offset=state_offset, ) next_states.append(next_state) y = F.relu(self.ln(a) @ self.decoder_y) * x if collect_stats: y_sparsity.append((y > 0).float().mean().detach()) y_mean.append(y.detach().abs().mean()) y = y.transpose(1, 2).reshape(batch_size, 1, idx.size(1), cfg.n_neuron) y = self.drop(y) v = self.ln(y @ self.encoder) x_prev = x if collect_stats: self.last_stats = { "activation/x_sparsity": torch.stack(x_sparsity).mean(), "activation/y_sparsity": torch.stack(y_sparsity).mean(), "activation/x_abs_mean": torch.stack(x_mean).mean(), "activation/y_abs_mean": torch.stack(y_mean).mean(), } else: self.last_stats = {} return v.squeeze(1), next_states def forward( self, idx: torch.Tensor, targets: Optional[torch.Tensor] = None, plastic_state: Optional[list[torch.Tensor]] = None, state_offset: Union[int, torch.Tensor] = 0, collect_stats: bool = False, ) -> tuple[torch.Tensor, Optional[torch.Tensor], list[torch.Tensor]]: v, next_state = self.forward_features( idx, plastic_state=plastic_state, state_offset=state_offset, collect_stats=collect_stats, ) logits = v @ self.readout loss = None if targets is not None: loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1)) return logits, loss, next_state @torch.no_grad() def generate( self, idx: torch.Tensor, max_new_tokens: int, temperature: float = 1.0, top_k: Optional[int] = None, ) -> torch.Tensor: self.eval() state = self.init_plastic_state(idx.size(0), idx.device) offset = 0 logits = None for t in range(idx.size(1)): logits, _, state = self( idx[:, t : t + 1], plastic_state=state, state_offset=offset, ) offset += 1 for _ in range(max_new_tokens): assert logits is not None next_logits = logits[:, -1, :] if temperature <= 0: idx_next = torch.argmax(next_logits, dim=-1, keepdim=True) else: next_logits = next_logits / temperature if top_k is not None: values, _ = torch.topk( next_logits, min(top_k, next_logits.size(-1)) ) next_logits = next_logits.masked_fill( next_logits < values[:, [-1]], float("-inf") ) probs = F.softmax(next_logits, dim=-1) if not torch.isfinite(probs).all(): raise RuntimeError( "non-finite sampling probabilities; try a higher temperature" ) idx_next = torch.multinomial(probs, num_samples=1) idx = torch.cat((idx, idx_next), dim=1) logits, _, state = self( idx_next, plastic_state=state, state_offset=offset, ) offset += 1 return idx