FLEX-MoE

Failure-Born Orthogonal Experts for Self-Improving Language Models.

This repository is the public code and reproducibility supplement for the preprint FLEX-MoE: Failure-Born Orthogonal Experts for Self-Improving Language Models (Alessio Rocchi, 2026). The compiled paper is in paper1_arxiv/main.pdf.

When a frozen language model keeps making the same verifier-checkable mistakes, FLEX-MoE grows a small, removable LoRA expert for them instead of re-training shared weights. The lifecycle is:

frozen base → verifier-grounded failures → failure clusters
            → one candidate LoRA expert per cluster
            → orthogonal "birth" + leakage audit
            → risk-certified route-cell deployment gate
            → promote / reject (reversible)

Two design choices make each expert auditable:

1. Activation-orthogonal birth. Each expert is born in the subspace orthogonal to the activations the base already uses correctly, giving a measurable preservation target before deployment.
2. Risk-certified deployment. A born expert is only allowed to fire on route cells whose per-cell utility passes a control-probe-calibrated lower confidence bound, under a realistic bounded-regression criterion rather than a strict zero-regression bar.

• Held-out positive deployment. On three held-out seeds the risk-certified route-cell gate yields +3.4 held-out points at ~24% expert coverage with 0/384 base-relative regressions (docs/paper_lcb_heldout.json).
• Repair is significant. Over five seeds, failure-born experts improve held-out accuracy over frozen Qwen2.5-1.5B with a seed-level paired t-test p ≈ 0.007 (marginal on 3B) — docs/paper1_multiseed_p4_summary.json.
• Orthogonal birth preserves by construction. Measured activation-leakage reduction tracks the geometric scale 1/√(1−ρ) set by the retained-energy threshold (docs/paper_experiments_summary.json).
• Cross-domain. The repair signal transfers to a second, non-arithmetic symbolic-counting domain; the deployment bottleneck transfers too.

The JSON files in docs/ are the result summaries that back the paper's tables.

python3 -m pip install -e .            # core (dependency-light)
python3 -m pip install -e '.[hf,dev]'  # + Hugging Face / PEFT / torch + tests

Requires Python ≥ 3.10. The mock backend runs CPU-only; the LoRA experiments require a CUDA GPU.

Dependency-light end-to-end run on the mock backend (no GPU):

PYTHONPATH=src python3 -m flexmoe.cli run --config configs/mock_mvp.yaml
PYTHONPATH=src python3 -m pytest            # unit tests

GPU experiments (single H200 used for the paper). Set up the stack, then run a config family:

bash scripts/runpod_setup.sh

# baselines + FLEX-MoE on the arithmetic stream
PYTHONPATH=src python -m flexmoe.cli baseline --kind base_only   --config configs/qwen_1_5b_centroid.yaml
PYTHONPATH=src python -m flexmoe.cli baseline --kind single_lora --config configs/qwen_1_5b_centroid.yaml
PYTHONPATH=src bash scripts/run_hf_mvp.sh configs/qwen_1_5b_centroid.yaml

The paper's specific runs are driven by:

Exact bit-for-bit numbers depend on the GPU / library stack. The committed docs/*.json summaries are the canonical record of the runs reported in the paper.

The benchmarks (arithmetic_hard, symbolic_hard) are synthetic streams generated by code with an exact-match verifier — there is no external dataset to download. Generation lives in src/flexmoe/data.py and is fully determined by the per-run seed.

src/flexmoe/    FLEX-MoE package (pipeline, clustering, orthogonal birth,
                route-cell gate, leakage audit, baselines, verifier)
configs/        Experiment configs (mock, baselines, ablations, exp/, multiseed/)
scripts/        Run drivers and aggregators
tests/          Unit tests
docs/           Result-summary JSONs backing the paper tables
paper1_arxiv/   LaTeX sources + compiled PDF of the preprint

Large run artifacts (LoRA adapter checkpoints, raw per-example reports) are not included here to keep the repository light; the docs/*.json summaries capture the reported numbers.

• pipeline.py — failure discovery → clustering → birth → routed evaluation.
• orthogonal.py / lora_projection.py — activation-orthogonal expert birth.
• activation_leakage.py — H_pass ΔWᵀ residual audit of saved experts.
• route_cells.py / active_evidence.py — risk-certified route-cell deployment.
• baselines.py — base-only, single-LoRA, oracle family-LoRA baselines.
• verification.py — exact-match verifier used by the benchmarks.

@misc{rocchi2026flexmoe,
  title  = {FLEX-MoE: Failure-Born Orthogonal Experts for Self-Improving Language Models},
  author = {Rocchi, Alessio},
  year   = {2026},
  note   = {Preprint}
}

(Update with the SSRN/arXiv identifier once assigned.)

Released under the MIT License — see LICENSE.