Akicou/LLaDA2.2-Flash-REAM-50B

This repository contains a REAM-compressed version of inclusionAI/LLaDA2.2-flash, produced with Akicou/ream — a REAM/REAP-style Mixture-of-Experts compression framework.

It is not an official inclusionAI release.

🧠 Methodology: REAM

REAM (Router Expert Activation Merging) — originally proposed by SamsungSAILMontreal/ream (Jha et al., 2026) — is a compression technique for MoE models that combines expert merging with pseudo-pruning:

1. Calibration — Run a small set of hardcoded prompts through the model to collect router logits and expert output activations.
2. REAP Saliency — Compute per-expert importance as S[i] = mean(||h_i(x)|| × p_i(x)) over tokens routed to expert i.
3. Pseudo-Grouping — Select the top-K most salient experts as centroids, then assign nearby experts using gated similarity (50% hidden state + 50% router logit distribution).
4. Merge — Within each group, merge experts with saliency-weighted averaging (--fast-merge).
5. Prune — Non-centroid experts are removed. The router is shrunk to match the new expert count, keeping only centroid rows.

This implementation follows the original REAM paper closely, including:

• Centroid-only router update (REAP-style)
• Zero saliency edge-case handling (unused experts get a small non-zero value)
• Single-layer sequential hooks to avoid OOM on large models

📊 Compression Details

Metric	Original (LLaDA2.2-flash)	REAM-50B	Change
Routed Experts/Layer	256	128	-128 (-50%)
MoE Layers	31 (layer 0 is dense)	31	—
Top-k per token	8	8	—
Calibration	100 samples, hardcoded prompts	—	—
Merge Method	Saliency-weighted avg (--fast-merge)	—	—
Grouping	REAM pseudo-group (group_size=16)	—	—

The per-layer sigmoid-router expert_bias vector (analogous to DeepSeek's e_score_correction_bias) is correctly shrunk alongside the router weights on every layer.

🛠 How It Was Created

python examples/compress_sequential.py \
    --model inclusionAI/LLaDA2.2-flash \
    --output ./LLaDA2.2-Flash-REAM-50B \
    --target-ratio 0.5 \
    --samples 100 \
    --max-seq-len 512 \
    --batch-size 4 \
    --max-tokens 2048 \
    --cpu-merge \
    --fast-merge \
    --seed 42

Hardware: 4× NVIDIA H100 (80GB each), 192 vCPUs, 2TB RAM

⚠️ Important Notes

• This is a calibrated merge (100 hardcoded prompts), not a full evaluation-grade calibration.
• No benchmark evaluation is claimed here — this is an experimental release intended as a starting point for fine-tuning.
• The model uses trust_remote_code=True (same as the base LLaDA2.2-flash).
• Shared experts and dense layer 0 are left untouched.
• The per-layer router expert_bias is correctly shrunk (256 → 128) alongside mlp.gate.weight.
• LLaDA2 is a block-diffusion (dLLM) model. Its custom generate() is a non-autoregressive, block-wise iterative refinement loop and targets transformers==5.2.0 (the base model's version); on newer transformers (≥ 5.3) the refactored GenerationMixin breaks it, so pin that version for inference. The forward / loss path works on the current stack, so the checkpoint is fine-tune ready regardless.
• The tokenizer files are copied verbatim from the base model (tokenizer.save_pretrained() was found to re-serialize the fast tokenizer incorrectly, producing byte-level fallback; the original tokenizer.json / tokenizer_config.json / chat_template.jinja / custom tokenization_llada2.py are used as-is).

🧪 Output Verification (Smoke Test)

Real generation from this checkpoint (transformers==5.2.0, temperature=0.0, block_length=32, gen_length=512, threshold=0.5, eos_early_stop=True). Note: like the base model, inputs must satisfy num_tokens % block_size(32) == 0.

User: What is the Heyting Algebra thingamajig?

Model Output:

The term "Heyting Algebra thingamajig" appears to be a fictional or made-up term, possibly derived from a combination of words such as "Heyting" and "Algebra" along with "thingamajig."

This is a calibrated-merge release: a 50%-compressed checkpoint with no fine-tuning yet. Generation is coherent but visibly lower-quality / more literal than the 256-expert base (which gives a full intuitionistic-logic explanation). Fine-tuning is expected to recover quality. Forward + cross-entropy loss verified (loss ≈ 0.19 on a sample sentence).

🔧 Integrity Checks

• All 31 MoE layers merged 256 → 128 experts (50.0% compression).
• model.config.num_experts == 128, num_experts_per_tok == 8.
• mlp.gate.weight (128, 4096) and mlp.gate.expert_bias (128,) confirmed per layer; shared experts present; lm_head.weight preserved (tie_word_embeddings=False).
• Tokenizer byte-for-byte identical to the base model (apply_chat_template → 19 tokens for the smoke-test prompt, matching the base).

📦 Basic Usage (forward / fine-tuning path)

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "Akicou/LLaDA2.2-Flash-REAM-50B"

tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    dtype=torch.bfloat16,
    device_map="auto",
    trust_remote_code=True,
)

# NOTE: block-diffusion models require token length to be a multiple of block_size (32).
text = "The quick brown fox jumps over the lazy dog. " * 3
inputs = tokenizer(text, return_tensors="pt").to(model.device)

outputs = model(**inputs, labels=inputs.input_ids)
print("loss:", float(outputs.loss))

For diffusion generation, follow the base model's recommended recipe (temperature=0.0, block_length=32, threshold=0.5, editing_threshold=0.0, gen_length=512, eos_early_stop=True) with a compatible transformers version.

📜 Citation

@article{jha2026ream,
  title={REAM: Merging Improves Pruning of Experts in LLMs},
  author={Jha, Saurav and Hashemzadeh, Maryam and Pasand, Ali Saheb and Parviz, Ali and Lee, Min-Joong and Knyazev, Boris},
  journal={arXiv preprint arXiv:2604.04356},
  year={2026}
}

@misc{llada2,
  title={LLaDA2.2-flash},
  author={inclusionAI},
  year={2026},
  url={https://huggingface.co/inclusionAI/LLaDA2.2-flash}
}

Attribution

All architecture, tokenizer, and original weights are from inclusionAI/LLaDA2.2-flash. This repository contains a REAM-compressed derivative checkpoint.