Skip to content

RESEARCH_BRIEF_2026_01_17.md

Latest commit

 

History

History
114 lines (81 loc) · 3.79 KB

RESEARCH_BRIEF_2026_01_17.md

File metadata and controls

114 lines (81 loc) · 3.79 KB
Raw

GroundThink Research Brief — January 17, 2026

For External Research Sessions

Project Summary

GroundThink is a hybrid language model architecture that runs RWKV-6 and Mamba-2 in parallel, with a learned arbiter that dynamically selects which pathway to use per-token.

Core Hypothesis: Different sequence modeling architectures excel at different patterns. A learned routing mechanism can achieve better performance than either alone.

Current State (v0.5.5.1) ✅ BREAKTHROUGH

Architecture

Input → Embedding → [TwinDebateBlock × N] → LM Head → Output

TwinDebateBlock:
├── Pre-Norm (RMSNorm)
├── RWKV-6 pathway (CUDA kernel, 0.02 output scaling)
├── Mamba-2 pathway (CUDA kernel, 0.05 output scaling, 0.5x residual)  
├── LinearRecurrenceArbiter (temperature=3.0 for exploration)
├── Weighted fusion: α_rwkv * h_rwkv + α_mamba * h_mamba
├── Post-Norm + FFN (GELU)
└── Residual connection

Validated Results

Metric Result
Hybrid vs RWKV-only 15.7% better (PPL 22.1 vs 36.0)
Hybrid vs Mamba-only 3% better (PPL 22.1 vs 22.8)
Arbiter Sensitivity 0.21 (21x above target)
Domain-Adaptive Routing ✅ Different routing for stories vs wikipedia
CUDA Kernels ✅ Both RWKV-6 and Mamba-SSM CUDA working
Routing Balance ✅ 64-89% RWKV (SOLVED!)

✅ SOLVED: Routing Collapse (Jan 17, 2026)

The Problem (Before)

With both CUDA kernels active, routing collapsed to 98-99% RWKV across all layers.

Root Causes Found

  1. Output magnitude mismatch: RWKV CUDA outputs 597x stronger than Mamba CUDA
  2. Softmax sharpening: Temperature=1.0 causes winner-take-all collapse

The Fix

# ops/twin_debate_block.py - Output balancing
CUDA_RWKV_SCALE = 0.02   # RWKV raw std ~2.5 → 0.02
CUDA_MAMBA_SCALE = 0.05  # Mamba raw std ~0.004 → 0.02  (2.5x boost needed)

# ops/linear_recurrence_arbiter.py - Exploration
temperature = 3.0  # Prevents softmax collapse

Results After Fix (100 steps)

Layer Before After
L0 99% RWKV 89% RWKV (entropy 0.35)
L2 99% RWKV 64% RWKV (36% Mamba!)
L5 99% RWKV 74% RWKV (entropy 0.57)

Entropy: 0.02-0.08 → 0.35-0.65 (healthy exploration!)

⚠️ NEW FINDING: Mamba-Only Wins at 81M Scale! (Jan 17, 2026)

Comparative Endurance Test (1000 steps)

Mode Params Final Loss Final PPL Speed
Mamba-only 64.51M 4.16 63.98 3.29/s
Hybrid 81.31M 4.29 72.98 1.67/s
RWKV-only 74.06M 4.59 98.53 2.26/s

Surprising Result: Mamba-only is 2x faster AND gets lower loss!

Why This Differs from Phase 2 Results

Metric Phase 2 (20M) Now (81M)
Kernels Prototype CUDA
Hybrid PPL 22.1 (best) 72.98 (2nd)
Mamba PPL 22.8 63.98 (best)
RWKV PPL 36.0 98.53

Hypothesis: The CUDA Mamba kernel is much stronger than the prototype. The hybrid overhead may not be worth it when Mamba is this capable.

Questions for Further Research

  1. Is this TinyStories-specific? Try WikiText or code datasets
  2. Does hybrid catch up with more steps? Run 10k step comparison
  3. Is temperature=3.0 hurting optimization? Try temp=1.5
  4. Does hybrid win on harder tasks? Long-range dependency tests

Balanced Configuration (Validated)

Component Setting Purpose
CUDA_RWKV_SCALE 0.02 Scale RWKV output to std ~0.02
CUDA_MAMBA_SCALE 0.05 Scale Mamba output to std ~0.02
temperature 3.0 Prevent softmax collapse
MAMBA_RESIDUAL_SCALE 0.5 Mode C: force Mamba signal