| --- |
| license: apache-2.0 |
| language: |
| - en |
| pipeline_tag: text-generation |
| tags: |
| - suprabrain |
| - gated-deltanet |
| - linear-attention |
| - sliding-window-attention |
| - custom-architecture |
| library_name: transformers |
| --- |
| |
| # 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. **Hybrid Layer Layout (3:1 Ratio):** |
| * **Gated DeltaNet (GDN):** 3 out of every 4 layers use Gated DeltaNet linear state-space recurrence for linear-time complexity and fast sequence processing. |
| * **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. |
| 2. **Surprise-Gated Updates ($\beta_t$):** |
| * Implements a scale-invariant residual prediction mechanism (`SurpriseBeta`) that dynamically scales learning updates based on local sequence surprise/prediction error. |
| 3. **Digit-Split Tokenizer:** |
| * 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. |
| 4. **Half-Untied Head & Low-Rank Gates:** |
| * 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. |
| 5. **Custom Muon + AdamW Hybrid Optimizer:** |
| * 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 |
|
|
| | Tasks |Version|Filter|n-shot| Metric | | Value | |Stderr| |
| |--------------|------:|------|-----:|---------------|---|------:|---|------| |
| |arc_challenge | 1|none | 0|acc |↑ | 0.2065|± |0.0118| |
| | | |none | 0|acc_norm |↑ | 0.2329|± |0.0124| |
| |arc_easy | 1|none | 0|acc |↑ | 0.4882|± |0.0103| |
| | | |none | 0|acc_norm |↑ | 0.4255|± |0.0101| |
| |boolq | 2|none | 0|acc |↑ | 0.4223|± |0.0086| |
| |hellaswag | 1|none | 0|acc |↑ | 0.2914|± |0.0045| |
| | | |none | 0|acc_norm |↑ | 0.3160|± |0.0046| |
| |lambada_openai| 1|none | 0|acc |↑ | 0.3072|± |0.0064| |
| | | |none | 0|perplexity |↓ |66.3967|± |2.8867| |
| |openbookqa | 1|none | 0|acc |↑ | 0.1920|± |0.0176| |
| | | |none | 0|acc_norm |↑ | 0.3160|± |0.0208| |
| |piqa | 1|none | 0|acc |↑ | 0.6295|± |0.0113| |
| | | |none | 0|acc_norm |↑ | 0.6175|± |0.0113| |
| |sciq | 1|none | 0|acc |↑ | 0.7020|± |0.0145| |
| | | |none | 0|acc_norm |↑ | 0.5990|± |0.0155| |
| |wikitext | 2|none | 0|bits_per_byte |↓ | 1.0489|± | N/A| |
| | | |none | 0|byte_perplexity|↓ | 2.0689|± | N/A| |
| | | |none | 0|word_perplexity|↓ |48.8008|± | N/A| |
| |winogrande | 1|none | 0|acc |↑ | 0.4878|± |0.0140| |
| |
| --- |
| |
| ## Model Configuration |
| |
| | Hyperparameter | Value | |
| | :--- | :--- | |
| | `hidden_size` | 384 | |
| | `num_hidden_layers` | 28 | |
| | `intermediate_size` | 1152 (2-Matrix MLP with `RationalAct`) | |
| | `gdn_num_heads` / `dim` | 3 heads / 128 dim | |
| | `attn_num_heads` / `kv_heads` | 6 query heads / 2 KV heads (GQA) | |
| | `attn_window` | 256 | |
| | `max_position_embeddings` | 1024 | |
| | `mlp_act` | Per-channel learnable Rational Activation | |
|
|
| --- |
|
|
| ## 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: |
| ```bash |
| wget https://huggingface.co/SupraLabs/SupraBrain-50M/resolve/main/modeling_suprabrain.py |
| ``` |
|
|
| Then, load the model: |
|
|
| ```python |
| 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. |
