⚡ BabyKilat

A Lightweight MoE Language Model with Hybrid Attention

140M parameters | Trained on ~1B tokens | 6,000 steps | Free to use

Model Description

• Hybrid Attention — Combines efficient global decay heads with precise latent MLA heads
• Mixture-of-Experts (MoE) — 4 experts (1 shared + 3 routed), 2 active per token
• No Positional Encoding (NoPE) — Relies on attention structure rather than explicit position embeddings
• Lightweight Design — Runs on consumer GPUs with just 6GB VRAM (RTX 4050 compatible)
• Built from Scratch — Trained using the custom Kilat framework

Architecture Details

Hybrid Attention Mechanism

                    Input x [B, N, D]
                           │
           ┌───────────────┴────────────────┐
           │                                │
           ▼                                ▼
   ╔═══════════════════╗           ╔════════════════════╗
   ║   PATH 1          ║           ║  PATH 2           ║
   ║   Global Decay    ║           ║   Latent MLA       ║
   ║   (O(N) linear)   ║           ║   (O(N²) exact)    ║
   ╚═══════════════════╝           ╚════════════════════╝
           │                                │
           ▼                                ▼
    Efficient recurrent                  High-precision
    state update                         full attention
           │                                │
           └───────────────┬────────────────┘
                           │
                           ▼
                    Learned gate fusion
                    (per-head, per-token)

• Global Decay Path — Linear-time attention with exponential decay, O(1) KV-cache per step
• Latent MLA Path — Standard softmax attention with 4x compressed KV-cache
• Recall Ratio — 0.5 means half of heads use each path, balanced for quality/efficiency

Quick Start

Option 1: Using Kilat (Recommended)

Option 2: Using TextGenerator Wrapper

Training Details

Dataset

• Source: FineWeb-Edu (educational content)
• Format: Memory-mapped numpy arrays
• Sequence Length: 2048 tokens
• Tokens Seen: ~123M tokens (6,000 steps × 32 batch × 8 grad accum × 2048 seq len / 1M)

Training Configuration

Training Results (at step 6,000)

Hardware & Estimated Training Time

Generation Strategies

Requirements

• Python 3.10+
• PyTorch 2.0+

Repository Files

BabyKilat/
├── config.json              # Model configuration
├── config.yaml              # Human-readable config
├── model.safetensors        # Model weights (140M params)
├── tokenizer_config.json    # Tokenizer config
├── tokenizer.json           # GPT-2 tokenizer
├── vocab.json               # Vocabulary
└── merges.txt               # BPE merges

Limitations

• English only. Trained exclusively on English text.
• No instruction tuning. Base language model only, not aligned for chat or task completion.
• Partially trained. This checkpoint represents 6,000 of 50,000 planned steps (~12% of full training). Quality will improve significantly with full training.
• Limited capacity. At 140M parameters, the model underperforms larger models on knowledge-intensive benchmarks. Outputs may be factually incorrect.
• Context window. 2,048-token context window.
• No formal evaluation. Benchmark results have not been reported for this release.

Citation

Author

License