Fused-expert abliteration via a runtime activation hook (no per-trial weight edits)

[rocker-zhang]

In #206 you said that if fused-expert abliteration could be made to work without direct weight modification, you'd be happy to hear about it. Here's a mechanism; I'd like your read before writing code.

Hook the output of each layer's Qwen3_5MoeSparseMoeBlock and project the refusal direction r_hat (unit vector) out of it: out -> out - alpha * (out . r_hat) * r_hat. Because down_proj is bias-free and the map is linear, it commutes with the top-k routing weights and the scatter, so projecting the block output is exactly equivalent to orthogonalizing the down_proj of every routed expert and the shared expert (delta_W = -alpha * r_hat * r_hat^T W); sparse routing adds no approximation.

This differs from the closed #206/#207: no weight is touched during the search. The per-layer alpha is a mutable scalar the hook reads, and reset_model zeroes it back to identity exactly like it zeroes lora_B today; the Optuna kernel and direction selection are reused unchanged. It's the limitation you noted in #339 — PEFT modifies the residual stream through Module hooks, which a raw nn.Parameter lacks — worked around by hooking the parent Module instead.

The one thing I want your call on before building: do you want this in model.py now, or held for the plugin direction in #332/#53?

Two things I'll handle rather than ask about. Efficacy: single-direction projection may underperform the closed direct-weight approach on Qwen3.5 (the multi-direction refusal results suggest it might), so before any PR I'll measure it on a small fused MoE against attention-only on the Pareto front, with a small multi-direction projector as fallback. Export: a hook isn't saved in safetensors, so I'd bake the projection into down_proj at export time, like the current LoRA merge — I'll confirm the details with you once efficacy holds.

[p-e-w]

When deciding whether to implement a major change to the abliteration or tensor handling logic, the most important thing to understand is that testing this logic is extremely time-consuming and very difficult to automate, because inference is computationally expensive, and the exact results are hardware- and platform-dependent. Before every release, I spend a considerable amount of time on making sure that things still work.

As a rule of thumb, I want a few things from such a change:

• It must bring a measurable improvement in the results, not just a theoretical advantage like "now we can also process this tensor". The improvement should plausibly extend to other and future models. The benchmark is ARA, not traditional abliteration. It has to improve upon ARA in order to be merged, since ARA will soon be the default.
• There should be little to no model-specific code. The model zoo keeps growing, and today's models will be obsolete tomorrow, so doing that is unmaintainable. 3-5 model-specific lines are okay, but things like separate methods for different models are not.
• Since Heretic is designed to work fully automatically by default, any change to the abliteration logic must work within that constraint. The user must never be required to figure out some magic parameter in order to get decent results.

Your proposal makes sense from a technical PoV, but it still has to fit into the above requirements. In particular, we want something we can generalize to other architectures with very little effort (GPT-OSS is another example of a model that has a fused expert tensor that is not an nn.Module).

A blueprint for a viable implementation looks like this:

• It works with both standard abliteration and LoRA-based ARA (you can ignore the direct-weight version of ARA, because we won't keep that in the eventual release).
• It works with both Qwen3.5 and GPT-OSS through essentially the same code path.
• It brings convincing improvements over LoRA-based ARA for both models.
• The amount of code added is reasonable (300 LoC net growth max).

On a side note, the testing problem I mentioned in the first paragraph is now solvable in principle with the reproducibility system combined with extremely small instruction-tuned models such as those from trl-internal-testing. It might make sense to implement or at least discuss a testing harness before we go on new architectural adventures.

[rocker-zhang]

Following up with measurable results and one decision for you. I also checked how ARA handles this, so the comparison below is not against a strawman.

Mechanism (changed from the hook above). I target the fused down_proj through PEFT target_parameters (the ParamWrapper path) rather than a forward hook; it stays generic across Qwen3.5 and GPT-OSS on a single code path. The adapter is the source of truth. During evaluation I bake its delta into down_proj and restore a pristine snapshot before the next trial (bit-exact), purely to avoid reconstructing the per-expert delta on every forward; base weights are unchanged across the search and export is a normal merge, so this is not a return to the closed direct-weight approach. Evaluation then runs at roughly attention-only speed (~5 min/trial on this model vs 15-20 min unbaked).

Setup. Qwen/Qwen3-30B-A3B, bf16. Two arms, identical config: 30 Optuna trials each (equal budget), n_startup_trials=8, seed 0, same eval prompts, row-normalization off, thinking disabled (caveat below). Un-abliterated refusal rate 100/100 on the harmful set (n=100).

Pareto fronts (refusals @ KL, lower-left is better):

attention-only (current behavior):  89@0.081  91@0.077  92@0.044  95@0.032  96@0.020  97@0.013  99@0.003
attention + fused down_proj:         39@0.080  58@0.040  73@0.035  92@0.032  93@0.018  95@0.016

At matched KL (fused vs attention-only): KL≤0.081 → 39 vs 89 (−50); KL≤0.077 → 58 vs 91 (−33); KL≤0.044 → 58 vs 92 (−34). The 95% binomial band on n=100 is ≈13, so the gaps are real. Attention-only never drops below 89/100 at any KL: mechanically, the refusal behavior lives in the expert MLPs, which the current loop no longer reaches now that experts are a single fused parameter (#339).

I checked whether ARA changes this baseline, and it does not. Both the direct-weight ARA (ara branch) and LoRA-ARA (#332) discover modules through the same get_layer_modules, whose for expert in layer.mlp.experts: try_add("mlp.down_proj", expert.down_proj) matches no nn.Linear when experts are a fused 3D nn.Parameter, and the resulting type-assert is swallowed by the enclosing suppress(Exception). #332 then does cast(Linear, module) / module.lora_A[...] and #211 does LBFGS([module.weight]), so neither ever sees the experts (PEFT also cannot LoRA-wrap a bare nn.Parameter). #332 adds no target_parameters. So ARA today is attention-only on fused-MoE for the same structural reason: it never reaches the experts either. (The Pareto numbers above are from the directional path; ARA optimizes differently, so its absolute floor may differ, but the same blind spot applies.) The fused targeting lives in the shared get_layer_modules, so it is mechanism-agnostic: if it lands on master, ARA and LoRA-ARA would inherit it when ara next rebases. The reversible bake is not part of that targeting layer; it is only the eval-speed trick for the directional path. LoRA-ARA (#332) keeps the base frozen and steers through live adapters, so what it would reuse is the targeting, not the bake.

Caveats. (1) The numbers above are traditional abliteration; the ARA-vs-ARA+fused run is next, but since ARA shares the blind spot I expect the gap to largely carry over. (2) On the KL metric for thinking models: this model generates <think> as its first token, so a raw first-token KL would be ≈0. Heretic already handles that — response_prefix auto-detection sets the prefix to the closed CoT block and moves the measurement past it; I disabled thinking here to get the same effect directly, so the KL above is meaningful. The only case the startswith detection would miss is a template that injects the opener into the prompt rather than generating it, which I have not seen in a real model.

Harness. I wrote the small CPU test suite you suggested: tiny trl-internal-testing fixtures (qwen3_moe; gpt_oss with transposed experts and a down_proj bias; a nested multimodal config; a dense negative control), asserting detection, closed-form correctness in both orientations, and bit-exact restore.

One question. Since the blind spot is in the shared get_layer_modules, would you prefer the fused targeting stays on master as the default fix (ARA picks it up when ara rebases onto master), or that I bring it onto the ara stack directly?

[p-e-w]

I target the fused down_proj through PEFT target_parameters (the ParamWrapper path) rather than a forward hook

This is almost certainly the correct approach because it outsources the complexity to PEFT; however, the PEFT docs specifically mention that this can add a lot of overhead for MoE experts.

During evaluation I bake its delta into down_proj and restore a pristine snapshot before the next trial (bit-exact)

I don't understand this part. Why not let PEFT handle this?

One question. Since the blind spot is in the shared get_layer_modules, would you prefer the fused targeting stays on master as the default fix (ARA picks it up when ara rebases onto master), or that I bring it onto the ara stack directly?

Assuming we actually merge this, it should go on master first, and ARA will be adapted to it.

 

Now let's get to the crux:

Two arms, identical config: 30 Optuna trials each (equal budget), n_startup_trials=8

It would be really, really nice if it were possible to work like this, but it just isn't. We have a 9-dimensional search space (9 effective parameters). 30 trials with 8 startup trials tell you nothing. You have fewer startup trials than parameters, so you basically exit the startup phase with a Gaussian blob that's almost indistinguishable from a uniform distribution. Thus you have no meaningful prior, and the remaining trials will effectively be fully random as well.

With this setup, the quality of the result will be dominated by sheer sampling luck. You get good results if the sampler happened to guess good parameters, and poor results otherwise. But we gain essentially no information about the overall shape of the objective function, and we don't know what its minima actually look like. I can pretty much guarantee that if you repeat the test with your setup a few times, you will see wildly different outcomes.

Even running 200 trials and repeating this process many times doesn't fully answer whether there is actually an improvement. For some decisions, I've done hundreds of runs (with hundreds of trials each) and I'm still not sure. This is exactly why this problem is so extremely hard.

[rocker-zhang]

You're right about the search. Here's what a multi-seed re-run actually shows, plus the bake.

On the search-quality objection. I re-ran it properly: 5 seeds, both arms, n_startup_trials=16 (above the treatment arm's effective dimensionality — it searches 10 params: direction_scope, direction_index, and 4 × {o_proj, down_proj}; baseline searches 6), n_trials=40, --no-think, identical config per arm, n=100 per arm.

The part of the effect that matters here doesn't depend on search quality. Baseline's parameter set is [attn.o_proj]; the experts are a fused down_proj parameter that isn't in it, so no point in baseline's search space ablates the experts at all. That's a reachability fact, not a tuning outcome. The data: across all 5 seeds baseline's running-min refusals plateaus high and never enters treatment's range. Mean running-min at 25/50/75/100% of the trial budget:

baseline :  93.2  91.4  91.0  88.2
treatment:  58.2  51.4  46.2  42.6

So baseline isn't search-starved at its floor; it's at a structural one.

You're right that the absolute treatment floor is luck-sensitive. Across seeds it ranges 8–57 (one seed found an 8-refusal point, but only at KL 0.56, a deep and expensive solution the others didn't hit). So I won't quote a floor.

What is stable across seeds is the matched-KL gap — baseline minus treatment refusals at the same KL budget — mean ± std over the 5 seeds, with the 95% binomial band for the difference of two counts on n=100:

KL<=0.02:   4.8 ± 4.2   (band 6.6)   within noise
KL<=0.05:  17.8 ± 4.5   (band 10.5)  beyond
KL<=0.10:  29.2 ± 6.0   (band 12.2)  beyond
KL<=0.20:  40.4 ± 6.3   (band 12.9)  beyond

Every seed favors treatment in every bucket at KL ≥ 0.05 (5/5; sign test p = 0.062 two-sided, 0.031 one-sided — the aggregator prints the two-sided value), and the per-seed spread at matched KL is small (std 4.5–6.3) even while the absolute floors are all over the place. The magnitude is noisy; the direction and the matched-KL gap are not.

This is still moderate depth — 24 optimizing trials/seed, not your hundreds — so I'm not claiming I've characterized the objective. But the reachability and the saturation are depth-invariant, and the matched-KL gap is the part I'd stake a claim on.

On the bake — "why not let PEFT handle this?" It's the same overhead the PEFT docs flag. Unmerged, ParamWrapper.forward re-runs the parametrization (the per-expert einsum across experts × layers) on every forward; for a generation-heavy eval that's the ~10–20× tax. merge() collapses it to a plain base_layer(x). So the bake avoids exactly the MoE overhead you mentioned, and PEFT already does it via merge()/unmerge(). The only non-stock part is that I snapshot the pristine tensor and restore it instead of calling unmerge(), because merge→unmerge on a fused bf16 parameter accumulates rounding each trial and I wanted the base bit-exact across the search. If you'd rather keep it stock and eat the (tiny) drift, swapping the snapshot/restore for per-trial unmerge() is a few-line change. Either way the base is unchanged across the search and export is a normal merge.

And noted on master-first — I'll target master.

[p-e-w]

Instead of running multiple studies with far fewer than the recommended number of trials (which is 200), I recommend running a single study with and without the change, but for at least 200 trials each. 40 and even 100 trials just isn't enough, and doing it multiple times doesn't change this. The fact that one of the configurations can't touch the down_proj tensor isn't an argument, because whether or not that tensor actually matters is what we are trying to establish.

Also, please make sure you compare against ARA, not standard abliteration. I've never seen a model where ARA only achieves 88% refusals, even if it can't modify the MLP (ARA reaches 3 refusals for GPT-OSS, which has the same fused expert architecture).

On the bake — "why not let PEFT handle this?" It's the same overhead the PEFT docs flag. Unmerged, ParamWrapper.forward re-runs the parametrization (the per-expert einsum across experts × layers) on every forward; for a generation-heavy eval that's the ~10–20× tax.

Hmm, but doesn't it inflate the VRAM requirements (or RAM, if you offload) when you keep copies of tensors around to be able to restore them later? VRAM is the single most critical resource constraint for many users, and often the deciding factor for whether they can process a model at all.

[rocker-zhang]

Ran the ARA comparison you asked for. On Qwen/Qwen3-30B-A3B, fused-expert coverage does not improve on LoRA-ARA at matched KL.

Setup. Two single studies, 200 trials each, n_startup_trials=60, ara_lora_rank=128, row_normalization=full, thinking auto-collapsed to <think></think>, n=400 harmful prompts.

• Arm 1: stock LoRA-ARA (feat: ARA, but it's LoRA #332). On this model it targets attn.o_proj only; the fused down_proj isn't reachable, as discussed.
• Arm 2: same LoRA-ARA, plus the fused down_proj covered. ARA's I/O capture sees the aggregate MoE input→output at the residual, not per-expert routed I/O, so ARA's input @ Wᵀ objective can't be applied to the fused experts without model-specific routing capture, which Arm 2 avoids. Instead Arm 2 ablates the experts with a closed-form directional step (difference-of-means, baked into down_proj) layered on ARA's attention, with the strength searched per trial. fused_weight=0 recovers Arm 1.

Matched-KL fronts (refusals @ KL):

Arm 1  LoRA-ARA:         99@.012  97@.019  95@.021  94@.029  93@.031  92@.033  65@.051  49@.077  48@.113  11@.228  1@.346  0@.538
Arm 2  +fused down_proj:  96@.035  95@.072  67@.085  35@.086  11@.194  9@.824   7@2.95   5@3.3    0@13.8

At matched KL the fused arm is worse at low KL (its floor is 96@.035, where ARA reaches ~92–93 at comparable KL), better in the 0.08–0.20 band (35 vs 49; 11 vs 48), and worse again past 0.25 (ARA reaches 1@.346; the fused arm needs KL≈0.8 to reach single digits). The 95% binomial band for a refusal-count difference at n=400 is ≈20 (it ranges ~16–28 with the refusal rate), and only the 11-vs-48 bucket clears it while the arm loses elsewhere, so this is no improvement beyond noise. Attention-only LoRA-ARA already removes refusals on this model; the experts add no measurable matched-KL value, consistent with ARA reaching ~3 on GPT-OSS without touching them.

Caveat. Arm 2 searches two extra dimensions (fused strength and direction index) on the same 200-trial budget, so it's under-searched relative to Arm 1. fused_weight=0 should recover Arm 1, but TPE didn't rediscover Arm 1's low-KL frontier within the budget, so part of the low-KL gap is plausibly search dilution rather than the experts being useless. One study can't separate the two. Either way, equal-budget fused coverage didn't win.

Net: I'm not proposing fused-expert coverage as an ARA improvement on this. The CPU test is the one part that stands alone: it pins the fused-nn.Parameter detection (the transformers-5.x regression where experts stop being a ModuleList). Can send it as a standalone test-only PR if that's useful.

[p-e-w]

Thanks for doing this in-depth test, that's pretty much what I feared, and matches my experience with GPT-OSS. The main problem is that modifying the MLP is generally highly damaging to model intelligence, so at similar KLD (intelligence preservation level), the optimum is often simply not touching the MLP.

Looking at the numbers, they seem indistinguishable to me, and while you are correct that the second setup has additional parameters, I doubt that makes a difference. I've run up to 500 ARA trials with GPT-OSS, and found convergence to a near-optimum within 150-200 trials.

Net: I'm not proposing fused-expert coverage as an ARA improvement on this. The CPU test is the one part that stands alone: it pins the fused-nn.Parameter detection (the transformers-5.x regression where experts stop being a ModuleList). Can send it as a standalone test-only PR if that's useful.

I don't understand this part. What is the proposed change?

[rocker-zhang]

Agreed — and your framing (modifying the MLP costs intelligence, so at matched KLD the optimum is usually not to touch it) is the cleaner way to put it. Fair on the budget too; I'll take the 150–200 convergence, the extra two dimensions aren't what's driving the gap.

On "what is the proposed change" — short answer: nothing anymore, I'm withdrawing it. To say what I was pointing at: the CPU test suite I described earlier (tests/test_fused_experts.py), which just checks that the fused experts are detected — the stacked 3D nn.Parameter transformers 5.x introduced when experts stopped being a ModuleList. On its own it guards nothing; it only has a point next to the model.py change that actually ablates those experts.

And that change isn't worth landing. I'd assumed the experts were structurally unreachable — that's why I closed #339 — but target_parameters does reach them; they just don't pay at matched KLD, as you've shown. So there's nothing to merge here, I'll close it out. Thanks for the detailed back-and-forth on this one.

[p-e-w]

Ok, thank you for the detailed investigation! I expect I will refer back to this discussion occasionally in the future.

Closing this for now; feel free to reopen if there is more to add.