kdqemre/FGA-GPT-v3

Overview

FGA-GPT v3 (Functional Geometric Attention) is an EXPERIMENTAL 67M-parameter Small Language Model (SLM) trained entirely on an Apple M1 Pro (16GB). Designed to test the efficacy of Pure Clifford (Geometric) Attention against standard dot-product attention. By replacing the standard attention mechanism with a hybrid dot-product (symmetric) and wedge-product (skew-symmetric) projection, the model natively learns directional and structural relationships between tokens.

To ensure a fair evaluation against State-of-the-Art (SOTA) standards, the architecture is modernized with components inspired from modern architecture families.

Architecture Highlights

1. Pure Clifford Attention: Utilizes a learnable skew-symmetric bivector matrix (W_skew) to project the exterior product (wedge product) of queries and keys into a scalar attention score.
2. Learnable Alpha ($\alpha$) Blending: Each attention head dynamically learns an $\alpha$ parameter to blend the dot-product (coherence/semantics) and wedge-product (syntax/directionality) scores before the softmax operation.
3. Rotary Position Embeddings (RoPE): Replaces absolute positional embeddings, injecting relative spatial geometry directly into the Q and K vectors (optimized with buffer caching).
4. SwiGLU FFN: Replaces the standard GELU multi-layer perceptron for higher capacity gating.
5. RMSNorm: Replaces LayerNorm for computationally efficient, mean-free normalization.
6. GPT-2 Tokenizer: I have used GPT-2 tokenzer.

💻 Hardware & Training Data

• Hardware Limit: Trained entirely locally on an Apple M1 Pro with 16GB Unified Memory using the MPS (Metal Performance Shaders) backend. To prevent Out-Of-Memory (OOM) crashes with the 67M parameter footprint and SwiGLU expansions, the batch size was strictly capped at 8 for pre-training and 4 for SFT.
• Pre-Training Datasets: A highly curated, locally cached mixture of:
• TinyStories
• WikiText-2
• Ultra-FineWeb

📐 Mathematical Motivation: Standard vs. Clifford Attention

In the standard Transformer architecture, the attention score relies exclusively on the inner (dot) product:

$$Scor{e}_{standard}=\frac{Q{K}^T}{\sqrt{d_k}}Score\_\{standard\}\; =\; \backslash frac\{Q\; K^T\}\{\backslash sqrt\{d\_k\}\}$$

The Limitation: The dot product is inherently symmetric ($a \cdot b = b \cdot a$). It measures "similarity" or "coherence" but lacks a native sense of direction or sequence. Transformers compensate for this by artificially injecting positional embeddings (like RoPE or Absolute encodings).

The Clifford (Geometric) Solution: In Geometric Algebra, the multiplication of two vectors (the Geometric Product) yields both a symmetric scalar and an asymmetric bivector (oriented area):

$$qk=q\cdot k+q\wedge kqk\; =\; q\; \backslash cdot\; k\; +\; q\; \backslash wedge\; k$$

The exterior (wedge) product ($q \wedge k$) is anti-symmetric ($q \wedge k = -k \wedge q$). This anti-symmetry naturally encodes direction, syntax, and structural flow. By integrating this into the attention mechanism, the model doesn't just learn if two tokens are related, but how they are oriented relative to each other in time and logic.

⚙️ How I Implemented It (Pure Clifford Core)

Calculating the full bivector for every token pair has an $O(T^2 d^2)$ complexity, which is computationally impossible for standard VRAMs. To implement this gracefully:

1. The Skew-Symmetric Projection: I project the bivector into a scalar score using a learnable skew-symmetric matrix $W_{skew}$. This mathematically mirrors the exact projection of a bivector.

$$W_{skew}=W-W^{T}W\_\{skew\}\; =\; W\; -\; W^T$$

$$Scor{e}_{wedge}=\frac{Q{W}_{skew}{K}^T}{\sqrt{d_k}}Score\_\{wedge\}\; =\; \backslash frac\{Q\; W\_\{skew\}\; K^T\}\{\backslash sqrt\{d\_k\}\}$$

2. Learnable Alpha ($\alpha$) Blending: Instead of forcing a fixed ratio, I introduce a learnable $\alpha$ parameter for each attention head. The model dynamically blends the symmetric (dot) and asymmetric (wedge) scores before the softmax operation.

$$Attention=Softmax(\alpha (Scor{e}_{dot})+(1-\alpha )(Scor{e}_{wedge}))VAttention\; =\; Softmax(\backslash alpha(Score\_\{dot\})\; +\; (1-\backslash alpha)(Score\_\{wedge\}))\; V$$

Experimental Findings & "The Bitter Lesson"

I subjected this model to Supervised Fine-Tuning (SFT) on the GSM8K dataset. The results yielded fascinating insights into the physical limits of Small Language Models:

• Exceptional Format Alignment: The geometric asymmetry provided by the Clifford mechanism allows the model to map and obey strict structural boundaries flawlessly. During greedy decoding tests, the model successfully synthesized the complex GSM8K #### <Answer> <|endoftext|> stopping template with high reliability at this scale.
• Semantic Hallucinations (Capacity Wall): While the model perfectly mimicked the syntax of mathematical reasoning, it lacked the semantics (world model). At 67M parameters, it failed to perform actual arithmetic, generating structurally perfect but logically absurd reasoning (e.g., calculating the "number of days" in an "egg").
• Conclusion: Pure Clifford Attention is highly efficient for syntax mapping and prompt alignment at minimal scales. However, true reasoning remains an emergent property of scale (parameter count and pre-training data volume), reinforcing Rich Sutton's "Bitter Lesson."

---

💻 THE CODE FOR PRE TRAINING

REMINDER: This is not a SOTA code. You can improve the code as you see fit. You need to configure the dataset paths, especially for the UltraFineWeb dataset. Since it is a very large dataset, I have downloaded and used 100K from it. Also, you'll probably need to use more data than I do.

import torch
import torch.nn as nn
import torch.nn.functional as F
import os
import tiktoken
from datetime import datetime
from datasets import load_dataset
from tqdm.auto import tqdm

# --- 1. CONFIGURATION & DEVICE ---
device = torch.device("cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu"))
print(f"Using Device: {device}")

BATCH_SIZE = 8    
SEQ_LEN = 256      
LEARNING_RATE = 5e-4 
EPOCHS = 1 

# --- 2. DATASET LOADING & PREPARATION ---
print("Loading Datasets...")
tinystories_dataset = load_dataset("roneneldan/TinyStories")
train_tinystories = tinystories_dataset['train']['text'][:5000]
val_tinystories = tinystories_dataset['validation']['text'][:500]

wikitext_dataset = load_dataset("zhengxuanzenwu/wikitext-2-split-128")
wiki_text_key = 'text' if 'text' in wikitext_dataset['train'].column_names else wikitext_dataset['train'].column_names[0]
train_wiki = wikitext_dataset['train'][wiki_text_key][:50000]
val_wiki = wikitext_dataset['validation'][wiki_text_key][:1000] if 'validation' in wikitext_dataset.keys() else wikitext_dataset['train'][wiki_text_key][40000:41000]

# Generic path for local datasets
fineweb_local_path = "./data/ultra_fineweb_100k.txt"
if os.path.exists(fineweb_local_path):
    print(f"-> Reading local Ultra-FineWeb data from {fineweb_local_path}...")
    with open(fineweb_local_path, "r", encoding="utf-8") as file:
        all_fineweb_samples = file.read().split("<|UF_SEP|>")
        if not all_fineweb_samples[-1]: 
            all_fineweb_samples.pop()
else:
    raise FileNotFoundError(f"ERROR: {fineweb_local_path} not found. Please ensure the dataset exists.")

train_fineweb = all_fineweb_samples[:50000]
val_fineweb = all_fineweb_samples[40000:41000]

# Sequential concatenation (Shuffle disabled for block processing)
train_texts = train_tinystories + train_wiki + train_fineweb
val_texts = val_tinystories + val_wiki + val_fineweb

combined_train_text = "\n".join(train_texts)
combined_val_text = "\n".join(val_texts)

tokenizer = tiktoken.get_encoding("gpt2")
vocab_size = tokenizer.n_vocab
encode = lambda s: tokenizer.encode(s, allowed_special={"<|endoftext|>"})
decode = lambda l: tokenizer.decode(l)

print("Tokenizing datasets...")
train_data = torch.tensor(encode(combined_train_text), dtype=torch.long)
val_data = torch.tensor(encode(combined_val_text), dtype=torch.long)

# DYNAMIC STEP CALCULATION
tokens_per_batch = BATCH_SIZE * SEQ_LEN
total_train_tokens = len(train_data)
max_iters = (total_train_tokens // tokens_per_batch) * EPOCHS
eval_interval = max(1, max_iters // 10) 

print(f"Total Training Tokens: {total_train_tokens}")
print(f"Calculated Total Iterations for {EPOCHS} Epoch(s): {max_iters}")

def get_batch(split):
    data = train_data if split == 'train' else val_data
    ix = torch.randint(len(data) - SEQ_LEN, (BATCH_SIZE,))
    x = torch.stack([data[i:i+SEQ_LEN] for i in ix])
    y = torch.stack([data[i+1:i+SEQ_LEN+1] for i in ix])
    return x.to(device), y.to(device)


# --- 3. ARCHITECTURE: PURE CLIFFORD + SOTA COMPONENTS ---

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight

class SwiGLU(nn.Module):
    def __init__(self, d_model):
        super().__init__()
        hidden_dim = int(2 * 4 * d_model / 3)
        hidden_dim = 256 * ((hidden_dim + 255) // 256) 
        
        self.w1 = nn.Linear(d_model, hidden_dim, bias=False)
        self.w2 = nn.Linear(d_model, hidden_dim, bias=False)
        self.w3 = nn.Linear(hidden_dim, d_model, bias=False)
        self.dropout = nn.Dropout(0.1)

    def forward(self, x):
        return self.dropout(self.w3(F.silu(self.w1(x)) * self.w2(x)))

class PureCliffordAttention(nn.Module):
    def __init__(self, d_model, num_heads, max_seq_len=256):
        super().__init__()
        self.num_heads = num_heads
        self.d_model = d_model
        self.head_dim = d_model // num_heads

        self.c_attn = nn.Linear(d_model, 3 * d_model, bias=False)
        self.c_proj = nn.Linear(d_model, d_model, bias=False)

        # Skew-symmetric bivector projection matrix
        self.bivector_proj = nn.Parameter(
            torch.randn(num_heads, self.head_dim, self.head_dim) * 0.02
        )

        # Learnable blend coefficient for dot vs. wedge products
        self.alpha = nn.Parameter(torch.ones(num_heads) * 0.5)

        mask = torch.tril(torch.ones(max_seq_len, max_seq_len)).view(1, 1, max_seq_len, max_seq_len)
        self.register_buffer('mask', mask, persistent=False)

        # RoPE Pre-computation and Caching
        t = torch.arange(max_seq_len, dtype=torch.float32)
        inv_freq = 1.0 / (10000.0 ** (torch.arange(0, self.head_dim, 2).float() / self.head_dim))
        freqs = torch.outer(t, inv_freq)
        
        cos_cached = torch.cos(freqs).view(1, 1, max_seq_len, self.head_dim // 2)
        sin_cached = torch.sin(freqs).view(1, 1, max_seq_len, self.head_dim // 2)
        
        self.register_buffer('cos_cached', cos_cached, persistent=False)
        self.register_buffer('sin_cached', sin_cached, persistent=False)

    def _apply_rope(self, x):
        B, nh, T, hd = x.shape
        cos = self.cos_cached[:, :, :T, :].to(x.dtype)
        sin = self.sin_cached[:, :, :T, :].to(x.dtype)
        
        x1 = x[..., :hd//2]
        x2 = x[..., hd//2:]
        return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)

    def forward(self, x):
        B, T, C = x.size()
        q, k, v = self.c_attn(x).split(self.d_model, dim=2)

        q = q.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)

        q = self._apply_rope(q)
        k = self._apply_rope(k)

        scale = self.head_dim ** -0.5

        # --- Dot Product (Symmetric Coherence) ---
        att_dot = (q @ k.transpose(-2, -1)) * scale

        # --- Wedge Product (Skew-Symmetric Structure) ---
        W_skew = (self.bivector_proj - self.bivector_proj.transpose(-1, -2)).contiguous()
        att_wedge = (q @ W_skew @ k.transpose(-2, -1)) * scale

        # Clifford Blend (Pre-Softmax)
        alpha = torch.sigmoid(self.alpha).view(1, self.num_heads, 1, 1)
        att = alpha * att_dot + (1 - alpha) * att_wedge

        att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float('-inf'))
        probs = F.softmax(att, dim=-1)

        out = (probs @ v).transpose(1, 2).contiguous().view(B, T, C)
        return self.c_proj(out)


class GeometricTransformerBlock(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        self.ln_1 = RMSNorm(d_model) 
        self.attn = PureCliffordAttention(d_model, num_heads)
        self.ln_2 = RMSNorm(d_model) 
        self.ffn = SwiGLU(d_model)   

    def forward(self, x):
        x = x + self.attn(self.ln_1(x))
        x = x + self.ffn(self.ln_2(x))
        return x


class GeometricGPT(nn.Module):
    def __init__(self, vocab_size, d_model=512, num_heads=8, num_layers=12): 
        super().__init__()
        self.token_emb = nn.Embedding(vocab_size, d_model)
        self.blocks = nn.Sequential(*[GeometricTransformerBlock(d_model, num_heads) for _ in range(num_layers)])
        self.ln_f = RMSNorm(d_model)
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
        self.token_emb.weight = self.lm_head.weight

    def forward(self, idx, targets=None):
        x = self.token_emb(idx)
        x = self.blocks(x)
        x = self.ln_f(x)
        logits = self.lm_head(x)
        
        loss = None
        if targets is not None:
            B, T, C = logits.shape
            logits_reshaped = logits.view(B*T, C)
            targets_reshaped = targets.view(B*T)
            loss = F.cross_entropy(logits_reshaped, targets_reshaped)
        return logits, loss    

    @torch.no_grad()
    def generate(self, idx, max_new_tokens):
        self.eval()
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -SEQ_LEN:]
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :]
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=1)
        self.train()
        return idx


# --- 4. TRAINING LOOP ---
model = GeometricGPT(vocab_size=vocab_size).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=max_iters, eta_min=1e-5)

total_params_m = sum(p.numel() for p in model.parameters()) / 1e6
print(f"\n[FGA-GPT v3] Total Parameters: {total_params_m:.2f} M")
print("Initiating Pre-Training...")

best_val_loss = float('inf')
best_train_loss = float('inf')
best_model_state = None
best_iter = 0

pbar = tqdm(range(max_iters), desc="Training FGA-GPT")
for current_iter in pbar:
    if current_iter % eval_interval == 0 or current_iter == max_iters - 1:
        model.eval()
        with torch.no_grad():
            inputs_train, targets_train = get_batch('train')
            _, train_loss = model(inputs_train, targets_train)
            inputs_val, targets_val = get_batch('val')
            _, val_loss = model(inputs_val, targets_val)
            
            if val_loss.item() < best_val_loss:
                best_val_loss = val_loss.item()
                best_train_loss = train_loss.item()
                best_iter = current_iter
                best_model_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}
                
        current_lr = scheduler.get_last_lr()[0]
        tqdm.write(f"\n[Evaluation] Step {current_iter}: Train Loss {train_loss.item():.4f}, Val Loss {val_loss.item():.4f} | Best Val: {best_val_loss:.4f} | LR: {current_lr:.6f}")
        
        with torch.no_grad():
            sample_prompt = "The most important thing to remember is"
            ctx_tokens = encode(sample_prompt)
            ctx_tensor = torch.tensor([ctx_tokens], dtype=torch.long, device=device)
            gen_idx = model.generate(ctx_tensor, max_new_tokens=40)[0].tolist()
            gen_text = decode(gen_idx)
            tqdm.write(f"🤖 [Generation Sample] {gen_text}\n" + "-"*50)

        model.train()

    inputs, targets = get_batch('train')
    logits, loss = model(inputs, targets)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()
    scheduler.step()
    
    pbar.set_postfix({'Current_Loss': f"{loss.item():.4f}"})

print(f"\nTraining Complete! Best model captured at step {best_iter} (Val Loss: {best_val_loss:.4f}).")


# --- 5. SAVING & REGISTRY ---
if best_model_state is not None:
    model.load_state_dict({k: v.to(device) for k, v in best_model_state.items()})

timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
os.makedirs("./models", exist_ok=True)
model_filename = f"./models/fga_gpt_v3_{total_params_m:.0f}M_{timestamp}.pt"

checkpoint = {
    'model_state_dict': best_model_state, 
    'vocab_size': vocab_size, 
    'tokenizer': 'gpt2',
    'architecture': 'pure_clifford_qwen_v3'
}
torch.save(checkpoint, model_filename)
print(f"-> Model saved to: {model_filename}")

registry_filename = "./models/training_registry.md"
new_row = f"| {datetime.now().strftime('%Y-%m-%d %H:%M:%S')} | {model_filename} | {512} | {12} | {8} | {SEQ_LEN} | {vocab_size} | {total_params_m:.2f} M | {best_train_loss:.4f} | {best_val_loss:.4f} |\n"

if not os.path.exists(registry_filename):
    with open(registry_filename, "w", encoding="utf-8") as f:
        f.write("| Date | File | d_model | layers | heads | seq_len | vocab | params | Train Loss | Val Loss |\n")
        f.write("|---|---|---|---|---|---|---|---|---|---|\n")

with open(registry_filename, "a", encoding="utf-8") as f:
    f.write(new_row)
print(f"-> Model parameters appended to {registry_filename}.")