A 597M-parameter transformer with biologically-inspired activation routing
Instead of a fixed activation function, each neuron dynamically selects among ReLU, Tanh, SiLU, and GELU — like biological neurons selecting neurotransmitters based on context.
PolyChromaticLM is a research language model built from scratch in PyTorch whose core innovation is PolyGLU (Polychromatic Gated Linear Unit) — a drop-in SwiGLU replacement that implements state-conditional activation routing. Rather than applying a single fixed activation function across all neurons, PolyGLU lets each FFN neuron dynamically choose among K=4 activation functions via a differentiable Gumbel-Softmax routing mechanism.
This is the base pre-trained checkpoint (no instruction tuning / SFT). It was trained on ~10B tokens with a math-heavy data mix on a single A100 80GB GPU.
Author: Daniel Nobrega (independent research)
Key Results
Routing converges to near-deterministic selections (entropy = 0.03% of maximum) without any explicit sparsity regularization — an emergent property
Clear depth-dependent activation specialization: early layers prefer GELU, deep layers strongly prefer Tanh
Achieves 62–89% of Qwen3-0.6B-Base benchmark performance using 3,600x fewer training tokens
The routing mechanism adds only 0.23% parameter overhead (~1.4M params)
where g_k = GumbelSoftmax(α_k + β_k · f(h̄), τ) and σ_k ∈ {ReLU, Tanh, SiLU, GELU}.
Each neuron has:
Static preference (α): a learned bias toward specific activations
Dynamic gating (β · f(h̄)): a lightweight MLP that reads the mean-pooled hidden state and modulates routing based on context
Temperature (τ): annealed from 1.0→0.1 during training, controlling routing sharpness
The biological analogy: just as neurons select specific neurotransmitters (glutamate, GABA, dopamine, acetylcholine) depending on circuit state, PolyGLU neurons select activation functions depending on input context.
Context: Qwen3-0.6B-Base was trained on ~36T tokens — approximately 3,600× our budget. Achieving 62–89% of its scores at 0.028% of the training compute demonstrates strong token efficiency for the PolyGLU architecture.
Domain Perplexity
Domain
Training Share
Perplexity
Bits/Token
Math
70% → 85%
3.56
1.83
Code
5%
7.08
2.82
STEM
25% → 10%
31.93
5.00
Domain perplexity across math, code, and STEM
Code perplexity (7.08) is significantly lower than STEM (31.93) despite receiving 5× less data — evidence that mathematical structure transfers effectively to code patterns.
Emergent Routing Behavior
The most striking finding from training: the routing mechanism converges to near-deterministic activation selections without any explicit sparsity loss or entropy regularization.
At convergence, mean dynamic routing entropy is 0.0004 (just 0.03% of the theoretical maximum), meaning the gate network makes near-one-hot activation choices for virtually every neuron.
Per-layer dynamic routing entropy at convergence
Layer-wise Activation Specialization
The model discovers a clear depth-dependent activation gradient:
Early layers (0–5): GELU-dominant (~35–40%) — smooth, probabilistic activations for initial feature extraction
Deep layers (15–27): Tanh-dominant (~50–65%) — bounded compression for deep representational processing
Activation function preference by layer
Three layers (9, 16, 17) maintain elevated routing entropy, suggesting they benefit from activation diversity. Layer 17 notably increases its entropy during the second half of training — counter to the global trend toward determinism.
Neurotransmitter map: preferred activation per neuron across all layers
Usage
This model was trained from scratch in pure PyTorch (no HuggingFace model wrappers). To load and use it:
python
1import torch
2from transformers import AutoTokenizer
34# Clone the training repo for model code5# git clone https://github.com/danielxmed/PolyGLU.git6from src.model.config import ModelConfig
7from src.model.model import load_checkpoint
89# Load model10config = ModelConfig(use_flash_attn=False)11model, step, tau = load_checkpoint("path/to/portable_final.pt", config, device="cuda")12model.eval()1314# Tokenize15tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B-Base")16input_ids = tokenizer("The derivative of x squared is", return_tensors="pt")["input_ids"].cuda()1718# Generate (greedy, no KV cache)19with torch.no_grad():20for _ inrange(50):21 logits = model(input_ids)22 next_token = logits[:,-1,:].argmax(dim=-1, keepdim=True)23 input_ids = torch.cat([input_ids, next_token], dim=1)2425print(tokenizer.decode(input_ids[0]))
Note: This is a base model — it produces raw continuations, not instruction-following responses. An SFT version fine-tuned on math problem-solving is forthcoming.
Limitations
Base model only — no instruction tuning, no chat capability, no RLHF. Outputs are raw text continuations.
10B token training budget — significantly less than comparable-size production models (Qwen3-0.6B: ~36T tokens). General knowledge and factual recall are limited.
Math-heavy distribution (70% math) — strong on mathematical language modeling, weaker on general NLU tasks.
No KV cache — inference requires the full training codebase; generation is slow without a dedicated inference implementation.
English only — trained exclusively on English-language data.
Citation
bibtex
1@misc{nobrega2026polychromaticLM,
2 title = {PolychromaticLM: State-Conditional Activation Routing via Neurotransmitter-Inspired Gated Linear Units},
3 author = {Daniel Nobrega},
4 year = {2026},
5 url = {https://huggingface.co/tylerxdurden/PolyChromaticLM-1.0-base-0.6B}
6}
