feat(ara): Clamped KNN objective & adaptive steering (resolves DiffusionGemma / MoE failures) - #400
feat(ara): Clamped KNN objective & adaptive steering (resolves DiffusionGemma / MoE failures)#400umran666 wants to merge 19 commits into
Conversation
This term was found experimentally to be 3-4 orders of magnitude smaller than the others in most runs, and have no meaningful effect on the result of the optimization.
Co-authored-by: kabachuha <artemkhrapov2001@yandex.ru>
…-e-w#214) Extends d79a443 — that commit correctly moves I/O tensors to the weight matrix's device before L-BFGS optimization, but the captured tensors remain on their original GPU between trials. When reset_model() reloads the model, device assignments can change, leaving orphaned tensors on GPUs that now need that VRAM for the reloaded weights. Moving to CPU at capture time in get_module_io ensures: - Zero VRAM wasted on stale device assignments between trials - Clean CPU→target transfer regardless of how devices shuffle on reload - No overhead on single-GPU (.cpu() is a no-op when already on CPU, and .to(device) in ara_abliterate handles the final placement)
Incorporates feedback from @joninco
* ARA, but it's LoRA * ARA, but it's LoRA: address Gemini's review * ARA, but it's LoRA: Gemini is stupid
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces Arbitrary-Rank Ablation (ARA), an optimization-based abliteration method, along with a two-stage abliteration process (scout and strike passes) and support for DiffusionGemma models. The review feedback identifies several critical bugs and robustness issues, including a broken and redundant redefinition of empty_cache in utils.py that causes a NameError, fragile indexing logic in get_module_io that can lead to assertion failures, potential KeyErrors in MoE models when experts do not activate for both datasets, and a potential RuntimeError in mean_distances_to_knn when token counts are low. Additionally, the reviewer points out dead code, shadowing of Python builtins, and multiple repository style guide violations regarding missing type annotations, unapproved abbreviations, comments lacking periods, unrelated code changes, and missing updates to config.default.toml.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| def empty_cache(): | ||
| # Collecting garbage is not an idempotent operation, and to avoid OOM errors, | ||
| # gc.collect() has to be called both before and after emptying the backend cache. | ||
| # See https://github.com/p-e-w/heretic/pull/17 for details. | ||
| gc.collect() | ||
|
|
||
| if torch.cuda.is_available(): | ||
| torch.cuda.empty_cache() | ||
| elif is_xpu_available(): | ||
| torch.xpu.empty_cache() | ||
| elif is_mlu_available(): | ||
| torch.mlu.empty_cache() # ty:ignore[unresolved-attribute] | ||
| elif is_sdaa_available(): | ||
| torch.sdaa.empty_cache() # ty:ignore[unresolved-attribute] | ||
| elif is_musa_available(): | ||
| torch.musa.empty_cache() # ty:ignore[unresolved-attribute] | ||
| elif torch.backends.mps.is_available(): | ||
| torch.mps.empty_cache() | ||
|
|
||
| gc.collect() |
There was a problem hiding this comment.
This redefinition of empty_cache is redundant because it is already defined in src/heretic/system.py and imported from there. Furthermore, this implementation is completely broken because it references gc, is_mlu_available, is_sdaa_available, and is_musa_available which are not imported or defined in this file, leading to a NameError at runtime. Please remove this redundant function entirely.
| module_io: ModuleIO = [] | ||
| temporal_io: list[dict[str, dict[int, dict[str, list[Tensor]]]]] = [] | ||
|
|
||
| def get_hook( | ||
| layer_index: int, | ||
| component: str, | ||
| module_index: int, | ||
| ): | ||
| def hook( | ||
| module: Module, | ||
| inputs: tuple[Tensor, ...], | ||
| outputs: Tensor, | ||
| ) -> None: | ||
| if len(temporal_io) == layer_index: | ||
| temporal_io.append({}) | ||
|
|
||
| assert len(temporal_io) == layer_index + 1 | ||
|
|
||
| if component not in temporal_io[layer_index]: | ||
| temporal_io[layer_index][component] = {} | ||
|
|
||
| if module_index not in temporal_io[layer_index][component]: | ||
| temporal_io[layer_index][component][module_index] = {"inputs": [], "outputs": []} | ||
|
|
||
| inp = inputs[0] | ||
| out = outputs | ||
|
|
||
| # Dimensional safety for standard dense (3D) vs flattened MoE tensors (2D) | ||
| if inp.dim() == 3: | ||
| inp = inp[:, -1, :] | ||
|
|
||
| if out.dim() == 3: | ||
| out = out[:, -1, :] | ||
|
|
||
| temporal_io[layer_index][component][module_index]["inputs"].append(inp.detach().cpu()) | ||
| temporal_io[layer_index][component][module_index]["outputs"].append(out.detach().cpu()) | ||
|
|
||
| return hook | ||
|
|
||
| hook_handles: list[RemovableHandle] = [] | ||
|
|
||
| for layer_index in range(len(self.get_layers())): | ||
| for component, modules in self.get_layer_modules(layer_index).items(): | ||
| for module_index, module in enumerate(modules): | ||
| hook_handles.append( | ||
| module.register_forward_hook( | ||
| get_hook(layer_index, component, module_index) | ||
| ) | ||
| ) | ||
|
|
||
| trajectory_steps = getattr(self.settings, "ara_trajectory_steps", 1) | ||
| self.generate(prompts, max_new_tokens=trajectory_steps) | ||
|
|
||
| for hook_handle in hook_handles: | ||
| hook_handle.remove() | ||
|
|
||
| for layer_index in range(len(temporal_io)): | ||
| module_io.append({}) | ||
| for component, io_map in temporal_io[layer_index].items(): | ||
| module_io[layer_index][component] = {} | ||
| for module_index, trajectory in io_map.items(): | ||
| if trajectory["inputs"]: | ||
| try: | ||
| avg_input = torch.stack(trajectory["inputs"]).mean(dim=0) | ||
| avg_output = torch.stack(trajectory["outputs"]).mean(dim=0) | ||
| except RuntimeError: | ||
| # MoE experts may process varying token counts per step, causing stack() to fail | ||
| # In this case, we concatenate across the temporal axis | ||
| avg_input = torch.cat(trajectory["inputs"], dim=0) | ||
| avg_output = torch.cat(trajectory["outputs"], dim=0) | ||
|
|
||
| module_io[layer_index][component][module_index] = (avg_input, avg_output) | ||
|
|
||
| return module_io |
There was a problem hiding this comment.
Using a list for temporal_io and appending to it based on layer_index is extremely fragile. If any layer is skipped (e.g., has no targeted components or is skipped during two-stage abliteration), len(temporal_io) == layer_index will be False, causing the assertion assert len(temporal_io) == layer_index + 1 to fail. Furthermore, if len(temporal_io) is less than the total number of layers, module_io will be shorter than len(self.get_layers()), causing an IndexError when accessed by layer_index in ara_abliterate.\n\nTo fix this, use a dictionary for temporal_io and initialize module_io as a list of empty dictionaries of length len(self.get_layers()).
def get_module_io(
self,
prompts: list[Prompt],
) -> ModuleIO:
module_io: ModuleIO = [{} for _ in range(len(self.get_layers()))]
temporal_io: dict[int, dict[str, dict[int, dict[str, list[Tensor]]]]] = {}
def get_hook(
layer_index: int,
component: str,
module_index: int,
):
def hook(
module: Module,
inputs: tuple[Tensor, ...],
outputs: Tensor,
) -> None:
if layer_index not in temporal_io:
temporal_io[layer_index] = {}
if component not in temporal_io[layer_index]:
temporal_io[layer_index][component] = {}
if module_index not in temporal_io[layer_index][component]:
temporal_io[layer_index][component][module_index] = {"inputs": [], "outputs": []}
input_tensor = inputs[0]
output_tensor = outputs
# Dimensional safety for standard dense (3D) vs flattened MoE tensors (2D).
if input_tensor.dim() == 3:
input_tensor = input_tensor[:, -1, :]
if output_tensor.dim() == 3:
output_tensor = output_tensor[:, -1, :]
temporal_io[layer_index][component][module_index]["inputs"].append(input_tensor.detach().cpu())
temporal_io[layer_index][component][module_index]["outputs"].append(output_tensor.detach().cpu())
return hook
hook_handles: list[RemovableHandle] = []
for layer_index in range(len(self.get_layers())):
for component, modules in self.get_layer_modules(layer_index).items():
for module_index, module in enumerate(modules):
hook_handles.append(
module.register_forward_hook(
get_hook(layer_index, component, module_index)
)
)
trajectory_steps = getattr(self.settings, "ara_trajectory_steps", 1)
self.generate(prompts, max_new_tokens=trajectory_steps)
for hook_handle in hook_handles:
hook_handle.remove()
for layer_index in range(len(self.get_layers())):
if layer_index not in temporal_io:
continue
for component, io_map in temporal_io[layer_index].items():
module_io[layer_index][component] = {}
for module_index, trajectory in io_map.items():
if trajectory["inputs"]:
try:
average_input = torch.stack(trajectory["inputs"]).mean(dim=0)
average_output = torch.stack(trajectory["outputs"]).mean(dim=0)
except RuntimeError:
# MoE experts may process varying token counts per step, causing stack() to fail.
# In this case, we concatenate across the temporal axis.
average_input = torch.cat(trajectory["inputs"], dim=0)
average_output = torch.cat(trajectory["outputs"], dim=0)
module_io[layer_index][component][module_index] = (average_input, average_output)
return module_io| for component, modules in self.get_layer_modules(layer_index).items(): | ||
| for module_index, module in enumerate(modules): |
There was a problem hiding this comment.
In Mixture-of-Experts (MoE) models, some experts are highly specialized and might only activate for harmful prompts and never for harmless prompts (or vice-versa). If an expert is not activated for both datasets, accessing good_module_io[layer_index][component][module_index] or bad_module_io[layer_index][component][module_index] will raise a KeyError. Please add a defensive check to skip modules that are not present in both good_module_io and bad_module_io.
| for component, modules in self.get_layer_modules(layer_index).items(): | |
| for module_index, module in enumerate(modules): | |
| for component, modules in self.get_layer_modules(layer_index).items(): | |
| for module_index, module in enumerate(modules): | |
| if ( | |
| component not in good_module_io[layer_index] | |
| or module_index not in good_module_io[layer_index][component] | |
| or component not in bad_module_io[layer_index] | |
| or module_index not in bad_module_io[layer_index][component] | |
| ): | |
| continue |
| for component, modules in self.get_layer_modules(layer_index).items(): | ||
| for module_index, module in enumerate(modules): |
There was a problem hiding this comment.
In Mixture-of-Experts (MoE) models, some experts are highly specialized and might only activate for harmful prompts and never for harmless prompts (or vice-versa). If an expert is not activated for both datasets, accessing good_module_io[layer_index][component][module_index] or bad_module_io[layer_index][component][module_index] will raise a KeyError. Please add a defensive check to skip modules that are not present in both good_module_io and bad_module_io.
| for component, modules in self.get_layer_modules(layer_index).items(): | |
| for module_index, module in enumerate(modules): | |
| for component, modules in self.get_layer_modules(layer_index).items(): | |
| for module_index, module in enumerate(modules): | |
| if ( | |
| component not in good_module_io[layer_index] | |
| or module_index not in good_module_io[layer_index][component] | |
| or component not in bad_module_io[layer_index] | |
| or module_index not in bad_module_io[layer_index][component] | |
| ): | |
| continue |
| def mean_distances_to_knn(a: Tensor, b: Tensor, k: int) -> Tensor: | ||
| distances = torch.cdist(a, b) | ||
| nearest_distances, _ = distances.topk(k, dim=1, largest=False) | ||
| return nearest_distances.mean(1) |
There was a problem hiding this comment.
If the number of active tokens in b is less than k (which is highly likely for sparse MoE experts), distances.topk(k, dim=1) will raise a RuntimeError: selected index k out of range. Please clamp k to min(k, b.shape[0]) and handle the k == 0 edge case to ensure robustness.
| def mean_distances_to_knn(a: Tensor, b: Tensor, k: int) -> Tensor: | |
| distances = torch.cdist(a, b) | |
| nearest_distances, _ = distances.topk(k, dim=1, largest=False) | |
| return nearest_distances.mean(1) | |
| def mean_distances_to_knn(a: Tensor, b: Tensor, k: int) -> Tensor: | |
| distances = torch.cdist(a, b) | |
| k = min(k, b.shape[0]) | |
| if k == 0: | |
| return torch.zeros(a.shape[0], device=a.device, dtype=a.dtype) | |
| nearest_distances, _ = distances.topk(k, dim=1, largest=False) | |
| return nearest_distances.mean(1) |
| input = torch.cat( | ||
| [input_output[0] for input_output in inputs_outputs], | ||
| dim=0, | ||
| ) | ||
| output = torch.cat( | ||
| [input_output[1] for input_output in inputs_outputs], | ||
| dim=0, | ||
| ) | ||
|
|
||
| # The key already exists, and replacing existing values | ||
| # in a dictionary while iterating over the same dictionary | ||
| # is safe in Python. | ||
| module_io[layer_index][component][module_index] = (input, output) |
There was a problem hiding this comment.
The local variables input and output shadow the Python builtin functions input and output (or standard names), which is a bad practice and can lead to confusion or bugs. Please rename them to input_tensor and output_tensor.
| input = torch.cat( | |
| [input_output[0] for input_output in inputs_outputs], | |
| dim=0, | |
| ) | |
| output = torch.cat( | |
| [input_output[1] for input_output in inputs_outputs], | |
| dim=0, | |
| ) | |
| # The key already exists, and replacing existing values | |
| # in a dictionary while iterating over the same dictionary | |
| # is safe in Python. | |
| module_io[layer_index][component][module_index] = (input, output) | |
| input_tensor = torch.cat( | |
| [input_output[0] for input_output in inputs_outputs], | |
| dim=0, | |
| ) | |
| output_tensor = torch.cat( | |
| [input_output[1] for input_output in inputs_outputs], | |
| dim=0, | |
| ) | |
| # The key already exists, and replacing existing values | |
| # in a dictionary while iterating over the same dictionary | |
| # is safe in Python. | |
| module_io[layer_index][component][module_index] = (input_tensor, output_tensor) |
| verification_mode=VerificationMode.NO_CHECKS, | ||
| # But also don't use cached data, as the dataset may have changed on disk. | ||
| download_mode=DownloadMode.FORCE_REDOWNLOAD, |
There was a problem hiding this comment.
The removal of download_mode=DownloadMode.FORCE_REDOWNLOAD is an existing code change that is not directly related to the PR's objective of implementing ARA and adaptive steering. This violates Rule 11 of the Repository style guide ('PRs must not change existing code unless the changes are directly related to the PR. This includes changes to formatting and comments'). Please revert this change or split it into a separate PR.
split=split_str,
# Don't require the number of examples (lines) per split to be pre-defined.
verification_mode=VerificationMode.NO_CHECKS,
# But also don't use cached data, as the dataset may have changed on disk.
download_mode=DownloadMode.FORCE_REDOWNLOAD,References
- PRs must not change existing code unless the changes are directly related to the PR. This includes changes to formatting and comments. (link)
|
|
||
| def get_merged_model(self) -> PreTrainedModel: | ||
| if getattr(self.settings, "use_ara", False): | ||
| # ARA models do not use PEFT LoRA wrappers |
There was a problem hiding this comment.
This comment does not end with a period, violating Rule 2 of the Repository style guide ('Comments should start with a capital letter and end with a period. They should use correct grammar and spelling'). Please add a period.
| # ARA models do not use PEFT LoRA wrappers | |
| # ARA models do not use PEFT LoRA wrappers. |
References
- Comments should start with a capital letter and end with a period. They should use correct grammar and spelling. (link)
| del good_input, good_output, bad_input, bad_output, optimizer | ||
| empty_cache() | ||
|
|
||
| # Ensure base weights are reloaded on the next trial since they were modified in-place |
There was a problem hiding this comment.
This comment does not end with a period, violating Rule 2 of the Repository style guide ('Comments should start with a capital letter and end with a period. They should use correct grammar and spelling'). Please add a period.
| # Ensure base weights are reloaded on the next trial since they were modified in-place | |
| # Ensure base weights are reloaded on the next trial since they were modified in-place. |
References
- Comments should start with a capital letter and end with a period. They should use correct grammar and spelling. (link)
|
@p-e-w @kabachuha I've pushed the initial implementation for the clamped KNN objective and the Two-Stage adaptive steering. Before I go through and patch the formatting/structural flags raised by the gemini-code-assist bot, would you mind pulling this branch and giving it a test run on your local terminals? Since this introduces some fairly complex L-BFGS logic and 3D tensor handling specifically for MoE architectures, I'd appreciate your eyes on it. If you catch any edge cases, critical bugs, or logic failures during your testing, just drop the logs here and I'll roll the fixes into the same patch pass. |
|
@umran666 Can you give the config toml or the key settings you are using for this? |
|
Please target the |
8d33003 to
bc0d85b
Compare
i will change present there are some fixes and critical bugs....So it take time after the fixes i will change to ara branch |
@kabachuha hey while testing it, change the target_components to and also because of standard ara on multiple layers can be slow so use the following config too also i didnt test it on TUI, so if there any errors in it can you fix them clearly.... |
|
@umran666 Would it be hard for you to adapt the code for the recently merged ara-lora technique for 4bit bitsandbytes compatibility? Raw bfloat16 DiffusionGemma 26b doesn't fit into my setup :( |
@kabachuha Try to use ara-lora technique on my code cause its not worked for me due to rank implementation and it dont work on complex models but it works on small.. Implement it once if its fixes then commit to this PR |
|
@kabachuha also i have updated some patches while testing use target_components = ["attn.o_proj", "mlp.down_proj"] |
|
@kabachuha just checking in on this since it's been a week. Did you have any luck getting the ara-lora 4-bit approach to converge on your end, or did you end up hitting the same Rank 128 optimizer trap I ran into? |
…d DiffusionGemma support
|
Hi, I've been busy last week, but now I have some time. I can try it with smaller (non-diffusion) models first. |
|
@umran666 I tested your vanilla code with your settings right now on Qwen 3 4b 2507 instruct and got these results I'd say, they are not very ideal compared to the current ARA and SOMA. Maybe your code is too aggressive? 
|
@kabachuha Yeah those ranges are definitely too aggressive for a 4B dense model those tightened Optuna bounds (start_layer, end_layer, steer_bad_behavior_weight) were specifically tuned for DiffusionGemma's 30-layer MoE architecture where the refusal circuit is smeared across 128 experts. For standard dense models like Qwen 3 4B, you'd want to revert to the upstream default search ranges and only keep target_components = ["attn.o_proj"] (single component, no mlp.down_proj and dense models don't need the MLP hammer). The core contribution of this PR is the MoE robustness fixes (KeyError guards for sparse experts, topk clamping, temporal_io dict migration) and the target_components config — not the narrowed search space. I should probably split the DiffusionGemma-specific tuning into a separate config preset so it doesn't pollute the defaults for other models. Let me know how it goes with wider ranges. |
abc001c to
61f2135
Compare
|
@kabachuha @p-e-w Thanks for reviewing the initial implementation! I realized that because this branch was originally based off master, it ended up pulling in a lot of unrelated commits and merges from master This made the diff extremely messy and hard to review.To keep the history perfectly clean and make reviewing much easier, I will do a clean port of our ARA enhancements directly on top of a fresh checkout of the upstream/ara branch. |
4 participants
I’ve been spending a lot of time running ARA on
google/diffusiongemma-26B-A4B-it(a large MoE model with 28 layers) and ran into a massive wall: the standard ARA approach was either failing to suppress refusals, or it was completely destroying the token distribution (KL > 0.3, sometimes even blowing up into the 20s).Analysis:
Analyzing the routing probabilities for harmless vs. harmful prompts.The refusal concept was heavily smeared across multiple experts as a distributed structural attractor.We couldn't just zero out specific experts, we had to use ARA to gently steer the weights of the entire targeted module.
expert_routing_report.json
DuoNeural/diffusiongemma-26B-A4B-it-abliterated. That attempt used standard Heretic geometric subtraction (RepE/Orthogonal Projection). The result was catastrophic capability collapse (massive KL divergence).
Multi-Component Mass Steering (attn.o_proj + mlp.down_proj) Because the refusal attractor is so distributed, we achieved a 4/100 refusal rate at ~0.1 KLD by explicitly targeting both the standard attention outputs and the MoE down projections simultaneously across large, contiguous blocks of layers.
I am using my down_proj specifically in identified layers (26-29) and PR feat: diffusion LLM abliteration support (DiffusionGemma) #378 helps me for more data MoE experts are stored as a single massive batched tensor [128, 2816, 704] and DiffusionGemma shares the exact same weights in memory between its encoder and decoder
Clamped KNN Multi-Directional Overcorrection I upgraded the objective function to use a KNN-based multi-directional overcorrection. Instead of just pulling the bad embeddings toward the good mean, we added a negative loss term multiplied by overcorrect_relative_weight to actively push them away. Crucially, I mathematically clamped this negative loss and tightened the Optuna search ranges. Without the clamp, the optimizer would exploit the unbounded negative term, driving the loss to negative infinity and causing catastrophic KL explosion.
Per-Layer Adaptive Steering
Previously, the engine applied the exact same steer_bad_behavior_weight to all layers. I introduced steer_core_weight and steer_late_weight. By splitting the active layer range into thirds, Optuna could apply heavy steering pressure specifically to the core layers where the refusal circuit lives, while preserving the integrity of the early and late layers.
The Results:
I uploaded the optuna preferred model and also test this model
Umranz/diffusiongemma-26B-A4B-it-abliteration
Honorable Results:
1)KLD: 0.09 and Refusals: 12/100
2)KLD: 0.06 and Refusals: 14/100...etc
Compatibility Note:
I made sure this doesn't break anything for normal Heretic users.
None, so old checkpoints will still deserialize perfectly and fall back to uniform steering.Qwen/Qwen2.5-1.5B-Instructand it runs exactly like the upstream branch.Disclaimer: While this pipeline works flawlessly for my tests (and backward compatibility has been verified with Qwen 1.5B), this introduces a lot of complex routing and optimization logic. There may be edge cases or unhandled bugs I've missed. Please review the 3D tensor handling and L-BFGS objective clamping closely. Happy to fix any critical bugs identified during review.