ayanami-kitasan/swe-pruner-pro-mimo-v2-flash-head

SWE-Pruner Pro — Mimo V2 Flash Head Pruning Head

A lightweight pruning head for MiMo-V2-Flash(309B MoE, XiaomiMiMo/MiMo-V2-Flash), from the paper:

SWE-Pruner Pro: The Coder LLM Already Knows What to Prune arXiv: 2607.18213

Overview

This head attaches to the frozen backbone and predicts per-token keep/prune logits from the backbone's hidden states. It is used to compress tool responses in long-horizon coding agent trajectories.

• Architecture: LayerNorm → Linear(d,d)→GELU→Dropout → Linear(d,d)→GELU→Dropout → Linear(d,1) + 8-bucket size embedding
• Parameters: ~18M
• Hidden size: 2048

Usage

from huggingface_hub import hf_hub_download

# Download checkpoint
for fname in ["best_model.pt", "model_config.json"]:
    hf_hub_download(repo_id="ayanami-kitasan/swe-pruner-pro-mimo-v2-flash-head", filename=fname, local_dir="./checkpoints/mimo-v2-flash-head")

# Load the pruning head
from swe_pruner_pro.model.head import PruningHead

head = PruningHead(checkpoint_dir="./checkpoints/mimo-v2-flash-head", hidden_size=2048, device="cuda:0")

For the full training and inference pipeline, please refer to the GitHub repository.

Links

• Paper: arXiv 2607.18213
• Code: GitHub
• Training corpus: ayanami-kitasan/swe-pruner-pro-training-corpus