MiniCPM-3B-Effective-Muon

🎯 Model Description

MiniCPM-3B-Effective-Muon is a custom Universal Transformer language model inspired by MiniCPM and Qwen architectures. It achieves approximately 3.3B effective parameter depth through recursive application of a single ~166M parameter block, trained with the novel Muon optimizer for matrix parameters.

Key Features

• 🧠 Universal Transformer: Single block reused 20 times for deep reasoning
• ⚡ Muon Optimizer: Newton-Schulz orthogonalization for matrix weights (arXiv:2502.16982)
• 🎯 GQA Attention: Grouped Query Attention with 4:1 ratio (16 Q heads, 4 KV heads)
• 🌀 RoPE: Rotary Position Embeddings for positional encoding
• 🔥 SwiGLU FFN: 2.67x expansion ratio (Qwen-style)
• 🌍 Bilingual: Trained on English and Russian text
• 💾 Memory Efficient: ~166M physical parameters, fits in 10GB VRAM

🏗️ Architecture

Model Specifications

Universal Transformer Architecture

Unlike traditional Transformers with N stacked layers, this model uses a single QwenBlock reused 20 times with depth embeddings:

for i in range(20):  # Universal passes
    x = x + depth_embedding[i]
    x = checkpoint(block, x, rope_cos, rope_sin)

This creates ~3.3B effective parameter compute while keeping the physical memory footprint at only ~166M parameters.

Components

1. Grouped Query Attention (GQA)

• Query projections: 16 heads
• Key/Value projections: 4 heads (shared across 4 query heads each)
• Efficiency: 4x reduction in KV cache memory vs Multi-Head Attention

2. SwiGLU Feed-Forward Network

def swiglu(x):
    return down(silu(gate(x)) * up(x))

• Expansion ratio: 2.67× (1024 → 2736 → 1024)
• No bias terms (Qwen-style)

3. RMSNorm

• Root Mean Square Layer Normalization
• Pre-norm architecture (before attention and FFN)
• No bias parameters

🎓 Training: Muon + AdamW Hybrid Optimizer

This model uses a hybrid optimization strategy inspired by Muon: MomentUm Orthogonalization by Newton-schulz:

Matrix Parameters (2D+): Muon Optimizer

• Momentum SGD with Nesterov acceleration
• Newton-Schulz iterations for polar decomposition (orthogonalization)
• 5 iterations of quintic polynomial approximation
• Applied to: attention projections, FFN layers

Vector Parameters (1D): AdamW

• Standard AdamW with weight decay 0.1
• Betas: (0.9, 0.95)
• Applied to: embeddings, RMSNorm weights, biases

Optimizer Hyperparameters

📚 Training Data

Dataset: openbmb/Ultra-FineWeb

• Source: High-quality web text corpus
• Language Filter: English and Russian
• Chunk Size: 1024 tokens
• Total Training Tokens: ~50M (2000 steps × 2 batch × 8 grad_accum × 1024 tokens)

Data Statistics

• English texts: 60,463
• Russian texts: 0 (filtered by language field)
• Total chunks: 50,000
• Train/Val split: 95% / 5%

Training Configuration

📈 Training Results

Final Metrics

Training Progress

Step   | Loss    | Val Loss | LR (AdamW)
-------|---------|----------|-----------
25     | 133.12  | -        | 2.39e-04
400    | ~80     | 7.57     | ~2e-04
1000   | ~65     | ~7.0     | ~1e-04
1701   | 55.13   | 7.57     | 4.79e-05
2000   | 54.66   | **6.61** | 8.69e-05

💻 How to Use This Model

⚠️ Important: This is a custom architecture. You must define the model classes before loading weights.

Step 1: Install Dependencies

pip install torch transformers

Step 2: Define Architecture

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint as grad_checkpoint

# Copy these classes from the training script:
# - RMSNorm
# - SwiGLU
# - GroupedQueryAttention
# - QwenBlock
# - UniversalQwen

CONFIG = {
    'vocab_size': 151665,
    'd_model': 1024,
    'n_heads': 16,
    'n_kv_heads': 4,
    'd_ff': 2736,
    'num_passes': 20,
    'max_seq_len': 1024,
    'rope_theta': 10000.0,
}

model = UniversalQwen(**CONFIG)

Step 3: Load Weights and Tokenizer

from huggingface_hub import hf_hub_download
from transformers import AutoTokenizer

repo_id = "antontuzovAI/minicpm-3b-effective-muon"

# Download and load weights
weights_path = hf_hub_download(repo_id=repo_id, filename="best_model_weights.pth")
checkpoint = torch.load(weights_path, map_location="cpu", weights_only=False)
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B", trust_remote_code=True)

Step 4: Generate Text

def generate(prompt, max_new_tokens=150, temperature=0.7):
    tokens = tokenizer.encode(prompt, add_special_tokens=False)
    x = torch.tensor([tokens], dtype=torch.long)
    
    with torch.no_grad():
        for _ in range(max_new_tokens):
            logits = model(x)
            logits = logits[:, -1, :] / temperature
            probs = F.softmax(logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
            x = torch.cat([x, next_token], dim=1)
            
            if next_token.item() == tokenizer.eos_token_id:
                break
    
    return tokenizer.decode(x[0].tolist(), skip_special_tokens=True)

# Test generation
print(generate("The future of artificial intelligence"))

⚠️ Limitations

1. Training Scale: This model was trained for only 2000 steps (~50M tokens). For production quality, 10,000+ steps on 1B+ tokens are recommended.
2. Language Coverage: Despite bilingual intent, the training data was predominantly English due to dataset filtering.
3. Research Prototype: This is an experimental model demonstrating Universal Transformer + Muon optimizer combination. Not intended for production deployment.
4. Hallucinations: Like all LLMs, may generate factually incorrect information.
5. Context Length: Limited to 1024 tokens due to training constraints.

🎯 Intended Use

This model is intended for:

• Research: Studying Universal Transformer architectures
• Education: Learning about Muon optimizer implementation
• Experimentation: Testing hybrid optimization strategies
• Prototyping: Base for further fine-tuning on specific domains

Not intended for:

• ❌ Production deployment without additional training
• ❌ Critical decision-making systems
• ❌ Generation of harmful, biased, or misleading content

📚 Citation

If you use this model or architecture in your research, please cite:

@misc{minicpm_3b_effective_muon,
  title={MiniCPM-3B-Effective: Universal Transformer with Muon Optimizer},
  author={Your Name},
  year={2026},
  publisher={Hugging Face},
  howpublished={\url{https://huggingface.co/antontuzovAI/minicpm-3b-effective-muon}}
}

Related Work

• Muon: MomentUm Orthogonalization by Newton-schulz
• MiniCPM: Revealing the Devil in Transformer's Details
• Qwen2.5 Technical Report
• Universal Transformer (Dehghani et al., 2018)

🙏 Acknowledgments

• OpenBMB Team: For the Ultra-FineWeb dataset
• Qwen Team (Alibaba): For the Qwen2.5 architecture and tokenizer
• Muon Optimizer Authors: For the novel optimization algorithm
• Hugging Face: For hosting infrastructure and datasets library

This model was trained from scratch using custom PyTorch code. No pre-trained weights from other models were used.