3D-DELTA

SteerMoE steers Mixture-of-Experts LLMs by clamping experts up or down at inference time, and it chooses which experts using a hand-engineered risk-difference statistic over router logits. 3D-DELTA replaces that statistic with a small learned low-rank head Δ ∈ ℝ^L×H×E, read from the last-token hidden state and trained contrastively on three datasets. It plugs into the identical inference rule, so it is an apples-to-apples, drop-in upgrade.

The problem

MoE routers encode behaviorally meaningful axes such as faithfulness, safety, and refusal. SteerMoE (Fayyaz et al., ICLR 2026) showed that small additive interventions on router logits move the model along those axes without retraining. Its detection step computes a static risk difference from contrastive SQuAD pairs; its deployment step force-activates the top-A positive cells and force-deactivates the top-D negative cells, globally across all (layer, expert) positions. The open question: can a hand-engineered population statistic be replaced by something learned that pools tokens better, combines several datasets, and exposes more than one usable axis?

The method

We fit eℓ = h_pool · Δℓ with Δℓ = W_proj · Δℓlr: a shared encoder W_proj ∈ ℝH×k and a per-layer factor Δℓlr ∈ ℝk×E, so capacity scales with the rank k rather than depth. The head is read from the last token (under causal attention it has seen the whole prompt) and trained with a contrast-aware pairwise hinge against the routing target. At inference it feeds SteerMoE's exact greedy A/D mask rule, so the comparison is like-for-like: same rule, same knobs, learned source.

Headline result (OLMoE-1B-7B)

We report Δ, the best-steered cell minus each method's own no-steer baseline, because Δ isolates what the steering does. Our heads share one baseline; the SteerMoE and adsrivatsa columns are quoted from their papers (different baselines), shown for reference.

The Combined head has the largest mean gain, about four times SteerMoE's, and matches or beats SteerMoE's gain on all six benchmarks. Its single largest result is FaithEval-Inconsistent, where it gains +.288 and every other method moves the benchmark by under +.08.

Three findings that make it a paper

1. Last-token pooling beats mean pooling, at every rank

The 36-variant (rank k × aggregation) sweep is monotone: last-token wins at every rank, with held-out ranking accuracy 0.946 against 0.926 for mean-over-all at k=64. Mechanistically, the mean-pool head puts 30 of its top 50 |Δ| cells at layer 0, an early lexical attractor, while the last-token Combined head puts only 3 there and spreads the rest across the middle layers, where the routing decision is behaviorally relevant.

2. FEVER drives the conflict win

The +.288 is not a generic "more data helps" effect. In four controlled training-mix experiments, SQuAD + FEVER alone reaches +.328 on FaithEval-Inconsistent, slightly above the full three-way mix, while SQuAD + HotpotQA gives only +.143 and SQuAD alone +.075. FEVER's claim-against-evidence pairs are themselves a conflict-detection task, so the head learns a routing axis aligned with exactly what that benchmark tests.

3. One head exposes several routing axes

Across every head we evaluated, the best (A, D) cell is different for each benchmark: an abstention regime (low A, moderate D) is best for FaithEval-Inconsistent and Unanswerable, while a parametric-recovery regime (high A) is best for MQuake. One head that exposes all of these is the signature of a multi-directional controller; a single static mask cannot do it.

Generalization across MoE families

A separate question is whether the learning recipe is OLMoE-specific. We ran the full pipeline on two larger, architecturally different MoEs. On all three the same last-token, k=64 configuration wins, reaching 0.836 to 0.976 held-out ranking accuracy across a sevenfold span in parameters and a sixteenfold span in expert count.

Why a learned head helps

SteerMoE's risk difference is a one-dimensional linear statistic from a single dataset, so it picks one direction per run. Our head has capacity for several directions, and the A×D knob selects which to surface, which is why the per-benchmark argmax cells genuinely differ. Multi-dataset training then lets the hinge loss prioritize the most discriminative axis (the FEVER axis for conflict) instead of averaging the objectives.

Limitations

• Steering results are complete for OLMoE-1B-7B; Qwen3 and Mixtral report the learning-side metric while their inference sweeps run (the work is compute-bound: each model is tens of GB and full sweeps run for days on a shared cluster).
• Numbers use a single greedy seed; multi-seed error bars are being added.
• Static-global scope only; a per-sample dynamic scope is future work.
• SteerMoE and adsrivatsa numbers are quoted from their papers, so we compare gains rather than absolute accuracy.

Reproducibility

The whole pipeline, from activation capture through the (k, aggregation) sweep, training, the A×D inference grid, and finalization, runs from one command per model:

python scripts/run_model_pipeline.py \
    --model allenai/OLMoE-1B-7B-0125-Instruct \
    --tag OLMOE --num-experts-per-tok 8

The three-model generalization result was produced this way.