adding sources

model.py

@@ -0,0 +1,162 @@
+#@title Architecture implementation
+# TODO: comment and rename variables / clean code
+
+
+# https://arxiv.org/abs/2410.01201v1
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# appendix B
+# https://github.com/glassroom/heinsen_sequence
+
+def heinsen_associative_scan_log(log_coeffs, log_values):
+    a_star = log_coeffs.cumsum(dim = 1)
+    log_h0_plus_b_star = (log_values - a_star).logcumsumexp(dim = 1)
+    log_h = a_star + log_h0_plus_b_star
+    return log_h.exp()
+
+# appendix B.3
+
+def g(x):     return torch.where(x >= 0, x + 0.5, x.sigmoid())
+def log_g(x): return torch.where(x >= 0, (F.relu(x) + 0.5).log(), -F.softplus(-x))
+
+# log-space version of minGRU - B.3.1
+# they enforce the hidden states to be positive
+
+class minGRU(nn.Module):
+    def __init__(self, d_model, d_inner):
+        super().__init__()
+
+        self.hidden_proj = nn.Linear(d_model, d_inner, bias=False)
+        self.gate_proj   = nn.Linear(d_model, d_inner, bias=False)
+        self.out_proj    = nn.Linear(d_inner, d_model, bias=False)
+
+
+    def step(self, x, h_prev=None):
+        hidden = self.hidden_proj(x)
+        gate   = self.gate_proj(x)
+
+        h_prev = h_prev.detach() if h_prev is not None else None
+
+        hidden = g(hidden)
+        gate   = gate.sigmoid()
+        out    = torch.lerp(h_prev, hidden, gate) if h_prev is not None else (hidden * gate)
+
+        h_next = out[:, -1:]
+        out    = self.out_proj(out)
+
+        return out, h_next
+
+
+    def forward(self, x, h_prev=None):
+        seq_len = x.shape[1]
+        hidden  = self.hidden_proj(x)
+        gate    = self.gate_proj(x)
+
+        h_prev = h_prev.detach() if h_prev is not None else None
+
+        log_coeffs  = -F.softplus(gate)
+        log_z       = -F.softplus(-gate)
+        log_tilde_h = log_g(hidden)
+        log_values  = log_z + log_tilde_h
+
+        if h_prev is not None:
+            log_values = torch.cat((h_prev.log(), log_values), dim=1)
+            log_coeffs = F.pad(log_coeffs, (0, 0, 1, 0))
+
+        out = heinsen_associative_scan_log(log_coeffs, log_values)
+        out = out[:, -seq_len:]
+        
+        h_next = out[:, -1:]
+        out    = self.out_proj(out)
+
+        return out, h_next
+
+
+
+
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, d_model: int, eps: float=1e-5):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(d_model))
+
+    def _norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        output = self._norm(x.float()).type_as(x)
+        return output * self.weight
+
+
+
+
+
+class minGRULM(nn.Module):
+    def __init__(self, vocab_size, d_model, d_inner, n_layers):
+        super().__init__()
+        self.embed = nn.Embedding(vocab_size, d_model)
+
+        self.layers = nn.ModuleList([])
+        for _ in range(n_layers):
+            self.layers.append(nn.ModuleList([
+                RMSNorm(d_model),
+                minGRU(d_model, d_inner)
+            ]))
+
+        self.norm_f = RMSNorm(d_model)
+        self.lm_head = nn.Linear(d_model, vocab_size, bias = False)
+
+
+
+    # One single step of minGRU, forwarding one token and outputting one token
+    def step(self, x, h_states=None):
+        x        = self.embed(x)
+
+        h_next   = []
+        h_states  = iter(h_states if h_states is not None else [])
+
+        for norm, mingru in self.layers:
+            h_prev   = next(h_states, None)
+            residual = x
+
+            x        = norm(x)
+            x, h_t   = mingru.step(x, h_prev)
+            x        = x + residual
+
+            h_next.append(h_t)
+
+        x      = self.norm_f(x)
+        logits = self.lm_head(x)
+
+        return logits, h_next
+
+
+
+    def forward(self, x, h_states=None):
+        x, labels = x[:, :-1], x[:, 1:]
+        x         = self.embed(x)
+    
+        h_next   = []
+        h_states  = iter(h_states if h_states is not None else [])
+
+        for norm, mingru in self.layers:
+            h_prev   = next(h_states, None)
+            residual = x
+
+            x        = norm(x)
+            x, h_t   = mingru.forward(x, h_prev)
+            x        = x + residual
+
+            h_next.append(h_t)
+
+        x      = self.norm_f(x)
+        logits = self.lm_head(x)
+        loss   = F.cross_entropy(logits.transpose(1, 2), labels)
+
+        return logits, h_next, loss

requirements.txt

@@ -0,0 +1,3 @@
+transformers>=4.44.2
+datasets>=3.0.1
+wandb>=0.18.3

train.py

@@ -0,0 +1,113 @@
+#@title Utility functions for sampling
+
+import torch
+import math
+from transformers import GPT2Tokenizer
+from datasets import load_dataset
+import numpy as np
+
+from model import minGRULM
+from util import generate_text
+
+
+
+
+
+
+
+dataset_path  = 'flpelerin/tinystories-100k'
+
+num_epochs    = 1
+batch_size    = 4
+seq_length    = 256
+learning_rate = 1e-4
+infer_step    = 50
+
+input_len     = 50
+num_predict   = 250
+
+reset_state_every = 16
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"total context size is {batch_size * seq_length} tokens");
+
+
+
+
+tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
+tokenizer.pad_token = tokenizer.eos_token
+vocab_size = tokenizer.vocab_size
+print(f"tokenizer has {vocab_size} unique tokens")
+
+
+
+
+dataset = load_dataset(dataset_path)
+
+def process_function(examples):
+    return tokenizer(examples['text'], padding='longest', truncation=True)
+
+tokenized_datasets = dataset.map(process_function, batched=True)
+print(f"dataset has {tokenized_datasets['train'].num_rows} rows of {batch_size} times {seq_length} tokens")
+
+
+
+
+#model = minGRULM(
+#    vocab_size = vocab_size,
+#    d_model = 768,
+#    d_inner = 1536,
+#    n_layers = 12
+#)
+
+model = minGRULM(
+    vocab_size = vocab_size,
+    d_model = 384,
+    d_inner = 768,
+    n_layers = 6
+)
+
+
+model.to(device)
+print(f"model has {sum(p.numel() for p in model.parameters()):,} parameters")
+
+
+
+
+optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
+
+h_states = None
+
+step = 0
+for epoch in range(num_epochs):
+    for i in range(0, len(tokenized_datasets['train']), batch_size):
+        batch = tokenized_datasets['train'][i:i + batch_size]
+        input_ids = torch.tensor(batch['input_ids']).to(device)
+
+        #if (i / batch_size) % reset_state_every == 0:
+        #    print(f"resetting state, {(i / batch_size)} % {reset_state_every} == 0")
+        #    h_states = None
+            
+        h_states = h_states if (i / batch_size) % reset_state_every != 0 else None
+        str_states = ''.join(['{:.3f}, '.format(h_states[0][0][0][i].item()) for i in range(10)]) if h_states is not None else 'None'
+
+        optimizer.zero_grad()
+        _, h_states, loss = model.forward(input_ids, h_states)
+        loss.backward()
+        optimizer.step()
+
+        step += 1
+        print(f"Epoch: {epoch} / {num_epochs}, Step: {step}, Loss: {loss.item():.4f}, Hidden State: {str_states}")
+
+        if step % infer_step == 0:
+            model.eval()
+
+            ids = input_ids[0][:input_len]
+            text = tokenizer.decode(ids)
+            print(f"input: {text}")
+
+            prompt = ids[None, ...]
+            text = generate_text(model, tokenizer, prompt, num_predict)
+            print(f"output: {text}")
+
+            model.train()

util.py

@@ -0,0 +1,45 @@
+import torch
+import math
+
+
+
+def log(t, eps = 1e-20):
+    return torch.log(t.clamp(min = eps))
+
+def gumbel_noise(t):
+    noise = torch.zeros_like(t).uniform_(0, 1)
+    return -log(-log(noise))
+
+def gumbel_sample(t, temperature = 1., dim = -1, keepdim = True):
+    return ((t / max(temperature, 1e-10)) + gumbel_noise(t)).argmax(dim = dim, keepdim = keepdim)
+
+def top_k(logits, thres = 0.9):
+    k = math.ceil((1 - thres) * logits.shape[-1])
+    val, ind = torch.topk(logits, k)
+    probs = torch.full_like(logits, float('-inf'))
+    probs.scatter_(-1, ind, val)
+    return probs
+
+
+
+def generate_text(model, tokenizer, prompt: torch.Tensor, seq_len: int):
+    prompt_seq_len = prompt.shape[-1]
+
+    h_states = None
+    logits   = None
+    text     = ""
+
+    for i in range(prompt_seq_len):
+        tok = prompt[:, i:i+1] # (1, 1)
+        logits, h_states = model.step(tok, h_states)
+
+    for _ in range(seq_len):
+        logits = top_k(logits, thres=.9)
+        token  = gumbel_sample(logits, temperature=.7, dim=-1)[0]
+
+        logits, h_states = model.step(token, h_states)
+
+        token = tokenizer.decode(token.item())
+        text += token
+
+    return text