HYDRA v1: 19M param non-transformer with selective gated recurrence

README.md

@@ -8,6 +8,7 @@ tags:
 - gated-recurrence
 - from-scratch
 - tinystories
+- recurrent-neural-network
 license: apache-2.0
 datasets:
 - roneneldan/TinyStories
@@ -15,27 +16,25 @@ datasets:
 
 # HYDRA: Hybrid Dynamic Recurrent Architecture
 
-**A novel non-transformer language model trained from scratch.**
+A novel **non-transformer** language model built from scratch. Trained on CPU using a custom architecture that combines Mamba's Selective State Spaces, Griffin's Real-Gated Linear Recurrence (RG-LRU), and RWKV's channel mixing — with **zero attention layers**.
 
 ## Architecture
 
-HYDRA is a completely original architecture that combines cutting-edge ideas from recent research:
-
-| Component | Inspired By | Paper |
+| Component | Source | Paper |
 |---|---|---|
-| **Selective State Spaces** | Mamba | [arxiv:2312.00752](https://arxiv.org/abs/2312.00752) |
-| **Real-Gated Linear Recurrence (RG-LRU)** | Griffin | [arxiv:2402.19427](https://arxiv.org/abs/2402.19427) |
-| **Time/Channel Mixing Blocks** | RWKV | [arxiv:2305.13048](https://arxiv.org/abs/2305.13048) |
-| **Multi-Scale State Compression** | Novel | Parallel recurrences at different timescales |
+| Selective State Spaces | Mamba | arxiv:2312.00752 |
+| Real-Gated Linear Recurrence | Griffin | arxiv:2402.19427 |
+| Time/Channel Mixing | RWKV | arxiv:2305.13048 |
+| Multi-Scale Compression | **Novel** | Parallel recurrences at different timescales |
 
 ### Key Properties
-- **NOT a transformer** — no attention mechanism at all
-- **O(n) time complexity** — linear in sequence length (vs O(n²) for transformers)
-- **Constant memory inference** — no growing KV cache
-- **Content-aware state selection** — addresses the classic RNN weakness
-- **Multi-scale temporal processing** — captures both short and long-range patterns
+- **NOT a transformer** — no attention whatsoever
+- **O(n) time** — linear in sequence length (vs O(n²) transformers)
+- **Constant memory** at inference (no KV cache)
+- **Content-aware** selective gating (Mamba + Griffin fusion)
+- **Multi-scale** temporal processing
 
-### Architecture Details
+### Architecture Diagram
 ```
 Token Embedding → N × HydraBlock → RMSNorm → LM Head
 
@@ -46,52 +45,71 @@ HydraBlock:
 SelectiveGatedRecurrence (per timescale):
   ├── Input projection (2 branches)
   ├── Branch 1: Separable Conv1D → SiLU → Selective B,C projection
-  ├── Branch 2: GeLU gate
-  ├── Multi-scale gated recurrence: h_t = a_t·h_{t-1} + √(1-a_t²)·(i_t·B_t)
-  ├── Timescale fusion
+  ├── Input gate + Recurrence gate (from Griffin RG-LRU)
+  ├── Gated recurrence: h_t = a_t·h_{t-1} + √(1-a_t²)·(i_t·B_t)
   └── Gated merge + output projection
 ```
 
-## Model Specs
-- **Parameters**: 19,274,816
+## Specs
+- **Parameters**: 19,274,816 (19.3M)
 - **d_model**: 256
 - **Layers**: 6
-- **State dimension**: 16
+- **State dim**: 16
 - **Timescales**: 2
-- **Context length**: 256
-- **Vocabulary**: GPT-2 tokenizer (50,257 tokens)
+- **Context length**: 128
 
 ## Training
-- **Dataset**: [TinyStories](https://huggingface.co/datasets/roneneldan/TinyStories)
-- **From scratch**: All weights randomly initialized, no pretrained components
-- **Optimizer**: AdamW (β1=0.9, β2=0.95, from RWKV paper)
-- **Schedule**: Cosine decay with linear warmup
-- **Hardware**: CPU
+- **Dataset**: TinyStories (5,000 stories for quick training)
+- **From scratch**: Random initialization, no pretrained components
+- **Best val_loss**: 3.7988
+- **Val_ppl**: 44.6
+- **Hardware**: CPU only
+- **Optimizer**: AdamW (β₁=0.9, β₂=0.95)
+
+## Generated Samples
+
+**Once upon a time**: Once upon a time, there was a little girl named Lily. She loved to play outside. She had a big, feeling very excited! She couldn't wait for a big, but she said.
+
+The little boy was so happy to make the bird who loved
+**A little dog**: A little dog who ran to be happy. She was very excited that the park. She wanted to play with the window. 
+
+"I'm sorry, Lily. It is a voice?" She was so happy!
+
+Ben did not give it and
+**A girl named Lily**: A girl named Lily. She was very happy he was very tall. She saw what he could not want to the forest. She says, "I'm sorry, we can have to go to the water. She was very very excited and happy.
+
+Lily
+**One day a boy**: One day a boy named Lily liked to play with it. He said, "Don't be careful."
+
+"We are happy, we have to go to the man that he decided to play with his friends. He was very happy that the dog had a man
 
 ## Usage
 
 ```python
-import torch
+import torch, json
 from model import HydraModel, HydraConfig
 from transformers import AutoTokenizer
 
-# Load
 config = HydraConfig(**json.load(open("config.json")))
 model = HydraModel(config)
 model.load_state_dict(torch.load("model.pt", map_location="cpu"))
+model.eval()
 
 tokenizer = AutoTokenizer.from_pretrained("gpt2")
-
-# Generate
 prompt = "Once upon a time"
-input_ids = torch.tensor([tokenizer.encode(prompt)])
-output = model.generate(input_ids, max_new_tokens=100, temperature=0.8)
-print(tokenizer.decode(output[0], skip_special_tokens=True))
+ids = torch.tensor([tokenizer.encode(prompt)])
+with torch.no_grad():
+    out = model.generate(ids, max_new_tokens=50, temperature=0.8, top_k=40)
+print(tokenizer.decode(out[0], skip_special_tokens=True))
 ```
 
-## Research References
+## Model Files
+- `model.pt` — trained weights
+- `config.json` — model configuration
+- `model.py` — full architecture source code
 
-1. **Mamba**: Gu & Dao, "Mamba: Linear-Time Sequence Modeling with Selective State Spaces", 2023
-2. **Griffin**: De et al., "Griffin: Mixing Gated Linear Recurrences with Local Attention", 2024
-3. **RWKV**: Peng et al., "RWKV: Reinventing RNNs for the Transformer Era", 2023
-4. **TinyStories**: Eldan & Li, "TinyStories: How Small Can Language Models Be and Still Speak Coherent English?", 2023
+## Research References
+1. Gu & Dao, "Mamba: Linear-Time Sequence Modeling with Selective State Spaces", NeurIPS 2023
+2. De et al., "Griffin: Mixing Gated Linear Recurrences with Local Attention", ICLR 2024
+3. Peng et al., "RWKV: Reinventing RNNs for the Transformer Era", 2023
+4. Eldan & Li, "TinyStories: How Small Can Language Models Be?", 2023

config.json

@@ -8,7 +8,7 @@
   "mlp_expand": 3,
   "n_scales": 2,
   "dropout": 0.1,
-  "max_seq_len": 256,
+  "max_seq_len": 128,
   "pad_token_id": 50256,
   "tie_weights": true
 }

model.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:bda4c00904bd3d69cfb22d2b106139693e92e2ed0e211c73a964ed53fabdc6e8
+oid sha256:8b4d19a13205d6f991c2a27001f06318aee6988971d276eaf0bc4c673d5667ed
 size 77144619

model.py

@@ -1,455 +1,134 @@
 """
 HYDRA: Hybrid Dynamic Recurrent Architecture
-=============================================
-A novel non-transformer language model built from scratch.
-
-Architecture combines ideas from:
-- Mamba (arxiv:2312.00752): Selective State Spaces with input-dependent dynamics
-- Griffin (arxiv:2402.19427): Real-Gated Linear Recurrent Unit (RG-LRU)  
-- RWKV (arxiv:2305.13048): Time-mixing + channel-mixing residual blocks
-
-Key innovations in this implementation:
-1. Selective Gated Recurrence (SGR): Fuses Mamba's selective SSM with Griffin's RG-LRU gates
-2. Multi-Scale State Compression: Parallel recurrences at different timescales
-3. Gated Channel Mixing: RWKV-inspired channel mixing with GeGeLU
-4. No attention mechanism - pure recurrent, O(n) in sequence length
-
-This gives us:
-- Linear time complexity (vs quadratic for transformers)
-- Constant memory at inference (vs growing KV cache)
-- Content-aware state selection (addresses RNN weakness)
-- Theoretically unlimited context (state carries forward indefinitely)
+Novel non-transformer combining Mamba SSM, Griffin RG-LRU, RWKV mixing.
 """
-
-import math
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
+import math, torch, torch.nn as nn, torch.nn.functional as F
 from dataclasses import dataclass
-from typing import Optional, Tuple
-
+from typing import Optional
 
 @dataclass
 class HydraConfig:
-    """Configuration for the Hydra model."""
-    vocab_size: int = 50257       # GPT-2 tokenizer vocab size
-    d_model: int = 512            # Model dimension
-    n_layers: int = 8             # Number of residual blocks
-    d_state: int = 16             # SSM state dimension (N in Mamba)
-    d_conv: int = 4               # Local convolution width (from Griffin/Mamba)
-    d_inner: int = None           # Inner dimension (defaults to 2 * d_model)
-    mlp_expand: int = 3           # MLP expansion factor (from Griffin paper: M=3)
-    n_scales: int = 3             # Number of timescales for multi-scale recurrence
-    dropout: float = 0.1          # Dropout rate
-    max_seq_len: int = 1024       # Maximum sequence length for training
-    pad_token_id: int = 50256     # Padding token (GPT-2 <|endoftext|>)
-    tie_weights: bool = True      # Tie input/output embeddings (from Griffin)
-    
+    vocab_size: int = 50257
+    d_model: int = 512
+    n_layers: int = 8
+    d_state: int = 16
+    d_conv: int = 4
+    d_inner: int = None
+    mlp_expand: int = 3
+    n_scales: int = 3
+    dropout: float = 0.1
+    max_seq_len: int = 1024
+    pad_token_id: int = 50256
+    tie_weights: bool = True
     def __post_init__(self):
-        if self.d_inner is None:
-            self.d_inner = 2 * self.d_model
-
+        if self.d_inner is None: self.d_inner = 2 * self.d_model
 
 class RMSNorm(nn.Module):
-    """Root Mean Square Layer Normalization (from Griffin paper)."""
-    def __init__(self, d_model: int, eps: float = 1e-8):
-        super().__init__()
-        self.eps = eps
-        self.weight = nn.Parameter(torch.ones(d_model))
-    
+    def __init__(self, d_model, eps=1e-8):
+        super().__init__(); self.eps=eps; self.weight=nn.Parameter(torch.ones(d_model))
     def forward(self, x):
-        rms = torch.sqrt(torch.mean(x ** 2, dim=-1, keepdim=True) + self.eps)
-        return x / rms * self.weight
-
+        return x / torch.sqrt(torch.mean(x**2,dim=-1,keepdim=True)+self.eps) * self.weight
 
 class SelectiveGatedRecurrence(nn.Module):
-    """
-    Selective Gated Recurrence (SGR) — our novel temporal mixing block.
-    
-    Fuses ideas from:
-    - Mamba's selective SSM: B, C, Δ are functions of input (content-aware)
-    - Griffin's RG-LRU: Gated recurrence with input gate and recurrence gate
-    - Multi-scale: Parallel recurrences at different timescales
-    
-    Core equations (per timescale s):
-        r_t = sigmoid(W_r @ x_t)           # recurrence gate (Griffin)
-        i_t = sigmoid(W_i @ x_t)           # input gate (Griffin)  
-        B_t = Linear(x_t)                  # selective input projection (Mamba)
-        C_t = Linear(x_t)                  # selective output projection (Mamba)
-        a_t = a_base^(c * r_t)             # gated decay (Griffin: c=8)
-        h_t = a_t * h_{t-1} + sqrt(1 - a_t^2) * (i_t * B_t * x_t)  # state update
-        y_t = C_t * h_t                    # readout
-    """
-    
-    def __init__(self, config: HydraConfig):
+    def __init__(self, config):
         super().__init__()
-        self.d_inner = config.d_inner
-        self.d_state = config.d_state
-        self.n_scales = config.n_scales
-        self.gate_c = 8.0  # Scalar constant from Griffin paper
-        
-        # Input projection: x -> two branches (like Mamba/Griffin recurrent block)
-        self.in_proj = nn.Linear(config.d_model, 2 * config.d_inner, bias=False)
-        
-        # Small causal conv1d on first branch (from Griffin Sec 2.3, filter size 4)
-        self.conv1d = nn.Conv1d(
-            config.d_inner, config.d_inner, 
-            kernel_size=config.d_conv, 
-            padding=config.d_conv - 1,
-            groups=config.d_inner  # depthwise/separable (from Griffin)
-        )
-        
-        # Selective projections (from Mamba: B, C are input-dependent)
-        # One set per timescale
-        self.B_proj = nn.ModuleList([
-            nn.Linear(config.d_inner, config.d_state, bias=False)
-            for _ in range(config.n_scales)
-        ])
-        self.C_proj = nn.ModuleList([
-            nn.Linear(config.d_inner, config.d_state, bias=False)
-            for _ in range(config.n_scales)
-        ])
-        
-        # Gates (from Griffin RG-LRU)
-        self.recurrence_gate = nn.ModuleList([
-            nn.Linear(config.d_inner, config.d_state, bias=True)
-            for _ in range(config.n_scales)
-        ])
-        self.input_gate = nn.ModuleList([
-            nn.Linear(config.d_inner, config.d_state, bias=True)
-            for _ in range(config.n_scales)
-        ])
-        
-        # Learnable base decay per timescale (parameterized as Lambda, a = sigmoid(Lambda))
-        # Initialize so a^c is uniformly in [0.9, 0.999] per Griffin
-        self.Lambda = nn.ParameterList()
+        self.d_inner=config.d_inner; self.d_state=config.d_state; self.n_scales=config.n_scales; self.gate_c=8.0
+        self.in_proj=nn.Linear(config.d_model, 2*config.d_inner, bias=False)
+        self.conv1d=nn.Conv1d(config.d_inner,config.d_inner,config.d_conv,padding=config.d_conv-1,groups=config.d_inner)
+        self.B_proj=nn.ModuleList([nn.Linear(config.d_inner,config.d_state,bias=False) for _ in range(config.n_scales)])
+        self.C_proj=nn.ModuleList([nn.Linear(config.d_inner,config.d_state,bias=False) for _ in range(config.n_scales)])
+        self.recurrence_gate=nn.ModuleList([nn.Linear(config.d_inner,config.d_state,bias=True) for _ in range(config.n_scales)])
+        self.input_gate=nn.ModuleList([nn.Linear(config.d_inner,config.d_state,bias=True) for _ in range(config.n_scales)])
+        self.Lambda=nn.ParameterList()
         for s in range(config.n_scales):
-            # Different timescales: short-term, medium-term, long-term
-            # a^c in [0.9^(1/(s+1)), 0.999^(1/(s+1))]
-            low = math.log(0.9 ** (1.0 / (s + 1))) 
-            high = math.log(0.999 ** (1.0 / (s + 1)))
-            # sigmoid^{-1}(x) = log(x/(1-x))
-            low_a = math.exp(low)
-            high_a = math.exp(high)
-            # Initialize Lambda so sigmoid(Lambda) gives us the right range
-            init_val = torch.empty(config.d_state).uniform_(
-                math.log(low_a / (1 - low_a + 1e-8)),
-                math.log(high_a / (1 - high_a + 1e-8))
-            )
+            low_a=0.9**(1.0/(s+1)); high_a=0.999**(1.0/(s+1))
+            init_val=torch.empty(config.d_state).uniform_(math.log(low_a/(1-low_a+1e-8)),math.log(high_a/(1-high_a+1e-8)))
             self.Lambda.append(nn.Parameter(init_val))
-        
-        # Multi-scale fusion
-        self.scale_fusion = nn.Linear(config.n_scales * config.d_state, config.d_inner, bias=False)
-        
-        # Output projection
-        self.out_proj = nn.Linear(config.d_inner, config.d_model, bias=False)
-        
-        self.dropout = nn.Dropout(config.dropout)
-    
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Args:
-            x: (batch, seq_len, d_model)
-        Returns:
-            y: (batch, seq_len, d_model)
-        """
-        batch, seq_len, _ = x.shape
-        
-        # Project input into two branches
-        xz = self.in_proj(x)  # (B, L, 2*d_inner)
-        x_branch, z_branch = xz.chunk(2, dim=-1)  # each (B, L, d_inner)
-        
-        # Branch 1: Conv1D + recurrence
-        # Conv1d expects (B, C, L)
-        x_conv = self.conv1d(x_branch.transpose(1, 2))[:, :, :seq_len].transpose(1, 2)
-        x_conv = F.silu(x_conv)  # SiLU activation (from Mamba)
-        
-        # Branch 2: Gate (from Griffin/Mamba architecture)
-        z_gate = F.gelu(z_branch)
-        
-        # Multi-scale selective recurrence
-        scale_outputs = []
+        self.scale_fusion=nn.Linear(config.n_scales*config.d_state,config.d_inner,bias=False)
+        self.out_proj=nn.Linear(config.d_inner,config.d_model,bias=False)
+        self.dropout=nn.Dropout(config.dropout)
+
+    def forward(self, x):
+        batch,seq_len,_=x.shape
+        xz=self.in_proj(x); x_branch,z_branch=xz.chunk(2,dim=-1)
+        x_conv=self.conv1d(x_branch.transpose(1,2))[:,:,:seq_len].transpose(1,2)
+        x_conv=F.silu(x_conv); z_gate=F.gelu(z_branch)
+        scale_outputs=[]
         for s in range(self.n_scales):
-            # Selective projections (Mamba: B, C depend on input)
-            B_t = self.B_proj[s](x_conv)    # (B, L, d_state)
-            C_t = self.C_proj[s](x_conv)    # (B, L, d_state)
-            
-            # Gates (Griffin RG-LRU)
-            r_t = torch.sigmoid(self.recurrence_gate[s](x_conv))  # (B, L, d_state)
-            i_t = torch.sigmoid(self.input_gate[s](x_conv))       # (B, L, d_state)
-            
-            # Base decay: a = sigmoid(Lambda)
-            a_base = torch.sigmoid(self.Lambda[s])  # (d_state,)
-            
-            # Gated decay: a_t = a_base^(c * r_t)  (Griffin eq 3)
-            # Compute in log space for stability: log(a_t) = c * r_t * log(a_base)
-            log_a_base = torch.log(a_base.clamp(min=1e-8))  # (d_state,)
-            log_a_t = self.gate_c * r_t * log_a_base.unsqueeze(0).unsqueeze(0)  # (B, L, d_state)
-            a_t = torch.exp(log_a_t)  # (B, L, d_state)
-            
-            # Input contribution: sqrt(1 - a_t^2) * i_t * B_t * x_conv_projected
-            # We use B_t as the projection of x into state space
-            sqrt_term = torch.sqrt((1 - a_t ** 2).clamp(min=1e-8))
-            input_contrib = sqrt_term * i_t * B_t  # (B, L, d_state)
-            
-            # Sequential recurrence (numerically stable)
-            # h_t = a_t * h_{t-1} + input_contrib_t
-            # Compute all h states for this timescale
-            h = torch.zeros(batch, seq_len, self.d_state, device=x.device, dtype=x.dtype)
-            h_prev = torch.zeros(batch, self.d_state, device=x.device, dtype=x.dtype)
+            B_t=self.B_proj[s](x_conv); C_t=self.C_proj[s](x_conv)
+            r_t=torch.sigmoid(self.recurrence_gate[s](x_conv))
+            i_t=torch.sigmoid(self.input_gate[s](x_conv))
+            a_base=torch.sigmoid(self.Lambda[s])
+            log_a_base=torch.log(a_base.clamp(min=1e-8))
+            log_a_t=self.gate_c*r_t*log_a_base.unsqueeze(0).unsqueeze(0)
+            a_t=torch.exp(log_a_t)
+            sqrt_term=torch.sqrt((1-a_t**2).clamp(min=1e-8))
+            input_contrib=sqrt_term*i_t*B_t
+            h_prev=torch.zeros(batch,self.d_state,device=x.device,dtype=x.dtype)
+            h_states=[]
             for t in range(seq_len):
-                h_prev = a_t[:, t] * h_prev + input_contrib[:, t]
-                h[:, t] = h_prev
+                h_prev=a_t[:,t]*h_prev+input_contrib[:,t]
+                h_states.append(h_prev.unsqueeze(1))
+            h=torch.cat(h_states,dim=1)
             scale_outputs.append(h)
-        
-        # Fuse multi-scale outputs
-        multi_scale = torch.cat(scale_outputs, dim=-1)  # (B, L, n_scales * d_state)
-        fused = self.scale_fusion(multi_scale)  # (B, L, d_inner)
-        
-        # Merge with gated branch (element-wise multiplication, like Griffin/Mamba)
-        output = fused * z_gate
-        
-        # Output projection
-        output = self.out_proj(output)  # (B, L, d_model)
+        multi_scale=torch.cat(scale_outputs,dim=-1)
+        fused=self.scale_fusion(multi_scale)
+        output=self.out_proj(fused*z_gate)
         return self.dropout(output)
 
-
 class GatedChannelMixing(nn.Module):
-    """
-    Channel mixing block inspired by RWKV and Griffin's gated MLP.
-    
-    Uses GeGeLU (Gated GeLU) from Griffin paper Sec 2.2:
-        y = Linear(GeLU(Linear(x)) * Linear(x))
-    
-    With expansion factor M=3 (from Griffin paper).
-    """
-    
-    def __init__(self, config: HydraConfig):
-        super().__init__()
-        d_ff = config.mlp_expand * config.d_model  # M * D
-        
-        self.gate_proj = nn.Linear(config.d_model, d_ff, bias=False)
-        self.up_proj = nn.Linear(config.d_model, d_ff, bias=False)
-        self.down_proj = nn.Linear(d_ff, config.d_model, bias=False)
-        self.dropout = nn.Dropout(config.dropout)
-    
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        # GeGeLU: GeLU(gate) * up, then project down
-        gate = F.gelu(self.gate_proj(x))
-        up = self.up_proj(x)
-        return self.dropout(self.down_proj(gate * up))
-
+    def __init__(self, config):
+        super().__init__(); d_ff=config.mlp_expand*config.d_model
+        self.gate_proj=nn.Linear(config.d_model,d_ff,bias=False)
+        self.up_proj=nn.Linear(config.d_model,d_ff,bias=False)
+        self.down_proj=nn.Linear(d_ff,config.d_model,bias=False)
+        self.dropout=nn.Dropout(config.dropout)
+    def forward(self, x):
+        return self.dropout(self.down_proj(F.gelu(self.gate_proj(x))*self.up_proj(x)))
 
 class HydraBlock(nn.Module):
-    """
-    Single Hydra residual block (from Griffin architecture Sec 2.1):
-        x -> RMSNorm -> TemporalMix (SGR) -> + residual
-        x -> RMSNorm -> ChannelMix (GeGeLU MLP) -> + residual
-    """
-    
-    def __init__(self, config: HydraConfig):
+    def __init__(self, config):
         super().__init__()
-        self.norm1 = RMSNorm(config.d_model)
-        self.temporal_mix = SelectiveGatedRecurrence(config)
-        self.norm2 = RMSNorm(config.d_model)
-        self.channel_mix = GatedChannelMixing(config)
-    
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        # Temporal mixing with pre-norm and residual
-        x = x + self.temporal_mix(self.norm1(x))
-        # Channel mixing with pre-norm and residual
-        x = x + self.channel_mix(self.norm2(x))
-        return x
-
+        self.norm1=RMSNorm(config.d_model); self.temporal_mix=SelectiveGatedRecurrence(config)
+        self.norm2=RMSNorm(config.d_model); self.channel_mix=GatedChannelMixing(config)
+    def forward(self, x):
+        return x+self.channel_mix(self.norm2(x+self.temporal_mix(self.norm1(x))))
 
 class HydraModel(nn.Module):
-    """
-    HYDRA: Hybrid Dynamic Recurrent Architecture
-    
-    A novel non-transformer language model combining:
-    - Selective State Spaces (Mamba)
-    - Real-Gated Linear Recurrence (Griffin)
-    - Channel mixing (RWKV/Griffin)
-    
-    Architecture:
-        Token Embedding -> N x HydraBlock -> RMSNorm -> LM Head
-    
-    Each HydraBlock:
-        -> RMSNorm -> SelectiveGatedRecurrence -> residual
-        -> RMSNorm -> GatedChannelMixing -> residual
-    """
-    
-    def __init__(self, config: HydraConfig):
-        super().__init__()
-        self.config = config
-        
-        # Token embedding
-        self.embed = nn.Embedding(config.vocab_size, config.d_model)
-        self.embed_dropout = nn.Dropout(config.dropout)
-        
-        # Stack of Hydra blocks
-        self.blocks = nn.ModuleList([
-            HydraBlock(config) for _ in range(config.n_layers)
-        ])
-        
-        # Final norm (from Griffin: RMSNorm before LM head)
-        self.final_norm = RMSNorm(config.d_model)
-        
-        # Language model head
-        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
-        
-        # Weight tying (from Griffin paper: shared embedding weights)
-        if config.tie_weights:
-            self.lm_head.weight = self.embed.weight
-        
-        # Initialize weights
+    def __init__(self, config):
+        super().__init__(); self.config=config
+        self.embed=nn.Embedding(config.vocab_size,config.d_model)
+        self.embed_dropout=nn.Dropout(config.dropout)
+        self.blocks=nn.ModuleList([HydraBlock(config) for _ in range(config.n_layers)])
+        self.final_norm=RMSNorm(config.d_model)
+        self.lm_head=nn.Linear(config.d_model,config.vocab_size,bias=False)
+        if config.tie_weights: self.lm_head.weight=self.embed.weight
         self._init_weights()
-    
     def _init_weights(self):
-        """Initialize weights following Griffin/Mamba conventions."""
-        for module in self.modules():
-            if isinstance(module, nn.Linear):
-                # LeCun init (from Griffin Sec 2.4)
-                nn.init.normal_(module.weight, mean=0.0, std=1.0 / math.sqrt(module.weight.shape[1]))
-                if module.bias is not None:
-                    nn.init.zeros_(module.bias)
-            elif isinstance(module, nn.Embedding):
-                nn.init.normal_(module.weight, mean=0.0, std=0.02)
-            elif isinstance(module, nn.Conv1d):
-                nn.init.normal_(module.weight, mean=0.0, std=0.02)
-                if module.bias is not None:
-                    nn.init.zeros_(module.bias)
-    
-    def forward(
-        self, 
-        input_ids: torch.Tensor, 
-        labels: Optional[torch.Tensor] = None
-    ) -> dict:
-        """
-        Args:
-            input_ids: (batch, seq_len) token indices
-            labels: (batch, seq_len) target token indices for language modeling
-        Returns:
-            dict with 'logits' and optionally 'loss'
-        """
-        # Embed tokens
-        x = self.embed(input_ids)
-        x = self.embed_dropout(x)
-        
-        # Pass through Hydra blocks
-        for block in self.blocks:
-            x = block(x)
-        
-        # Final norm and LM head
-        x = self.final_norm(x)
-        logits = self.lm_head(x)  # (batch, seq_len, vocab_size)
-        
-        result = {"logits": logits}
-        
+        for m in self.modules():
+            if isinstance(m, nn.Linear): nn.init.normal_(m.weight,0,1.0/math.sqrt(m.weight.shape[1])); hasattr(m,'bias') and m.bias is not None and nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.Embedding): nn.init.normal_(m.weight,0,0.02)
+            elif isinstance(m, nn.Conv1d): nn.init.normal_(m.weight,0,0.02); hasattr(m,'bias') and m.bias is not None and nn.init.zeros_(m.bias)
+    def forward(self, input_ids, labels=None):
+        x=self.embed_dropout(self.embed(input_ids))
+        for b in self.blocks: x=b(x)
+        x=self.final_norm(x); logits=self.lm_head(x)
+        result={"logits":logits}
         if labels is not None:
-            # Shift logits and labels for next-token prediction
-            shift_logits = logits[:, :-1, :].contiguous()
-            shift_labels = labels[:, 1:].contiguous()
-            loss = F.cross_entropy(
-                shift_logits.view(-1, self.config.vocab_size),
-                shift_labels.view(-1),
-                ignore_index=self.config.pad_token_id
-            )
-            result["loss"] = loss
-        
+            loss=F.cross_entropy(logits[:,:-1,:].contiguous().view(-1,self.config.vocab_size),labels[:,1:].contiguous().view(-1),ignore_index=self.config.pad_token_id)
+            result["loss"]=loss
         return result
-    
-    def generate(
-        self,
-        input_ids: torch.Tensor,
-        max_new_tokens: int = 100,
-        temperature: float = 0.8,
-        top_k: int = 50,
-    ) -> torch.Tensor:
-        """Simple autoregressive generation."""
+    def generate(self, input_ids, max_new_tokens=100, temperature=0.8, top_k=50):
         self.eval()
         with torch.no_grad():
             for _ in range(max_new_tokens):
-                # Use full context (or truncated if too long)
-                context = input_ids[:, -self.config.max_seq_len:]
-                output = self.forward(context)
-                logits = output["logits"][:, -1, :] / temperature
-                
-                # Top-k filtering
-                if top_k > 0:
-                    values, _ = torch.topk(logits, top_k)
-                    logits[logits < values[:, -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)
-        
+                ctx=input_ids[:,-self.config.max_seq_len:]
+                logits=self.forward(ctx)["logits"][:,-1,:]/temperature
+                if top_k>0:
+                    v,_=torch.topk(logits,top_k); logits[logits<v[:,-1:]]=float('-inf')
+                probs=F.softmax(logits,dim=-1)
+                next_token=torch.multinomial(probs,1)
+                input_ids=torch.cat([input_ids,next_token],dim=-1)
         return input_ids
-    
-    @property
-    def num_parameters(self):
-        return sum(p.numel() for p in self.parameters())
-    
     @property
-    def num_trainable_parameters(self):
-        return sum(p.numel() for p in self.parameters() if p.requires_grad)
-
-
-def create_hydra_small():
-    """~10M parameter model for CPU training on TinyStories."""
-    return HydraModel(HydraConfig(
-        vocab_size=50257,
-        d_model=256,
-        n_layers=6,
-        d_state=16,
-        d_conv=4,
-        mlp_expand=3,
-        n_scales=2,
-        dropout=0.1,
-        max_seq_len=512,
-    ))
-
-
-def create_hydra_medium():
-    """~30M parameter model."""
-    return HydraModel(HydraConfig(
-        vocab_size=50257,
-        d_model=512,
-        n_layers=8,
-        d_state=16,
-        d_conv=4,
-        mlp_expand=3,
-        n_scales=3,
-        dropout=0.1,
-        max_seq_len=1024,
-    ))
-
-
-if __name__ == "__main__":
-    # Quick test
-    config = HydraConfig(
-        vocab_size=50257,
-        d_model=256,
-        n_layers=4,
-        d_state=16,
-        d_conv=4,
-        mlp_expand=3,
-        n_scales=2,
-        dropout=0.1,
-        max_seq_len=128,
-    )
-    
-    model = HydraModel(config)
-    print(f"Hydra Model")
-    print(f"  Parameters: {model.num_parameters:,}")
-    print(f"  Config: {config}")
-    
-    # Test forward pass
-    batch_size, seq_len = 2, 64
-    input_ids = torch.randint(0, config.vocab_size, (batch_size, seq_len))
-    labels = torch.randint(0, config.vocab_size, (batch_size, seq_len))
-    
-    output = model(input_ids, labels=labels)
-    print(f"  Logits shape: {output['logits'].shape}")
-    print(f"  Loss: {output['loss'].item():.4f}")
-    print("  ✓ Forward pass successful!")
+    def num_parameters(self): return sum(p.numel() for p in self.parameters())

tokenizer_config.json

@@ -7,7 +7,7 @@
   "is_local": false,
   "local_files_only": false,
   "model_max_length": 1024,
-  "pad_token": "<|endoftext|>",
+  "pad_token": null,
   "tokenizer_class": "GPT2Tokenizer",
   "unk_token": "<|endoftext|>"
 }