README.md

---
library_name: transformers
pipeline_tag: text-generation
tags:
- custom_code
- causal-lm
- cma
- small-language-model
- generalist
- 4k-tokenizer
datasets:
- HuggingFaceFW/fineweb_edu_100BT-shuffled
- HuggingFaceTB/smollm-corpus
- epfml/FineWeb-HQ
- HuggingFaceTB/finemath
---

<div align="center">

# CMA 8M

### Channel-Mixing Attention at 7.85 million parameters.

A compact generalist language model built to study channel mixing at small scale.

**7.85M parameters | 1,024-token context | 4,096-token vocabulary | Safetensors**

**Benchmark-selected at step 40,000 | Open SLM Leaderboard-style average: 36.21%**

</div>

---

## Overview

CMA 8M is the benchmark-selected checkpoint of a **7,849,161-parameter**
Channel-Mixing Attention generalist causal language model. It pairs a compact
nine-layer architecture with the unmodified native GPT-S 4,096-token tokenizer.

| Detail | Value |
|---|---|
| **Architecture** | Channel-Mixing Attention causal LM |
| **Parameters** | 7,849,161, with tied input/output embeddings |
| **Checkpoint** | Step 40,000 |
| **Training tokens** | 20,971,520,000 |
| **Context length** | 1,024 tokens |
| **Vocabulary** | 4,096 native tokens |
| **Weights** | Safetensors |

## How CMA works

In each of the nine transformer blocks, causal token attention is followed by
CMA in the feed-forward position. For every token, CMA reshapes the 288-channel
hidden state into **12 chunks of 24**, routes between those chunks with **3
heads**, applies a signed blend and SiLU gate, then projects back to 288 channels.

![CMA 8M architecture](https://huggingface.co/User01110/cma-8M/resolve/main/assets/cma-architecture.svg?v=4)

## Quick start

This repository contains a custom Transformers architecture. Load it with
`trust_remote_code=True` to register CMA and the model's rotary convention.

```python
from transformers import AutoModelForCausalLM, AutoTokenizer

repo_id = "User01110/cma-8M"

tokenizer = AutoTokenizer.from_pretrained(
    repo_id,
    trust_remote_code=True,
)
model = AutoModelForCausalLM.from_pretrained(
    repo_id,
    trust_remote_code=True,
    dtype="auto",
)
model.eval()
```

Generate a continuation:

```python
import torch

prompt = "Channel-mixing attention is"
inputs = tokenizer(prompt, return_tensors="pt")

with torch.inference_mode():
    output = model.generate(
        **inputs,
        max_new_tokens=80,
        do_sample=True,
        temperature=0.8,
        top_p=0.95,
    )

print(tokenizer.decode(output[0], skip_special_tokens=True))
```

## Architecture

| Component | Configuration |
|---|---|
| Total parameters | **7,849,161** |
| Tied embedding parameters | 1,179,648 |
| Non-embedding parameters | 6,669,513 |
| Width / layers | 288 / 9 |
| Context length | 1,024 tokens |
| Token-attention heads | 6 query / 2 KV |
| CMA | chunk 24 / 3 heads / expansion 2 |
| Positional encoding | Contiguous-half RoPE |
| Tokenizer | `AxiomicLabs/GPT-S-5M` at revision `df47402` |
| Vocabulary | 4,096 native tokens |
| Runtime | PyTorch 2.5+ for native SDPA grouped-query attention |
| Weights | Safetensors |

Input and output embeddings are tied.

## Training mixture

All four sources are streamed natural-text corpora.

| Source | Share of trained tokens |
|---|---:|
| FineWeb-Edu 100BT shuffled | **55.00%** |
| Cosmopedia v2 | **25.00%** |
| FineWeb-HQ | **10.00%** |
| FineMath 4+ | **10.00%** |

FineMath 4+ supplies high-quality mathematical explanations and reasoning as
ordinary causal-language-model text.

## Training trajectory

Across the logged evaluation checkpoints, normalized bits per byte (BPB)
decreased from **3.7838** to **1.3520**, while the Open SLM Leaderboard-style
average increased from **30.66%** to **36.21%**.

![CMA 8M training trajectory](https://huggingface.co/User01110/cma-8M/resolve/main/assets/training-trajectory.svg?v=4)

## Evaluation

### Zero-shot results at step 40,000

The four lm-eval tasks use normalized accuracy when supplied by `lm-eval`
0.4.12, with float32 weights and softmax. ArithMark uses float32 weights and its
official raw continuation log-likelihood-sum rule.

| Benchmark | Accuracy |
|---|---:|
| HellaSwag | **28.19%** |
| ARC-Easy | **35.35%** |
| ARC-Challenge | **23.29%** |
| PIQA | **58.22%** |
| ArithMark-2 | **29.12%** |
| ARC mean | **29.32%** |
| **Open SLM Leaderboard-style average** | **36.21%** |

The selection average is:

```text
(HellaSwag + mean(ARC-Easy, ARC-Challenge) + PIQA + ArithMark-2) / 4
```

This matches the Open SLM Leaderboard formula. The checkpoint was selected at
step 40,000 for its **36.21%** average.

### WikiText-103 validation

| Metric | Value |
|---|---:|
| Loss | **2.9903** |
| Perplexity | **19.89** |
| Normalized BPB | **1.3520** |
| Scored tokens | 358,911 |
| Normalized UTF-8 bytes | 1,145,226 |

WikiText normalized BPB is reported separately and was **not** used for
checkpoint selection.