tinymind-native-colab-handoff/bundle/model/axiom_weave.py

"""TinyMind AxiomWeave: routed synthesis architecture.

AxiomWeave is a research architecture that combines complementary model
mechanisms inside one block:

- local exact-ish mixing through gated linear attention,
- dynamical compression through selective state space updates,
- bounded long-memory through PureField recurrent memory,
- symbolic/tool compatibility through stable routing and evidence hooks,
- parameter efficiency through shared low-rank PureField weights and KAN FFN.

It is new code in this repository, but it is not a world-best claim. The claim
gate lives in the dossier and requires external benchmark evidence.
"""

from __future__ import annotations

import torch
import torch.nn as nn

from .config import OmegaConfig
from .layers import GatedLinearAttention, KANFeedForward, RMSNorm, SelectiveSSM
from .purefield import PureFieldBlock, PureFieldShared


class AxiomWeaveBlock(nn.Module):
    """One routed block that fuses attention, SSM, PureField, and KAN paths."""

    def __init__(self, cfg: OmegaConfig, layer_index: int, shared: PureFieldShared | None = None):
        super().__init__()
        self.cfg = cfg
        self.layer_index = int(layer_index)
        self.norm = RMSNorm(cfg.dim)
        self.local_attention = GatedLinearAttention(cfg)
        self.state_space = SelectiveSSM(cfg)
        self.purefield = PureFieldBlock(cfg, layer_index=layer_index, shared=shared)
        self.router = nn.Linear(cfg.dim, 3, bias=True)
        self.router_temp = max(float(getattr(cfg, "contractive_eps", 1e-3)), 1e-4) ** 0.25
        self.post_norm = RMSNorm(cfg.dim)
        self.ffn = KANFeedForward(cfg)
        self.residual_scale = min(float(getattr(cfg, "residual_alpha", 0.2)), cfg.n_layers ** -0.5)

    def forward(
        self,
        x: torch.Tensor,
        cache: dict | None = None,
        mask: torch.Tensor | None = None,
        return_stats: bool = False,
    ) -> tuple[torch.Tensor, dict] | tuple[torch.Tensor, dict, dict[str, torch.Tensor]]:
        cache = cache or {}
        u = self.norm(x)
        attn_out, attn_cache = self.local_attention(u, cache.get("attention"), mask)
        ssm_out, ssm_cache = self.state_space(u, cache.get("ssm"), mask)
        field_out, field_cache, field_stats = self.purefield(
            u,
            kv_cache=cache.get("purefield"),
            mask=mask,
            return_stats=True,
        )
        weights = torch.softmax(self.router(u) / self.router_temp, dim=-1)
        mixed = (
            weights[..., 0:1] * attn_out
            + weights[..., 1:2] * ssm_out
            + weights[..., 2:3] * field_out
        )
        y = x + self.residual_scale * torch.tanh(mixed)
        y = y + self.residual_scale * torch.tanh(self.ffn(self.post_norm(y)))
        new_cache = {"attention": attn_cache, "ssm": ssm_cache, "purefield": field_cache}
        if not return_stats:
            return y, new_cache
        stats = {
            "route_weights_mean": weights.detach().mean(dim=(0, 1)),
            "route_entropy": (-(weights * torch.log(weights.clamp_min(1e-8))).sum(dim=-1)).detach().mean(),
            "branch_norms": torch.stack(
                [
                    attn_out.detach().norm(dim=-1).mean(),
                    ssm_out.detach().norm(dim=-1).mean(),
                    field_out.detach().norm(dim=-1).mean(),
                ]
            ),
            "purefield_memory_norm": field_stats["memory_norm"].detach().mean(),
        }
        return y, new_cache, stats


class AxiomWeaveModel(nn.Module):
    """Small LM wrapper for AxiomWeave blocks."""

    def __init__(self, cfg: OmegaConfig):
        super().__init__()
        self.cfg = cfg
        cfg.architecture_mode = "axiomweave"
        self.embed = nn.Embedding(cfg.vocab_size, cfg.dim, padding_idx=cfg.pad_token_id)
        self.shared = PureFieldShared(cfg)
        self.blocks = nn.ModuleList([AxiomWeaveBlock(cfg, i, self.shared) for i in range(cfg.n_layers)])
        self.norm = RMSNorm(cfg.dim)
        self.lm_head = nn.Linear(cfg.dim, cfg.vocab_size, bias=False)
        if cfg.tie_word_embeddings:
            self.lm_head.weight = self.embed.weight
        self._init_weights()

    def _init_weights(self) -> None:
        nn.init.normal_(self.embed.weight, std=0.02)
        for module in self.modules():
            if isinstance(module, nn.Linear):
                nn.init.normal_(module.weight, std=0.02)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        labels: torch.Tensor | None = None,
        caches: list[dict] | None = None,
        return_stats: bool = False,
    ) -> dict:
        x = self.embed(input_ids)
        new_caches = []
        stats = []
        for i, block in enumerate(self.blocks):
            cache_in = caches[i] if caches else None
            if return_stats:
                x, cache_out, block_stats = block(x, cache=cache_in, mask=attention_mask, return_stats=True)
                stats.append(block_stats)
            else:
                x, cache_out = block(x, cache=cache_in, mask=attention_mask)
            new_caches.append(cache_out)
        logits = self.lm_head(self.norm(x))
        result = {"logits": logits, "caches": new_caches}
        if labels is not None:
            result["loss"] = nn.functional.cross_entropy(
                logits[..., :-1, :].contiguous().view(-1, self.cfg.vocab_size),
                labels[..., 1:].contiguous().view(-1),
                ignore_index=-100,
            )
        if return_stats:
            result["stats"] = stats
        return result

    def count_params_int(self) -> int:
        return sum(p.numel() for p in self.parameters())