|
|
--- |
|
|
language: |
|
|
- en |
|
|
license: apache-2.0 |
|
|
library_name: transformers |
|
|
pipeline_tag: text-generation |
|
|
tags: |
|
|
- mixture-of-attentions |
|
|
- distance-attention |
|
|
- metric-attention |
|
|
- mqa |
|
|
- hyperffn |
|
|
- router-gating |
|
|
datasets: |
|
|
- nvidia/Nemotron-Math-HumanReasoning |
|
|
- WeMake/Intelligent-Content-Understanding |
|
|
--- |
|
|
|
|
|
# MoAMetricLM‑100M — Mixture of Attentions (MoA) |
|
|
|
|
|
**A geometry‑aware Transformer that mixes several attention mechanisms and routes them with a metric‑based router.** |
|
|
- **Parameters:** ~185 M (≈ 100 M effective due to the mixture) |
|
|
- **Task:** Causal language modeling (decoder‑only) |
|
|
- **Library:** 🤗 Transformers |
|
|
- **KV cache:** Not yet implemented (generation recomputes the full context at every step) |
|
|
|
|
|
--- |
|
|
|
|
|
## Model card |
|
|
|
|
|
| **Model ID** | `reaperdoesntknow/MoA-100M` | |
|
|
|--------------|-------------------------------------| |
|
|
| **Architecture** | `moa_metric` (custom) | |
|
|
| **Tokenizer** | GPT‑2 (`gpt2`) – `pad_token` set to `eos_token` | |
|
|
| **Context length** | 2048 tokens | |
|
|
| **Training data** | 2 × ≈ 256 k tokens from the datasets listed above | |
|
|
| **Training compute** | CPU‑only (Intel), FP32 | |
|
|
| **Training hyper‑parameters** | LR = 5e‑4 (AdamW), batch = 4, seq ≤ 512, 500 k total tokens | |
|
|
| **Final loss** | ≈ 0.30 (train) | |
|
|
| **License** | Apache‑2.0 | |
|
|
| **Safety** | No alignment or safety fine‑tuning – outputs may be biased or inaccurate. | |
|
|
| **Intended use** | Research on geometry‑aware attention, structured sparsity, and mixture‑of‑attention models. | |
|
|
| **Limitations** | • No KV‑cache → slower generation. <br>• Small token budget → not a general‑purpose LM. <br>• No safety/alignment training. | |
|
|
| **Out‑of‑scope** | High‑stakes applications (medical, legal, etc.) without further evaluation. | |
|
|
|
|
|
--- |
|
|
|
|
|
## Overview |
|
|
|
|
|
MoA replaces the classic dot‑product attention with **metric‑based attention** and blends **four** distinct heads per Transformer block: |
|
|
|
|
|
| Head type | Description | |
|
|
|-----------|-------------| |
|
|
| **LocalConvHead** | Depthwise‑separable 1‑D convolution → captures short‑range context. | |
|
|
| **Metric Multi‑Head Attention (MetricMHAttention)** | Soft‑min over **L2 / cosine / diagonal‑Mahalanobis** distances: <br> \(\displaystyle \text{attn}_{h}(i,j) \propto \exp\!\big(-\alpha_h\|q_i-k_j\|^2\big)\) | |
|
|
| **Metric MQA** | Multi‑Query attention (shared K/V) in the same metric space – cheaper than full MHA. | |
|
|
| **ChannelMixHead** | Per‑token MLP that mixes channel dimensions (no positional mixing). | |
|
|
|
|
|
A **token‑wise router** decides, for each token, which head(s) to use and applies **feature‑gates** (FiLM‑style) and **router‑bias gates** for up/down‑scaling. |
|
|
|
|
|
The **FFN** is a **HyperFFN** – three parallel branches (SwiGLU MLP, separable‑conv, low‑rank) combined by a **branch router**. LayerScale and optional DropPath keep training stable. |
|
|
|
|
|
### Regularisation (optional) |
|
|
|
|
|
* **Triangle‑inequality (TI) penalty** on sampled triples to encourage true‑metric behaviour. |
|
|
* **Ball pruning** – each head learns an **origin** \(o_h\) and **radius** \(r_h\); keys outside the ball are masked, giving structured sparsity. |
|
|
|
|
|
--- |
|
|
|
|
|
## Architecture diagram (high‑level) |
|
|
|
|
|
``` |
|
|
Input → Embedding → (PreNorm) → Block₁ → … → Blockₙ → LM‑Head → Output |
|
|
│ |
|
|
├─ LocalConvHead |
|
|
├─ MetricMHAttention |
|
|
├─ MetricMQA |
|
|
└─ ChannelMixHead |
|
|
(router decides per‑token) |
|
|
|
|
|
Each Block also contains: |
|
|
→ HyperFFN (SwiGLU | Conv | Low‑rank) ← branch router |
|
|
→ LayerScale + DropPath |
|
|
``` |
|
|
|
|
|
--- |
|
|
|
|
|
## Configuration (example) |
|
|
|
|
|
```json |
|
|
{ |
|
|
"model_type": "moa_metric", |
|
|
"vocab_size": 50257, |
|
|
"dim": 768, |
|
|
"num_layers": 12, |
|
|
"attn_heads": 8, |
|
|
"mqa_q_heads": 8, |
|
|
"mixer_hidden": 3072, |
|
|
"ffn_hidden": 3072, |
|
|
"metric": "l2", // "l2" | "cosine" | "maha_diag" |
|
|
"alpha_init": 1.0, |
|
|
"learn_alpha": true, |
|
|
"use_balls": true, |
|
|
"radius_init": 3.0, |
|
|
"learn_radius": true, |
|
|
"origin_init_scale": 0.0, |
|
|
"maha_init": 1.0, |
|
|
"ti_reg_weight": 0.0, |
|
|
"ti_reg_samples": 0, |
|
|
"router_hidden": 128, |
|
|
"router_dropout": 0.1, |
|
|
"router_temperature": 1.0, |
|
|
"attn_drop": 0.1, |
|
|
"proj_drop": 0.1, |
|
|
"drop_path": 0.0, |
|
|
"max_position_embeddings": 2048, |
|
|
"pad_token_id": 50256, |
|
|
"bos_token_id": 50256, |
|
|
"eos_token_id": 50256 |
|
|
} |
|
|
``` |
|
|
|
|
|
> **Tip:** If you use the GPT‑2 tokenizer, set `pad_token = eos_token` and make sure `vocab_size` matches the tokenizer (50257). |
|
|
|
|
|
--- |
|
|
|
|
|
## Quick‑start (inference) |
|
|
|
|
|
```python |
|
|
>>> from transformers import AutoTokenizer, AutoModelForCausalLM |
|
|
|
|
|
>>> model_id = "reaperdoesntknow/MoA-100M" |
|
|
>>> tokenizer = AutoTokenizer.from_pretrained(model_id) |
|
|
>>> tokenizer.pad_token = tokenizer.eos_token # needed for the GPT‑2 tokenizer |
|
|
|
|
|
>>> model = AutoModelForCausalLM.from_pretrained(model_id) |
|
|
|
|
|
>>> prompt = "Explain metric‑based attention in simple terms:" |
|
|
>>> inputs = tokenizer(prompt, return_tensors="pt") |
|
|
>>> output_ids = model.generate( |
|
|
... **inputs, |
|
|
... max_new_tokens=128, |
|
|
... do_sample=False, # deterministic; set temperature>0 for sampling |
|
|
... pad_token_id=tokenizer.pad_token_id, |
|
|
... ) |
|
|
>>> print(tokenizer.decode(output_ids[0], skip_special_tokens=True)) |
|
|
``` |
|
|
|
|
|
*Note:* Because KV‑cache is not implemented, generation time grows linearly with the total context length. |
|
|
|
|
|
--- |
|
|
|
|
|
## Training (custom loop sketch) |
|
|
|
|
|
```python |
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM, DataCollatorForLanguageModeling |
|
|
from torch.utils.data import DataLoader |
|
|
import torch, torch.nn.functional as F |
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained("gpt2") |
|
|
tokenizer.pad_token = tokenizer.eos_token |
|
|
|
|
|
def collate_fn(examples): |
|
|
batch = tokenizer( |
|
|
[ex["text"] for ex in examples], |
|
|
padding="max_length", |
|
|
truncation=True, |
|
|
max_length=512, |
|
|
return_tensors="pt", |
|
|
) |
|
|
labels = batch["input_ids"].clone() |
|
|
labels[batch["attention_mask"] == 0] = -100 |
|
|
batch["labels"] = labels |
|
|
return batch |
|
|
|
|
|
# dataset = load_dataset(..., split="train") # must contain a 'text' field |
|
|
# loader = DataLoader(dataset, batch_size=4, shuffle=True, collate_fn=collate_fn) |
|
|
|
|
|
model = AutoModelForCausalLM.from_pretrained("reaperdoesntknow/MoA-100M") |
|
|
optimizer = torch.optim.AdamW( |
|
|
model.parameters(), |
|
|
lr=5e-4, |
|
|
betas=(0.9, 0.95), |
|
|
weight_decay=0.01, |
|
|
) |
|
|
|
|
|
for batch in loader: |
|
|
out = model(**batch) |
|
|
out.loss.backward() |
|
|
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.2) |
|
|
optimizer.step() |
|
|
optimizer.zero_grad() |
|
|
``` |
|
|
|
|
|
--- |
|
|
|
|
|
## Evaluation checklist |
|
|
|
|
|
* **Perplexity** on a held‑out split of the two training datasets. |
|
|
* **Ablation studies** (keep total token budget constant): |
|
|
* L2 vs. cosine vs. diagonal‑Mahalanobis distance. |
|
|
* With / without ball pruning. |
|
|
* With / without HyperFFN branch router. |
|
|
* With / without TI regulariser. |
|
|
* **Speed / memory** comparison against a vanilla GPT‑2‑size model (same `dim`/`layers`). |
|
|
|
|
|
--- |
|
|
|
|
|
## Efficiency notes |
|
|
|
|
|
| Feature | What it does | |
|
|
|---------|--------------| |
|
|
| **Ball pruning** | Masks keys that lie outside a learned radius → reduces the quadratic attention cost. | |
|
|
| **Metric MQA** | Shares K/V across heads → fewer projection matrices, lower FLOPs. | |
|
|
| **HyperFFN branch router** | Token‑wise top‑k routing means only the most useful branch is evaluated per token. | |
|
|
| **CPU tips** | Set `OMP_NUM_THREADS` / `MKL_NUM_THREADS` to the number of physical cores; use `torch.set_num_threads()` if needed. | |
|
|
|
|
|
Future roadmap: metric‑aware KV‑cache, kernelised distance approximations (e.g., Random Fourier Features), quantisation & mixed‑precision inference. |
|
|
|
|
|
--- |
|
|
|
|
|
## Safety, Bias & Risks |
|
|
|
|
|
* The model **has not been fine‑tuned for safety or alignment**. |
|
|
* Outputs may contain **biases, profanity, or factual errors**. |
|
|
* Do **not** deploy in high‑stakes contexts without additional evaluation, moderation, and possibly further fine‑tuning. |
|
|
|
|
|
--- |
|
|
|
|
|
## License |
|
|
|
|
|
Apache‑2.0 – see the `LICENSE` file in the repository. |
|
|
|
|
|
--- |
|
|
|
|
|
## Citation |
|
|
|
|
|
```bibtex |
|
|
@misc{moametriclm185m, |
|
|
title = {reaperdoesntknow/MoA-100M: A Geometry-Aware Mixture-of-Attentions Language Model}, |
|
|
author = {Colca, Roy Shawn and collaborators}, |
|
|
year = {2025}, |
|
|
url = {https://huggingface.co/reaperdoesntknow/MoA-100M} |
|
|
} |
|
|
``` |
|
|
|
|
|
--- |
|
|
|
|
|
## Changelog |
|
|
|
|
|
| Version | Date | Notes | |
|
|
|---------|------|-------| |
|
|
| **v0.2** | 2025‑09‑20 | 500 k‑token CPU run, GPT‑2 tokenizer, LR = 5e‑4, final loss ≈ 0.30. | |
|
|
| **v0.1** | 2025‑09‑20 | Initial public release: metric heads, MQA, ball pruning, HyperFFN, router & gates; HF‑compatible; no KV cache. | |
|
|
|
|
|
--- |
|
|
|
|
|
## Maintainers |
|
|
|
|
|
* **Author:** reaper (Convergent Intelligence LLC) |
|
|
* **Contact:** *Email* (convergentintelligencenyc@gmail.com)* |
|
|
|
|
|
|
|
|
--- |
|
|
|
|
|
## Special Remarks |
|
|
|
|
|
- This models still in an extremely experimental state. As are most of them, but im working on stabilizing this one for general inference. |
|
|
- I design create and train all of my models using my mathematical research and pure disgust for the dot product! |
|
|
- For those of you who actually read this and use my models, you make my day everytime I see another download, so thank you for being awesome! |
|
|
|
|
|
|
