model.py

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