ai-law-society-lab/gpt-oss-20b-controller-ckpt

Temporally Extended MoE Controller Checkpoints

Controller and LoRA adapter checkpoints for Temporally Extended Mixture-of-Experts Models (paper · project page · code).

These checkpoints add a learned option controller on top of a frozen gpt-oss-20b MoE base model. The controller decides when to switch the active expert set and which experts to activate, using the options framework with deliberation costs. This reduces expert switch rates from >50% to as low as ~1–4%, while retaining most of the base model's accuracy.

Checkpoints

All checkpoints use k = 16 (16 out of 32 experts selected per layer). Three deliberation cost (η) settings are provided: eta0.02/, eta0.03/, and eta0.04/.

Each directory contains:

• activation_controller.pt — per-layer option controller weights (termination head, selection head, value heads)
• lora/ — LoRA adapter (rank 16, α = 16) applied to attention projections (q/k/v/o) and expert MLPs (gate_up_proj, down_proj)

Usage

See the code repository for full instructions. In brief:

1. Load gpt-oss-20b with the custom transformers patches (see install.sh).
2. Load the LoRA adapter from eta{X}/lora/.
3. Load the controller weights from eta{X}/activation_controller.pt.
4. Run inference with the controller enabled.

Citation

@article{shen2025temporally,
  title={Temporally Extended Mixture-of-Experts Models},
  author={Shen, Zeyu and Henderson, Peter},
  year={2025}
}