README.md

---
license: apache-2.0
library_name: pytorch
language:
  - en
pipeline_tag: text-generation
tags:
  - mlp-mixer
  - byte-level
  - causal-lm
  - attention-free
  - ccd-mixer
  - content-gated-dilated-conv
  - tiny-stories
  - pretrained-backbone
  - micro-language-model
  - sub-1m-parameters
model_name: MicroMixer-4-50K-TinyStories
datasets:
  - roneneldan/TinyStories
metrics:
  - perplexity
---

<div align="center">

<img src="https://raw.githubusercontent.com/llaa33219/MicroMixer-4/main/logo.svg" width="300" alt="MicroMixer-4 Logo"/>

# MicroMixer-4-50K-TinyStories

<img src="https://img.shields.io/badge/Parameters-48%2C684-blue?style=for-the-badge&logo=python&logoColor=white&color=%23007BFF" alt="Parameters"/>
<img src="https://img.shields.io/badge/Architecture-CCD--Mixer-purple?style=for-the-badge&color=%23AE00FF" alt="Architecture"/>
<img src="https://img.shields.io/badge/Stage-Pretrain-green?style=for-the-badge&color=%2300D620" alt="Pretrain"/>

<br/>
<br/>

<table>
<tr>
<td align="center" style="padding: 20px;">
<strong style="color: #007BFF; font-size: 1.2em;">Micro Language Model</strong><br/><em>Attention-Free • MLP-Only • Byte-Level • Content-Gated Dilated Convolution</em>
</td>
</tr>
</table>

[![GitHub](https://img.shields.io/badge/GitHub-MicroMixer--4-blue?style=for-the-badge&logo=github&color=%23007BFF)](https://github.com/llaa33219/MicroMixer-4)

</div>

<div style="background: linear-gradient(135deg, #007BFF22, #AE00FF22); padding: 20px; border-radius: 10px; border-left: 4px solid #007BFF;">

## 📋 Overview

**MicroMixer-4-50K-TinyStories** is a **48,684-parameter** pure MLP-Mixer causal language model — **no attention, no recurrence, no SSM** — **pretrained on [TinyStories](https://huggingface.co/datasets/roneneldan/TinyStories)** as part of the MicroMixer-4 **dataset-efficiency comparison study** ([analysis](https://github.com/llaa33219/MicroMixer-4/blob/main/DATASET_COMPARISON_ANALYSIS.md)). TinyStories is a corpus of ~200K short synthetic children's stories (GPT-3.5/4-generated, constrained grammar and vocabulary).

⚠️ This repo is **pretrain-only** — it is **NOT FMSP-fine-tuned**. TinyStories contains free-flowing prose with no `User:`/`Assistant:` dialogue markers, so the FMSP answer-only cross-entropy recipe does not apply (there is no answer region to isolate). What you get is the **raw pretrained backbone** (seed 42, V76 recipe, 3 epochs). It continues stories; it does not answer questions or follow instructions.

The backbone is **V87 Final**, the project's champion CCD-Mixer architecture (the V83-RPG champion frozen and scaled to six budgets). The 50K preset reproduces the champion recipe at its budget.

</div>

## 🏗️ Architecture

<div align="center">

```mermaid
graph TD
    A[Byte Input] --> B[Embed 256→32 NoPE]
    B --> C[CCD-Mixer Block × 4]
    C --> D[RMSNorm]
    D --> E[LM Head Tied with Embed]
    E --> F[Byte Output]

    subgraph "CCD-Mixer Block"
        X[Input 32] --> U["Linear d→2d → split v, g"]
        U --> RP[Full RoPE on v AND g]
        RP --> M["Shared-weight dilated conv<br/>dilations 1·2·4·8, k=65"]
        M --> G["Per-position 4-way gate<br/>softmax(Linear_dil(x)/τ)"]
        G --> O["W_o(v ⊙ g)  — zero-init"]
        O --> SW[SwiGLU Channel-Mix]
        SW --> RM[ReMixerLayer sidecar]
    end

    style A fill:#007BFF,color:#fff
    style F fill:#00D620,color:#fff
    style G fill:#AE00FF,color:#fff
    style M fill:#FF6600,color:#fff
```

</div>

### Model Configuration

<table>
<tr>
<th style="background-color: #007BFF; color: white;">Parameter</th>
<th style="background-color: #AE00FF; color: white;">Value</th>
</tr>
<tr><td>Hidden Dimension (d_model)</td><td><code>32</code></td></tr>
<tr><td>Number of Blocks</td><td><code>4</code></td></tr>
<tr><td>Token-Mix</td><td><code>GLCTokenMixCCD</code> (content-gated mixture of shared-weight dilated causal conv)</td></tr>
<tr><td>Dilations</td><td><code>(1, 2, 4, 8)</code> — one shared depthwise kernel, zero extra conv params</td></tr>
<tr><td>Depthwise Kernel Size</td><td><code>65</code></td></tr>
<tr><td>RoPE</td><td>Full RoPE on <b>both</b> v and g (V76 "RPG" pattern)</td></tr>
<tr><td>Channel-Mix</td><td><code>SwiGLU</code></td></tr>
<tr><td>Sidecar</td><td><code>ReMixerLayer</code> per block (label_dim 16, pool_heads 4)</td></tr>
<tr><td>Max Sequence Length</td><td><code>1024</code></td></tr>
<tr><td>Vocabulary Size</td><td><code>256</code> (byte-level)</td></tr>
<tr><td>Position Encoding</td><td>RoPE inside token-mix only; no position embedding table</td></tr>
<tr><td>Normalization</td><td><code>RMSNorm</code> (pre-norm)</td></tr>
<tr><td>Output Head</td><td>Tied with input embedding</td></tr>
<tr><td>Zero-Init</td><td><code>W_o</code>, <code>dil_gate</code>, <code>log_τ</code> — silent at init</td></tr>
</table>

### Core Components

```
┌──────────────────────────────────────────────────────────────┐
│                 CCD-Mixer Block (×4)                        │
│  u    = Linear(d → 2d)(x)                                    │
│  v, g = u.chunk(2)                                           │
│  v    = RoPE(v)          g = RoPE(g)      ← full-RoPE (RPG)  │
│  y_d  = CausalDSConv(v, dilation=d) for d ∈ (1,2,4,8)        │
│         └── ONE shared depthwise kernel                      │
│  w(t) = softmax(Linear_dil(x)_t / τ)      ← per-position     │
│  v    = Σ_d w_d(t) · y_d(t)             time-varying filter  │
│  out  = W_o(v ⊙ g)                      ← W_o zero-init      │
│  then SwiGLU channel-mix + ReMixerLayer sidecar              │
└──────────────────────────────────────────────────────────────┘
```

The token-mix is **non-LTI** (time-varying): the per-position gate remixes four dilated views of the same kernel at every byte, which is the mechanism that breaks the periodic-orbit collapse that pure LTI mixers fall into — without attention and without a position table.

---

## 🎯 Generation Examples

<div style="background-color: #FF050515; padding: 15px; border-radius: 8px; border-left: 4px solid #FF6600;">

**Story-continuation probes** (raw text — no `User:/Assistant:` prefix, since TinyStories is a plain LM; greedy decoding, `repetition_penalty=1.2`, `no_repeat_ngram_size=4`, `max_new_tokens=200`):

```
[Prompt] Once upon a time, there was a little girl named Lily.
[Output] She loved to play with her friends. One day, Lily was very happy and saw a big box of the store…
```

<sub>coherent — on-topic story continuation in the TinyStories register</sub>

```
[Prompt] One day, a boy named Tom found a big red balloon.
[Output] They were very happy and he was so happy. One day, he was so happy to help him and said, "I wan…
```

<sub>coherent — on-topic story continuation in the TinyStories register</sub>

```
[Prompt] The sun was shining and the birds were singing.
[Output] The bird was so happy and said, "I want to see the store with the store with the store with the…
```

<sub>drifts — starts as story but falls into a repetitive loop</sub>

```
[Prompt] A little dog ran into the garden and saw
[Output] a big box of fun. He was so happy and he wanted to see the box and said, "I want to see the bea…
```

<sub>coherent — on-topic story continuation in the TinyStories register</sub>


</div>

---

## 📊 Results

<div style="background-color: #007BFF15; padding: 15px; border-radius: 8px; border-left: 4px solid #007BFF;">

### Pretraining (TinyStories, V76 recipe, 3 epochs)

| Metric | 1 ep | 2 ep | 3 ep |
|--------|------|------|------|
| Val PPL | 2.83 | 2.75 | **2.60** |

AdamW lr 3e-3 · WSD (warmup 500) · wd 0.01 · bs 16 · seq 1024 · seed 42 · plain CE on non-pad bytes.

### MicroMixer-4 TinyStories family (pretrain-only, all sizes)

| Size | Params | 3ep Val PPL |
|------|--------|------------|
| 1M | 996,873 | **1.78** |
| 500K | 491,742 | **1.90** |
| 300K | 292,525 | **2.00** |
| 100K | 95,084 | **2.31** |
| **50K** | 48,684 | **2.60** |
| 10K | 9,666 | **4.06** |

<sub>Pretrain-only family — FMSP-based axes (chatter fluency, full-988 EM, q-relevance, OOD, unanswerable fabrication) are N/A: TinyStories has no `User:/Assistant:` markers, so the answer-only-CE recipe does not apply.</sub>

</div>

---

## 📚 Training Data

<div style="background-color: #00D62015; padding: 15px; border-radius: 8px; border-left: 4px solid #00D620;">

1. **Pretraining**: [TinyStories](https://huggingface.co/datasets/roneneldan/TinyStories) — synthetic short stories generated by GPT-3.5/4 with a constrained vocabulary and simple grammar, ~200K stories sampled, flattened to 1024-byte sequences, 3 epochs. No `User:/Assistant:` dialogue structure.

</div>

---

## 🔧 Usage

### Files in this repository
- `epoch_{0,1,2}.safetensors` — per-epoch pretrained backbone weights (pickle-free safetensors). **`epoch_2.safetensors` is the final (3rd-epoch) checkpoint.** No FMSP adapter — this is the plain backbone.

### Load and generate (local clone)

```python
import torch
from safetensors.torch import load_file
from src.model_v87_final import MicroMixerV87Final, v87_final_50k
from src.tokenizer import ByteTokenizer

# Clone the code repository first:
# git clone https://github.com/llaa33219/MicroMixer-4.git && cd MicroMixer-4

cfg = v87_final_50k()
model = MicroMixerV87Final(cfg)                         # plain backbone — NO attach_adapter (pretrain-only)
model.load_state_dict(load_file("epoch_2.safetensors"), strict=True)
model.eval()

tok = ByteTokenizer()
prompt = "Once upon a time, there was a little girl named Lily."
ids = tok.encode(prompt)
if ids and ids[-1] == tok.eos_token_id:
    ids = ids[:-1]                       # ByteTokenizer appends EOS; the prompt must end open
ids = torch.tensor([ids])
with torch.no_grad():
    out = model.generate(
        ids, max_new_tokens=200,
        temperature=0.0,                 # greedy
        repetition_penalty=1.2,
        no_repeat_ngram_size=4,
        eos_token_id=tok.eos_token_id,
    )
print(prompt + tok.decode(out[0].tolist()[len(ids):]))
```

> Note the differences from the FMSP cards: (1) **no `attach_adapter`** — the backbone is loaded
> as-is; (2) the prompt is **raw story text**, not the `User: …\n\nAssistant: ` dialogue format.

### Load from Hugging Face Hub (no clone of the weights needed)

```python
import torch
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from src.model_v87_final import MicroMixerV87Final, v87_final_50k

REPO = "llaa33219/MicroMixer-4-50K-TinyStories"

cfg = v87_final_50k()
model = MicroMixerV87Final(cfg)
model.load_state_dict(
    load_file(hf_hub_download(REPO, "epoch_2.safetensors")), strict=True)
model.eval()
# ... continue a story as above
```

---

## ⚠️ Limitations

<div style="background-color: #FF050515; padding: 15px; border-radius: 8px; border-left: 4px solid #FF0505;">

| Limitation | Description |
|------------|-------------|
| **Pretrain-only — no instruction/QA ability** | Not FMSP-fine-tuned; it only continues TinyStories-style prose. It cannot answer questions or follow instructions. |
| **Micro parameters** | 48,684 parameters; capacity is the binding constraint |
| **Knows only TinyStories** | Distribution is synthetic children's stories; no real-world knowledge |
| **Byte-level noise** | 256-vocab byte tokenizer; PPL not comparable to BPE baselines |
| **Research use only** | Architecture/pretraining research artifact, not a production model |

</div>

---

## 🧬 Context

This is the **50K** TinyStories-pretrained arm of the **dataset-efficiency comparison study** in the [MicroMixer-4](https://github.com/llaa33219/MicroMixer-4) project — the pretrain-only third corpus alongside the UltraChat and SmolTalk2 FMSP arms (TinyStories is excluded from the FMSP/eval battery because it has no `User:/Assistant:` markers). Sibling repos: `llaa33219/{MicroMixer-4,MicroT-test1}-{1M..10K}-TinyStories`, plus the UltraChat/SmolTalk2 arms `…-{UltraChat,SmolTalk2}` and the discord-pretrained baselines `llaa33219/{MicroMixer-4,MicroT-test1}-{1M..10K}`. Full analysis: [DATASET_COMPARISON_ANALYSIS.md](https://github.com/llaa33219/MicroMixer-4/blob/main/DATASET_COMPARISON_ANALYSIS.md).

---

<div align="center">

[![GitHub](https://img.shields.io/badge/Back_to_Repository-MicroMixer--4-blue?style=for-the-badge&logo=github&color=%23007BFF)](https://github.com/llaa33219/MicroMixer-4)

<sub>Part of the <a href="https://github.com/llaa33219/MicroMixer-4">MicroMixer-4</a> research project — V87 Final (CCD-Mixer) family, 50K preset, TinyStories pretraining (pretrain-only)</sub>

</div>