SupraLabs/SupraBrain-50M

SupraBrain 50M v0.1

SupraBrain 50M v0.1 is an experimental 50-million parameter hybrid language model engineered by SupraLabs. It combines Gated DeltaNet linear recurrence with Sliding-Window Attention and Surprise-Gated update mechanisms, optimized using a custom Muon + AdamW hybrid optimizer schedule.

Model Summary

• —Developer: SupraLabs
• —Architecture: Hybrid Gated DeltaNet (3:1) + Sliding-Window Attention + Surprise Gating
• —Parameters: ~50M (Sub-50M budget constraint)
• —Vocabulary Size: 23,808 (GEMM-friendly: 186*128, Byte-Level BPE with Digit-Splitting)
• —Context Length: 1,024 tokens (Supports sliding-window attention)
• —Primary Training Data: FineWeb-Edu & Cosmopedia-v2 (5B tokens total)
• —License: Apache 2.0

Key Architectural Innovations

1. 1.Hybrid Layer Layout (3:1 Ratio):
2. 2.Gated DeltaNet (GDN): 3 out of every 4 layers use Gated DeltaNet linear state-space recurrence for linear-time complexity and fast sequence processing.
3. 3.Sliding-Window Attention (SWA): Every 4th layer incorporates localized attention (Window size = 256) with QK-Normalization to maintain strong long-range associative recall. Layer 19 features full global attention.
4. 4.Surprise-Gated Updates ($\beta_t$):
5. 5.Implements a scale-invariant residual prediction mechanism (SurpriseBeta) that dynamically scales learning updates based on local sequence surprise/prediction error.
6. 6.Digit-Split Tokenizer:
7. 7.Custom Byte-Level BPE tokenizer trained on FineWeb-Edu. Enforces single-digit splitting (individual_digits=True) to dramatically boost arithmetic and numerical reasoning performance in sub-100M parameter models.
8. 8.Half-Untied Head & Low-Rank Gates:
9. 9.Utilizes an unembedding rank adapter (unembed_rank=32) and low-rank output gating (gdn_gate_rank=32) to conserve parameter count while maintaining model capacity in the core layers.
10. 10.Custom Muon + AdamW Hybrid Optimizer:
11. 11.2D weight matrices in the body are optimized using the Muon optimizer (Newton-Schulz orthogonalization updates), while embeddings, norms, and 1D vectors are updated via AdamW over a WSD (Warmup-Stable-Decay) schedule.

Benchmarks

Model Configuration

Usage

Since SupraBrain uses a custom architecture without standard Hugging Face native integration, you must register the model class locally before loading it with AutoModelForCausalLM.

Quickstart (Inference Script)

First, download the modeling script:

wget https://huggingface.co/SupraLabs/SupraBrain-50M/resolve/main/modeling_suprabrain.py

Then, load the model:

import importlib
import torch
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer

# Import custom model and config classes from the model script (modeling_suprabrain.py)
sb_module = importlib.import_module("modeling_suprabrain")
SupraBrainConfig = sb_module.SupraBrainConfig
SupraBrainForCausalLM = sb_module.SupraBrainForCausalLM

# Register custom architecture with Hugging Face AutoClasses
AutoConfig.register("suprabrain", SupraBrainConfig)
AutoModelForCausalLM.register(SupraBrainConfig, SupraBrainForCausalLM)

model_id = "SupraLabs/SupraBrain-50M"

print("[*] Loading tokenizer and model...")
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    trust_remote_code=True,
    torch_dtype=torch.bfloat16
).to("cuda")

# Prompt setup
prompt = "The mitochondrion produces"
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

# Generation with repetition penalty control
print("[*] Generating text...")
with torch.no_grad():
    outputs = model.generate(
        **inputs,
        max_new_tokens=100,
        temperature=0.7,
        top_p=0.9,
        do_sample=True,
        no_repeat_ngram_size=3,  # Prevents 3-gram repetitions
        pad_token_id=tokenizer.pad_token_id,
        eos_token_id=tokenizer.eos_token_id
    )

generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
print("\n--- Output ---")
print(generated_text)

Training Details

Training Pipeline & Schedule

• —Dataset: 5 Billion Tokens Total
• —Stable Phase (3.6B Tokens): FineWeb-Edu (sample-100BT)
• —Anneal Phase (1.4B Tokens): 65% FineWeb-Edu (Score $\ge 4.2$) + 35% Cosmopedia-v2

• —Schedule: Warmup-Stable-Decay (WSD) with square-root decay during the annealing phase.
• —Batch Size: Micro-batch size 16 with Gradient Accumulation 8 ($\approx 262,144$ tokens/step over sequence length 1024).

Hardware Requirements & Optimization

• —Dependencies: Optimized with flash-linear-attention (fla) for Gated DeltaNet kernels and PyTorch flex_attention for masked sliding-window operations.
• —FP32 Logit Chunking: Uses memory-checkpointed chunked Cross-Entropy loss to avoid VRAM allocation spikes.