MoE-Study-Applied

Python PyTorch License: MIT Transformers Github Repo Github Repo Github Repo


Table of contents


Overview

A supervised fine-tune of the 201M-parameter sparse Mixture-of-Experts model from MoE-Study-Remastered into an abstractive summarizer. Six hours on a single RTX 4060 and 181,773 BBC article/summary pairs from XSum. This is the third study in the MoE-Study trilogy, continuing from MoE-Study-Remastered — where the base model was trained.

At a glance

Base checkpoint MoE-Study-Remastered
Total parameters 201,269,248
Context length 1024
Fine-tuning data EdinburghNLP/xsum, 181,773 train pairs
Task Abstractive single-sentence summarization
Hardware NVIDIA GeForce RTX 4060, 8GB VRAM
Fine-tuning time 21,844s (6.07 h)
Code MoE-Study-Applied

Results

base sft
Test cross-entropy (held-out XSum) 3.3070 2.3281
Behaviour on an article never terminates stops after 18 tokens

At first glance the fine-tune specialized the model at no cost. The following sections are about what it actually cost.

The behavioural change

Both checkpoints, one BBC article each, up to 500 tokens, top_k=50, temp=0.8.

The base model, given a piece about England's T20 World Cup squad:

', 4.
"It's going to be exciting and exciting," said Stanislav Khurr, who plays for The Daily
Express. "It's going to be very exciting and it's going to be very much looking forward
to the final game."cious has said that he has been disappointed that he didn't want to be
part of the club and was disappointed when they had to return the favour.
"I guess on the road I'm going to be a little bit disappointed. I always thought it was a
good place to play but I've been there for a bit. I've been in my home for four and a
half years. I'm not happy, I'm sad. It's not the right thing for me to do.
[...runs to the full 500 tokens]

It never stops, never summarizes, and invents a journalist. It is continuing text, which is all a base model knows how to do.

The fine-tuned model, given a Charlton Athletic transfer roundup that mentions a season-long loan at Ajax, defender Rod Fanni signing from Al-Arabi, and cancelled contracts for Franck Moussa and Ricardo Vaz Te:

Charlton Athletic have recalled striker Philippe Fanni from fellow side Blackpool on loan.

Eighteen tokens, then <|end|>, unprompted. Right register, right length, right entity domain. Also wrong: the surname is lifted off Rod Fanni, "Philippe" is invented outright, and Blackpool appears nowhere in the article.

The shape was learned completely. The grounding was not learned at all.

Benchmarks

Seven tasks, lm-evaluation-harness v0.4.13, full splits, batch size 8. Few-shot counts are inherited from Remastered. (If you want to know the reasoning for the few-shot counts for each metric, please read MoE-Study-Remastered)

Benchmark Shots Metric base sft Δ
wikitext 0 word_perplexity 75.20 106.39 +31.19
lambada_openai 0 acc 19.02 15.54 −3.48
piqa 0 acc_norm 62.35 60.28 −2.07
hellaswag 5 acc_norm 29.22 28.57 −0.65
arc_easy 0 acc_norm 35.44 35.27 −0.17
arc_challenge 15 acc_norm 22.78 23.12 +0.34
winogrande 5 acc 51.22 52.88 +1.66

Sorted by damage. The top two moved, the bottom five did not.

WikiText: 75.20 vs 106.39 word_perplexity Lambada: 19.02 vs 15.54 acc PIQA: 62.35 vs 60.28 acc_norm HellaSwag: 29.22 vs 28.57 acc_norm ARC-Easy: 35.44 vs 35.27 acc_norm ARC-Challenge: 22.78 vs 23.12 acc_norm WinoGrande: 51.22 vs 52.88 acc

Wikitext and lambada are pure language modelling: continue this text, predict this final word. Both are exactly the capability the fine-tune stopped training. This is textbook catastrophic forgetting, and three decisions made it unavoidable:

  1. Full-parameter SFT. Every weight moved, including the ones holding general language modelling.

  2. A single narrow distribution. XSum is BBC news and one-sentence summaries, with no replay of pretraining data mixed in. Nothing pulled the model back toward where it started.

  3. Loss masked to the summary span. The model got gradient only on short, formulaic sentences and was explicitly not trained to model the long input articles. Open-ended continuation had no reason to survive.

Secondary metrics
Benchmark Metric base sft
lambada_openai perplexity 201.27 378.81
wikitext byte_perplexity 2.243 2.394
wikitext bits_per_byte 1.166 1.259
arc_easy acc 37.88 37.12
piqa acc 62.68 61.59
hellaswag acc 27.95 27.74
arc_challenge acc 17.83 18.17

Full splits throughout: arc_easy 2,376, piqa 1,838, wikitext 62, lambada_openai 5,153, winogrande 1,267, hellaswag 10,042, arc_challenge 1,172.

Speed

KV cache vs no cache

Time-to-first-token is flat within each model, as it should be: the first token needs a full forward pass over the prompt either way, so there is nothing yet to reuse. The gap between the two rows is prompt length rather than caching, since each run drew its own random document.

Where it actually stands

It reliably produces well-formed summaries. One sentence, newswire register, correct length, terminated by the model itself. A capability the base model did not have in any form.

It hallucinates freely. No faithfulness metric (ROUGE, factual consistency, entity overlap) was run, so the rate is unmeasured.

It is worse at everything else. If you need a general-purpose small model, the base checkpoint is strictly better and it lives in the other repo.

It is not a general summarizer. Every claim here is about XSum, which is BBC news in one house style. The single-sentence output length is a property of the dataset the model absorbed, not a switch to turn on/off.

Most benchmark numbers sit near chance for both checkpoints. Hellaswag 28.57 against 25.00 chance, arc_challenge 23.12 against 25.00, winogrande 52.88 against 50.00. Only piqa and arc_easy clear the floor by a real margin.

What was not tried. No ablations, so the individual contributions of the cache, chat template, masking and learning rate are unknown. No parameter-efficient variant to compare forgetting against. No summarization-specific metric, the most obvious gap. No RLHF.


Using it

Unlike Remastered, you do not have to write your own sampling loop, because LLM.generate ships inside the checkpoint. It is not the Hugging Face generate: it takes a raw string, applies the chat template itself, and returns a (text, tokens_per_second, time_to_first_token) tuple, where text is the whole decoded sequence with the formatted prompt still in it.


import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset

def get_random_prompt(DATA_DIR, SEED, SPLIT, tokenizer, context_length, max_new_tokens):
    ds = iter(load_dataset(DATA_DIR, split=SPLIT).shuffle(SEED))

    while True:
        prompt = next(ds)["document"]
        formatted_prompt = "<|user|>" + prompt + "<|end|>" + "<|assistant|>"
        ids = tokenizer(formatted_prompt).input_ids
        if len(ids) <= context_length - max_new_tokens:
            break

    return prompt

MODEL = "OliverSundaram/MoE-Study-Applied"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

tokenizer = AutoTokenizer.from_pretrained(MODEL)
model = AutoModelForCausalLM.from_pretrained(MODEL, trust_remote_code=True)

# Keeps the embedding table and the special tokens in sync; required for the base checkpoint.
model.resize_token_embeddings(len(tokenizer))
model.config.eos_token_id = tokenizer.eos_token_id
model.config.pad_token_id = tokenizer.pad_token_id
model.to(device).eval()

prompt = get_random_prompt(
    DATA_DIR="EdinburghNLP/xsum",
    SEED=42,
    SPLIT="test",
    tokenizer=tokenizer,
    context_length=model.config.context_length,
    max_new_tokens=100
)

text, tok_per_sec, ttft = model.generate(
    prompt,
    tokenizer,
    device,
    max_new_tokens=100,
    top_k=50,
    temp=0.8,
    use_cached=True,
    print_text=True
)

trust_remote_code=True is required, because model_type is custom_llm and the architecture loads from the modules.py shipped inside the checkpoint.


Citation

@misc{moe-study-applied,
  author       = {Oliver Sundaram},
  title        = {MoE-Study-Applied: Supervised Fine-Tuning a Sparse Mixture-of-Experts Language Model for Abstractive Summarization},
  year         = {2026},
  publisher    = {GitHub},
  howpublished = {\url{https://github.com/OliverSundaram/MoE-Study-Applied}}
}

References


License

MIT — see LICENSE.

Downloads last month
25
Safetensors
Model size
0.2B params
Tensor type
F32
·
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Dataset used to train OliverSundaram/MoE-Study-Applied

Papers for OliverSundaram/MoE-Study-Applied