Upload cognet_model.py with huggingface_hub

cognet_model.py

@@ -0,0 +1,646 @@
+"""
+CogNet1B: Non-Transformer Language Model with Cognitive Routing
+================================================================
+Replaces self-attention with O(n) cognitive routing and
+hierarchical memory, enabling linear-time sequence processing.
+
+Key architectural innovations:
+- CognitiveChannel: Depthwise separable conv + SwiGLU FFN (O(n) per channel)
+- CoherenceRouter: O(n) routing via learned coherence scoring
+- SharedHierarchicalMemory: 3-tier key-value memory (Working/Episodic/Semantic)
+- AdaptiveComputationBlock: Variable-depth processing per token
+- CompositionalReasoner: Hyperdimensional computing for role-filler binding
+"""
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Dict, Optional, Tuple
+
+
+# ─── Token Encoder ───────────────────────────────────────────────────────────
+
+class TokenEncoder(nn.Module):
+    """Token embedding + learned positional encoding."""
+
+    def __init__(self, vocab_size: int, hidden_dim: int, max_seq_len: int, dropout: float = 0.1):
+        super().__init__()
+        self.token_emb = nn.Embedding(vocab_size, hidden_dim)
+        self.pos_emb = nn.Embedding(max_seq_len, hidden_dim)
+        self.dropout = nn.Dropout(dropout)
+        self.norm = nn.LayerNorm(hidden_dim)
+        self._init_weights()
+
+    def _init_weights(self):
+        nn.init.normal_(self.token_emb.weight, mean=0.0, std=0.02)
+        nn.init.normal_(self.pos_emb.weight, mean=0.0, std=0.02)
+
+    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
+        B, T = input_ids.shape
+        positions = torch.arange(T, device=input_ids.device).unsqueeze(0).expand(B, -1)
+        x = self.token_emb(input_ids) + self.pos_emb(positions)
+        return self.dropout(self.norm(x))
+
+
+# ─── Cognitive Channel ───────────────────────────────────────────────────────
+
+class CognitiveChannel(nn.Module):
+    """Depthwise separable convolution + SwiGLU FFN — O(n) per channel."""
+
+    def __init__(self, channel_dim: int, ff_dim: int, dropout: float = 0.1):
+        super().__init__()
+        # Depthwise separable conv
+        self.dw_conv = nn.Conv1d(
+            channel_dim, channel_dim, kernel_size=3, padding=1,
+            groups=channel_dim
+        )
+        self.pw_conv = nn.Conv1d(channel_dim, channel_dim, kernel_size=1)
+        self.conv_norm = nn.LayerNorm(channel_dim)
+        self.conv_dropout = nn.Dropout(dropout)
+
+        # SwiGLU FFN
+        self.ff_gate = nn.Linear(channel_dim, ff_dim, bias=False)
+        self.ff_up = nn.Linear(channel_dim, ff_dim, bias=False)
+        self.ff_down = nn.Linear(ff_dim, channel_dim, bias=False)
+        self.ff_norm = nn.LayerNorm(channel_dim)
+        self.ff_dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: (B, T, D)
+        residual = x
+        # Conv path
+        h = x.transpose(1, 2)                    # (B, D, T)
+        h = self.dw_conv(h)
+        h = self.pw_conv(h)
+        h = h.transpose(1, 2)                     # (B, T, D)
+        h = self.conv_norm(h)
+        x = residual + self.conv_dropout(h)
+
+        # FFN path (SwiGLU)
+        residual = x
+        gate = F.silu(self.ff_gate(x))
+        up = self.ff_up(x)
+        h = gate * up
+        h = self.ff_down(h)
+        h = self.ff_norm(h)
+        x = residual + self.ff_dropout(h)
+        return x
+
+
+# ─── Coherence Router ────────────────────────────────────────────────────────
+
+class CoherenceRouter(nn.Module):
+    """O(n) routing: compute which channel handles each token."""
+
+    def __init__(self, hidden_dim: int, num_channels: int, routing_iters: int = 1):
+        super().__init__()
+        self.num_channels = num_channels
+        self.routing_iters = routing_iters
+        self.query = nn.Linear(hidden_dim, num_channels, bias=False)
+        self.key = nn.Linear(hidden_dim, num_channels, bias=False)
+        self.refine = nn.Linear(num_channels, num_channels, bias=False)
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Args:
+            x: (B, T, D)
+        Returns:
+            routing_weights: (B, T, num_channels)  — soft assignment
+            channel_masks:   (B, T, num_channels)  — hard top-k for efficiency
+        """
+        B, T, D = x.shape
+        q = self.query(x)  # (B, T, C)
+        k = self.key(x)    # (B, T, C)
+
+        # O(n) coherence: dot-product of each token's query with mean key
+        mean_key = k.mean(dim=1, keepdim=True)  # (B, 1, C)
+        scores = q * mean_key                    # (B, T, C)
+        scores = scores + self.refine(scores) * 0.1  # one refinement step
+        routing_weights = F.softmax(scores, dim=-1)   # (B, T, C)
+
+        # Hard routing: top-2 channels per token
+        _, top_idx = routing_weights.topk(2, dim=-1)
+        channel_masks = torch.zeros_like(routing_weights)
+        channel_masks.scatter_(-1, top_idx, 1.0)
+
+        return routing_weights, channel_masks
+
+
+# ─── Shared Hierarchical Memory ──────────────────────────────────────────────
+
+class SharedHierarchicalMemory(nn.Module):
+    """3-tier memory: Working → Episodic → Semantic with key-value attention."""
+
+    def __init__(self, hidden_dim: int, key_dim: int,
+                 working_slots: int, episodic_slots: int, semantic_slots: int,
+                 dropout: float = 0.1):
+        super().__init__()
+        self.key_dim = key_dim
+        self.working_slots = working_slots
+        self.episodic_slots = episodic_slots
+        self.semantic_slots = semantic_slots
+
+        # Key / value projections
+        self.q_proj = nn.Linear(hidden_dim, key_dim, bias=False)
+        self.k_proj = nn.Linear(hidden_dim, key_dim, bias=False)
+        self.v_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.out_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
+
+        # Learnable memory slots
+        self.working_keys = nn.Parameter(torch.randn(working_slots, key_dim) * 0.02)
+        self.working_vals = nn.Parameter(torch.randn(working_slots, hidden_dim) * 0.02)
+        self.episodic_keys = nn.Parameter(torch.randn(episodic_slots, key_dim) * 0.02)
+        self.episodic_vals = nn.Parameter(torch.randn(episodic_slots, hidden_dim) * 0.02)
+        self.semantic_keys = nn.Parameter(torch.randn(semantic_slots, key_dim) * 0.02)
+        self.semantic_vals = nn.Parameter(torch.randn(semantic_slots, hidden_dim) * 0.02)
+
+        # Gating between tiers
+        self.tier_gate = nn.Linear(hidden_dim * 3, 3, bias=False)
+        self.norm = nn.LayerNorm(hidden_dim)
+        self.dropout = nn.Dropout(dropout)
+
+    def _read_tier(self, queries: torch.Tensor, keys: torch.Tensor,
+                   values: torch.Tensor) -> torch.Tensor:
+        """
+        Read from one memory tier.
+        queries: (B, T, key_dim)
+        keys:    (S, key_dim)
+        values:  (S, hidden_dim)
+        Returns: (B, T, hidden_dim)
+        """
+        B = queries.shape[0]
+        # BUG FIX: expand keys/vales to batch dim without transposing last two dims
+        keys_expanded = keys.unsqueeze(0).expand(B, -1, -1)    # (B, S, key_dim)
+        values_expanded = values.unsqueeze(0).expand(B, -1, -1) # (B, S, hidden_dim)
+
+        # Scaled dot-product attention (O(n*S) but S is small)
+        scale = math.sqrt(self.key_dim)
+        # (B, T, key_dim) @ (B, key_dim, S) → (B, T, S)
+        attn = torch.bmm(queries, keys_expanded.transpose(1, 2)) / scale
+        attn = F.softmax(attn, dim=-1)
+        # (B, T, S) @ (B, S, hidden_dim) → (B, T, hidden_dim)
+        out = torch.bmm(attn, values_expanded)
+        return out
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+        B, T, D = x.shape
+        queries = self.q_proj(x)  # (B, T, key_dim)
+
+        # Read from each tier
+        w_out = self._read_tier(queries, self.working_keys, self.working_vals)
+        e_out = self._read_tier(queries, self.episodic_keys, self.episodic_vals)
+        s_out = self._read_tier(queries, self.semantic_keys, self.semantic_vals)
+
+        # Gated combination
+        gate_input = torch.cat([w_out, e_out, s_out], dim=-1)  # (B, T, D*3)
+        gates = F.softmax(self.tier_gate(gate_input), dim=-1)   # (B, T, 3)
+        combined = (gates[..., 0:1] * w_out +
+                    gates[..., 1:2] * e_out +
+                    gates[..., 2:3] * s_out)
+
+        # Project and residual
+        out = self.out_proj(self.v_proj(x) + combined)
+        out = self.norm(out)
+        x = x + self.dropout(out)
+
+        stats = {
+            'mem_w_gate': gates[..., 0].mean(),
+            'mem_e_gate': gates[..., 1].mean(),
+            'mem_s_gate': gates[..., 2].mean(),
+        }
+        return x, stats
+
+
+# ─── Gated FFN (SwiGLU) ─────────────────────────────────────────────────────
+
+class GatedFFN(nn.Module):
+    """SwiGLU feed-forward network."""
+
+    def __init__(self, hidden_dim: int, ff_dim: int, dropout: float = 0.1):
+        super().__init__()
+        self.gate_proj = nn.Linear(hidden_dim, ff_dim, bias=False)
+        self.up_proj = nn.Linear(hidden_dim, ff_dim, bias=False)
+        self.down_proj = nn.Linear(ff_dim, hidden_dim, bias=False)
+        self.norm = nn.LayerNorm(hidden_dim)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+        gate = F.silu(self.gate_proj(x))
+        up = self.up_proj(x)
+        h = gate * up
+        h = self.down_proj(h)
+        h = self.norm(h)
+        return residual + self.dropout(h)
+
+
+# ─── Adaptive Computation Block ──────────────────────────────────────────────
+
+class AdaptiveComputationBlock(nn.Module):
+    """Variable-depth processing: each token may take 1..max_adaptive_steps."""
+
+    def __init__(self, hidden_dim: int, ff_dim: int, max_adaptive_steps: int,
+                 dropout: float = 0.1):
+        super().__init__()
+        self.max_steps = max_adaptive_steps
+        self.layers = nn.ModuleList([
+            GatedFFN(hidden_dim, ff_dim, dropout) for _ in range(max_adaptive_steps)
+        ])
+        self.halt_prob = nn.Linear(hidden_dim, 1, bias=False)
+        self.norm = nn.LayerNorm(hidden_dim)
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+        B, T, D = x.shape
+        output = torch.zeros_like(x)
+        total_weight = torch.zeros(B, T, 1, device=x.device)
+
+        stats = {'avg_steps': torch.tensor(0.0, device=x.device)}
+
+        for step_idx in range(self.max_steps):
+            x = self.layers[step_idx](x)
+
+            # Halting probability
+            p = torch.sigmoid(self.halt_prob(x))  # (B, T, 1)
+
+            # BUG FIX: clamp to avoid going over 1.0
+            remaining = 1.0 - total_weight
+            # Compute max allowed p (leave at least 0.01 for remaining steps)
+            steps_left = self.max_steps - step_idx
+            min_remaining = 0.01 * max(steps_left - 1, 0)
+            max_val = torch.clamp(remaining - min_remaining, min=0.01)
+            p = torch.clamp(p, min=torch.tensor(0.01, device=x.device), max=max_val)
+
+            # On last step, use all remaining weight
+            if step_idx == self.max_steps - 1:
+                p = torch.clamp(remaining, min=0.01)
+
+            output = output + p * x
+            total_weight = total_weight + p
+
+        output = self.norm(output)
+
+        avg_steps = torch.tensor(float(self.max_steps), device=x.device)
+        stats['avg_steps'] = avg_steps
+        return output, stats
+
+
+# ─── Compositional Reasoner ─────────────────────────────────────────────────
+
+class CompositionalReasoner(nn.Module):
+    """Hyperdimensional computing for role-filler binding."""
+
+    def __init__(self, hidden_dim: int, key_dim: int, dropout: float = 0.1):
+        super().__init__()
+        self.role_proj = nn.Linear(hidden_dim, key_dim, bias=False)
+        self.filler_proj = nn.Linear(hidden_dim, key_dim, bias=False)
+        self.unbind_proj = nn.Linear(key_dim, hidden_dim, bias=False)
+        self.norm = nn.LayerNorm(hidden_dim)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+        roles = self.role_proj(x)     # (B, T, K)
+        fillers = self.filler_proj(x) # (B, T, K)
+
+        # Circular convolution as binding operation (element-wise multiply in frequency domain)
+        bound = roles * fillers  # (B, T, K) — simplified HDC binding
+
+        # Shift-based unbinding for positional awareness
+        bound_shifted = torch.roll(bound, shifts=1, dims=1)
+        composed = bound + bound_shifted
+
+        out = self.unbind_proj(composed)
+        out = self.norm(out)
+        return residual + self.dropout(out)
+
+
+# ─── Cognitive Router ────────────────────────────────────────────────────────
+
+class CognitiveRouter(nn.Module):
+    """Routes tokens to channels based on coherence scores."""
+
+    def __init__(self, hidden_dim: int, num_channels: int, channel_dim: int,
+                 routing_iters: int = 1):
+        super().__init__()
+        self.num_channels = num_channels
+        self.channel_dim = channel_dim
+        self.coherence_router = CoherenceRouter(hidden_dim, num_channels, routing_iters)
+
+        # Per-channel projections
+        self.to_channels = nn.Linear(hidden_dim, num_channels * channel_dim, bias=False)
+        self.from_channels = nn.Linear(num_channels * channel_dim, hidden_dim, bias=False)
+
+        # Channel processing
+        self.channels = nn.ModuleList([
+            CognitiveChannel(channel_dim, channel_dim * 4) for _ in range(num_channels)
+        ])
+
+        self.norm = nn.LayerNorm(hidden_dim)
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+        B, T, D = x.shape
+
+        # Route
+        routing_weights, channel_masks = self.coherence_router(x)  # (B, T, C)
+
+        # Project to channel space
+        channel_input = self.to_channels(x)  # (B, T, C*CD)
+        channel_input = channel_input.view(B, T, self.num_channels, self.channel_dim)
+
+        # Process each channel
+        channel_outputs = []
+        for c in range(self.num_channels):
+            # Weighted input for this channel
+            w = routing_weights[:, :, c:c+1].unsqueeze(-1)  # (B, T, 1, 1)
+            ch_in = (channel_input[:, :, c, :] * routing_weights[:, :, c:c+1])  # (B, T, CD)
+            ch_out = self.channels[c](ch_in)  # (B, T, CD)
+            channel_outputs.append(ch_out)
+
+        # Combine channels
+        combined = torch.cat(channel_outputs, dim=-1)  # (B, T, C*CD)
+        out = self.from_channels(combined)
+        out = self.norm(out)
+        x = x + out
+
+        stats = {
+            'routing_entropy': -(routing_weights * (routing_weights + 1e-8).log()).sum(-1).mean(),
+        }
+        return x, stats
+
+
+# ─── CogNet Block ────────────────────────────────────────────────────────────
+
+class CogNetBlock(nn.Module):
+    """Router + Memory Read + FFN with residual connections."""
+
+    def __init__(self, hidden_dim: int, num_channels: int, channel_dim: int,
+                 ff_dim: int, key_dim: int, routing_iters: int,
+                 max_adaptive_steps: int,
+                 working_slots: int, episodic_slots: int, semantic_slots: int,
+                 dropout: float = 0.1):
+        super().__init__()
+        self.router = CognitiveRouter(hidden_dim, num_channels, channel_dim, routing_iters)
+        self.memory = SharedHierarchicalMemory(
+            hidden_dim, key_dim, working_slots, episodic_slots, semantic_slots, dropout
+        )
+        self.adaptive_ffn = AdaptiveComputationBlock(
+            hidden_dim, ff_dim, max_adaptive_steps, dropout
+        )
+        self.composer = CompositionalReasoner(hidden_dim, key_dim, dropout)
+        self.norm = nn.LayerNorm(hidden_dim)
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+        stats = {}
+
+        x, r_stats = self.router(x)
+        stats.update(r_stats)
+
+        x, m_stats = self.memory(x)
+        stats.update(m_stats)
+
+        x, a_stats = self.adaptive_ffn(x)
+        stats.update(a_stats)
+
+        x = self.composer(x)
+        x = self.norm(x)
+
+        return x, stats
+
+
+# ─── CogNet1B ────────────────────────────────────────────────────────────────
+
+class CogNet1B(nn.Module):
+    """Non-transformer language model with cognitive routing."""
+
+    def __init__(
+        self,
+        vocab_size: int = 256,
+        hidden_dim: int = 2048,
+        num_blocks: int = 13,
+        num_channels: int = 8,
+        channel_dim: int = 256,
+        ff_dim: int = 4096,
+        routing_iters: int = 1,
+        max_adaptive_steps: int = 2,
+        max_seq_len: int = 2048,
+        working_slots: int = 64,
+        episodic_slots: int = 128,
+        semantic_slots: int = 256,
+        key_dim: int = 256,
+        dropout: float = 0.1,
+    ):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.hidden_dim = hidden_dim
+        self.num_blocks = num_blocks
+        self.num_channels = num_channels
+        self.channel_dim = channel_dim
+        self.ff_dim = ff_dim
+        self.max_seq_len = max_seq_len
+
+        # Encoder
+        self.encoder = TokenEncoder(vocab_size, hidden_dim, max_seq_len, dropout)
+
+        # Blocks
+        self.blocks = nn.ModuleList([
+            CogNetBlock(
+                hidden_dim, num_channels, channel_dim, ff_dim,
+                key_dim, routing_iters, max_adaptive_steps,
+                working_slots, episodic_slots, semantic_slots, dropout
+            )
+            for _ in range(num_blocks)
+        ])
+
+        # Final norm
+        self.final_norm = nn.LayerNorm(hidden_dim)
+
+        # Output head (weight-tied with token embedding)
+        self.output_proj = nn.Linear(hidden_dim, vocab_size, bias=False)
+        self.output_proj.weight = self.encoder.token_emb.weight
+
+        # Initialize
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.LayerNorm):
+            torch.nn.init.ones_(module.weight)
+            torch.nn.init.zeros_(module.bias)
+
+    def forward(self, input_ids: torch.Tensor,
+                return_stats: bool = False) -> Dict[str, torch.Tensor]:
+        """
+        Args:
+            input_ids: (B, T) integer token ids
+            return_stats: whether to collect intermediate statistics
+        Returns:
+            dict with 'logits' (B, T, vocab_size) and optional 'stats'
+        """
+        x = self.encoder(input_ids)
+
+        all_stats = {} if return_stats else None
+
+        for i, block in enumerate(self.blocks):
+            x, block_stats = block(x)
+            if return_stats:
+                for k, v in block_stats.items():
+                    key = f'block{i}_{k}'
+                    # BUG FIX: clamp NaN/Inf in stats
+                    if isinstance(v, torch.Tensor):
+                        v = v.detach().float()
+                        if torch.isnan(v) or torch.isinf(v):
+                            v = torch.tensor(0.0)
+                    all_stats[key] = v
+
+        x = self.final_norm(x)
+        logits = self.output_proj(x)
+
+        result = {'logits': logits}
+        if return_stats:
+            result['stats'] = all_stats
+        return result
+
+    @torch.no_grad()
+    def generate(self, input_ids: torch.Tensor, max_new_tokens: int = 50,
+                 temperature: float = 1.0, top_k: int = 0,
+                 ) -> torch.Tensor:
+        """Autoregressive generation."""
+        self.eval()
+        for _ in range(max_new_tokens):
+            # Crop to max_seq_len
+            idx = input_ids[:, -self.max_seq_len:]
+            result = self(idx)
+            logits = result['logits'][:, -1, :] / max(temperature, 1e-8)
+
+            # Top-k filtering
+            if top_k > 0:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = float('-inf')
+
+            probs = F.softmax(logits, dim=-1)
+            next_token = torch.multinomial(probs, num_samples=1)
+            input_ids = torch.cat([input_ids, next_token], dim=1)
+
+        return input_ids
+
+    def count_parameters(self) -> Dict[str, int]:
+        total = sum(p.numel() for p in self.parameters())
+        trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
+        return {'total': total, 'trainable': trainable}
+
+    def get_complexity_analysis(self) -> Dict[str, str]:
+        return {
+            'architecture': 'CogNet (Non-Transformer)',
+            'routing': f'O(n) coherence routing x {self.num_channels} channels',
+            'memory': '3-tier hierarchical (Working/Episodic/Semantic)',
+            'attention': 'None (replaced by cognitive routing + memory)',
+            'ffn': 'SwiGLU with adaptive computation',
+            'composition': 'Hyperdimensional role-filler binding',
+            'sequence_complexity': 'O(n) per layer (vs O(n^2) for transformers)',
+            'params': f'{self.count_parameters()["total"]:,}',
+        }
+
+
+# ─── Factory Functions ───────────────────────────────────────────────────────
+
+def create_cognet_1b_small(vocab_size: int = 256, max_seq_len: int = 2048,
+                           dropout: float = 0.1) -> CogNet1B:
+    """Create ~87M parameter model."""
+    return CogNet1B(
+        vocab_size=vocab_size,
+        hidden_dim=1024,
+        num_blocks=8,
+        num_channels=8,
+        channel_dim=128,
+        ff_dim=2048,
+        routing_iters=1,
+        max_adaptive_steps=2,
+        max_seq_len=max_seq_len,
+        working_slots=32,
+        episodic_slots=64,
+        semantic_slots=128,
+        key_dim=256,
+        dropout=dropout,
+    )
+
+
+def create_cognet_1b(vocab_size: int = 256, max_seq_len: int = 2048,
+                     dropout: float = 0.1) -> CogNet1B:
+    """Create ~1B parameter model."""
+    return CogNet1B(
+        vocab_size=vocab_size,
+        hidden_dim=2048,
+        num_blocks=13,
+        num_channels=8,
+        channel_dim=256,
+        ff_dim=4096,
+        routing_iters=1,
+        max_adaptive_steps=2,
+        max_seq_len=max_seq_len,
+        working_slots=64,
+        episodic_slots=128,
+        semantic_slots=256,
+        key_dim=256,
+        dropout=dropout,
+    )
+
+
+# ─── Self-Test ───────────────────────────────────────────────────────────────
+
+if __name__ == '__main__':
+    print("=" * 60)
+    print("CogNet1B Self-Test")
+    print("=" * 60)
+
+    # Small model for quick test
+    model = CogNet1B(
+        vocab_size=128,
+        hidden_dim=128,
+        num_blocks=2,
+        num_channels=4,
+        channel_dim=32,
+        ff_dim=256,
+        routing_iters=1,
+        max_adaptive_steps=2,
+        max_seq_len=64,
+        working_slots=8,
+        episodic_slots=16,
+        semantic_slots=32,
+        key_dim=64,
+        dropout=0.1,
+    )
+
+    params = model.count_parameters()
+    print(f"\nParameters: {params['total']:,} total, {params['trainable']:,} trainable")
+
+    # Forward pass
+    x = torch.randint(0, 128, (2, 16))
+    result = model(x, return_stats=True)
+    logits = result['logits']
+    print(f"Input shape: {x.shape}")
+    print(f"Output logits shape: {logits.shape}")
+    print(f"Stats keys: {len(result.get('stats', {}))}")
+
+    # Backward pass
+    loss = logits.sum()
+    loss.backward()
+    print("Backward pass OK")
+
+    # Generate test
+    gen = model.generate(x[:, :4], max_new_tokens=8, temperature=0.8, top_k=10)
+    print(f"Generated shape: {gen.shape}")
+
+    # Complexity analysis
+    analysis = model.get_complexity_analysis()
+    for k, v in analysis.items():
+        print(f"  {k}: {v}")
+
+    print("\n✓ All self-tests passed!")
+
+    # Test factory functions
+    small = create_cognet_1b_small(vocab_size=128, max_seq_len=64)
+    print(f"\nSmall model params: {small.count_parameters()['total']:,}")