Skip to content

feat: Arbitrary-Rank Ablation (ARA) - #211

Draft
p-e-w wants to merge 19 commits into
masterfrom
ara
Draft

feat: Arbitrary-Rank Ablation (ARA)#211
p-e-w wants to merge 19 commits into
masterfrom
ara

Conversation

@p-e-w

@p-e-w p-e-w commented Mar 4, 2026

Copy link
Copy Markdown
Owner

Arbitrary-Rank Ablation (ARA) is a radically new abliteration method that I've been developing for the past two months or so. I believe that it can replace all currently implemented methods in Heretic, including MPOA, once the remaining issues are worked out. Its only serious competitor at this time is @kabachuha's implementation of multi-directional refusal suppression with Self-Organizing Maps (#196).

ARA doesn't use refusal directions at all, neither a single direction like traditional abliteration, nor multiple directions like SOMA. Instead, ARA works by capturing input/output tensors at each individual transformer module using PyTorch hooks, then uses direct, unconstrained matrix optimization to modify those modules, based on an objective function that captures the essence of what we want to (and don't want to) change.

Intuitively, the objective encodes three competing optimization goals:

  1. The outputs of the module for inputs associated with "harmless" prompts should change as little as possible.
  2. The outputs of the module for inputs associated with "harmful" prompts should become as similar to those associated with "harmless" prompts as possible.
  3. The outputs of the module for inputs associated with "harmful" prompts should become as dissimilar to those previously associated with "harmful" prompts as possible. In combination with (2), this overcorrects away from the original residuals, which results in stronger steering that can overcome more complex refusal mechanisms.

Unlike other abliteration methods, this approach doesn't assume a particular rank for the refusal manifold, or that the centroid of the outputs must shift in a specific manner. This gives the optimizer more freedom to modify the matrix in the best possible way. Please see the code for implementation details.

The objective is affine-convex and the initial value (the original matrix) is already very close to the optimum, so L-BFGS makes short work of it, typically converging in 2-3 iterations. Because the matrices are optimized one-by-one, the total memory requirements are barely higher than for regular abliteration. The abliteration process takes longer, but the time per trial is still dominated by counting refusals. Combined with the fact that ARA has fewer optimizable parameters than our current approach (meaning that fewer trials are needed for good results), this might actually make ARA faster than regular abliteration.

Results

For demonstration purposes, I have processed openai/gpt-oss-20b with the exact code currently in this pull request. The result is p-e-w/gpt-oss-20b-heretic-ara-v3:

ara-results

This is dramatically better than any existing abliteration of gpt-oss-20b (see this table), with the possible exception of the brand new kabachuha/gpt-oss-20b-SOMbliterated, which has the same refusal count but higher KL divergence.

TODO

ARA isn't quite ready for mainstream use yet, but it's getting close. The remaining issues are:

  • Figure out how to obtain good values for the weight constants in steer_bad_behavior.
  • Decide whether to have separate abliteration parameter sets for each component, like in our current implementation.
  • Fix multi-GPU issues
  • Write a proper article explaining the mathematics and motivation behind ARA.

Feedback welcome!

@spikymoth
@kabachuha
@red40maxxer

@gemini-code-assist

Copy link
Copy Markdown
Contributor

Summary of Changes

Hello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces Arbitrary-Rank Ablation (ARA), a radically new method for model steering that moves beyond traditional directional ablation. ARA leverages direct matrix optimization within individual transformer modules, guided by an objective function designed to minimize changes to harmless outputs while aggressively modifying harmful ones. This approach offers greater flexibility by not assuming a fixed refusal manifold rank, potentially leading to more robust and efficient abliteration results.

Highlights

  • Arbitrary-Rank Ablation (ARA) Introduction: Introduced a novel ablation method, Arbitrary-Rank Ablation (ARA), which uses direct, unconstrained matrix optimization on transformer modules instead of traditional directional ablation. This method aims to provide more flexible and effective model steering.
  • New Optimization Objective: Implemented a new objective function for ARA that balances three goals: preserving outputs for 'harmless' prompts, making 'harmful' prompt outputs similar to 'harmless' ones, and pushing 'harmful' prompt outputs away from their original state for stronger steering.
  • Conditional Ablation Logic: Integrated conditional logic throughout the codebase to switch between ARA and traditional ablation methods, affecting gradient handling, parameter optimization, module I/O capture, and model saving/loading.
  • Module I/O Capture with PyTorch Hooks: Developed functionality to capture input/output tensors at each transformer module using PyTorch hooks, which is crucial for the ARA optimization process.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog
  • src/heretic/config.py
    • Added target_components setting to specify which model components to abliterate.
    • Added use_ara setting to toggle between ARA and traditional ablation.
  • src/heretic/main.py
    • Modified gradient enabling to be conditional on the use_ara setting.
    • Updated the main run function to conditionally calculate refusal directions (traditional) or module I/O (ARA).
    • Adjusted the objective function in Optuna to suggest different parameters based on whether use_ara is enabled.
    • Modified model reset, abliteration call, and model saving/uploading logic to be conditional on use_ara.
  • src/heretic/model.py
    • Added imports for LBFGS, RemovableHandle, Callable, and TypeAlias.
    • Defined ModuleIO type alias for storing module input/output tensors.
    • Made _apply_lora and reset_model calls conditional on use_ara.
    • Updated get_layer_modules to filter components based on settings.target_components.
    • Implemented the ara_abliterate method, which performs matrix optimization using L-BFGS based on captured module I/O.
    • Added get_module_io and get_module_io_batched methods to capture input/output tensors from model modules using PyTorch hooks.
  • src/heretic/utils.py
    • Imported Tensor from torch.
    • Added mean_distances_to_knn utility function for calculating mean Euclidean distances to k-nearest neighbors.
    • Modified get_trial_parameters to return different sets of parameters depending on whether ARA is enabled.
Activity
  • The author, p-e-w, developed and implemented the Arbitrary-Rank Ablation (ARA) method over approximately two months.
  • Demonstrated results with openai/gpt-oss-20b processed using ARA, showing significantly improved refusal counts compared to existing methods.
  • Requested feedback from specific reviewers (@spikymoth, @kabachuha, @red40maxxer).
  • Outlined remaining TODOs for ARA, including optimizing weight constants, deciding on separate abliteration parameter sets for components, and writing a mathematical article.
Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help /gemini help Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces a new abliteration method called Arbitrary-Rank Ablation (ARA). The changes are extensive, touching configuration, main application logic, and the model implementation to support this new method. The implementation uses PyTorch hooks to capture module I/O and an L-BFGS optimizer to modify module weights directly. The changes are mostly gated behind a new use_ara setting.

My feedback focuses on ensuring consistency with the repository's style guide and improving maintainability. Specifically, I've pointed out missing configuration updates, inconsistent trial parameter handling, a missing type annotation, and some minor style guide violations.

Comment thread src/heretic/config.py
Comment thread src/heretic/main.py
Comment thread src/heretic/model.py
Comment thread src/heretic/model.py
@kabachuha

kabachuha commented Mar 4, 2026

Copy link
Copy Markdown
Contributor

Congratulations on the conception!

Expected all tensors to be on the same device, but got mat2 is on cuda:0, different from other tensors on
cuda:1 (when checking argument in method wrapper_CUDA_mm)

Needs fix for multiGPU

@p-e-w

p-e-w commented Mar 4, 2026

Copy link
Copy Markdown
Owner Author

Needs fix for multiGPU

Thanks for pointing this out. I don't have a multi-GPU setup myself, but I'll rent one to figure out where the problem is.

@kabachuha

Copy link
Copy Markdown
Contributor

Pareto frontier for Qwen3-4B-Instruct-2507.

Qwen series are reportedly notoriously hard to decensor, so I decided to test it.

? Which trial do you want to use? (Use arrow keys)
   [Trial 156] Refusals:  3/100, KL divergence: 2.9901
   [Trial  25] Refusals:  4/100, KL divergence: 1.1411
   [Trial 148] Refusals:  5/100, KL divergence: 0.2406
   [Trial 147] Refusals:  6/100, KL divergence: 0.2359
   [Trial 168] Refusals:  7/100, KL divergence: 0.1521
   [Trial 151] Refusals:  8/100, KL divergence: 0.1384
 » [Trial 154] Refusals: 10/100, KL divergence: 0.1345
   [Trial 170] Refusals: 17/100, KL divergence: 0.0770
   [Trial 174] Refusals: 31/100, KL divergence: 0.0749
   [Trial  73] Refusals: 45/100, KL divergence: 0.0390
   [Trial 142] Refusals: 56/100, KL divergence: 0.0237
   [Trial  86] Refusals: 87/100, KL divergence: 0.0188
   Run additional trials

Comparison with the other methods:

Model Refusals for "harmful" prompts KL divergence from original model for "harmless" prompts
Qwen/Qwen3-4B-Instruct-2507 (original) 100/100 0 (by definition)
kabachuha/Qwen3-4B-Instruct-2507-SOMbliterated 3/100 0.08
heretic-org/Qwen3-4B-Instruct-2507-heretic 5/100 0.07
Goekdeniz-Guelmez/Qwen3-4B-Instruct-2507-gabliterated 4/100 0.25
p-e-w/gpt-oss-20b-heretic-ara (this) 8/100 0.14

I'd say, the results are somewhere in-between

@kabachuha

kabachuha commented Mar 4, 2026

Copy link
Copy Markdown
Contributor

@p-e-w Can you submit the gpt-oss model to the UGI leaderboard? https://huggingface.co/spaces/DontPlanToEnd/UGI-Leaderboard/discussions

@spikymoth

Copy link
Copy Markdown
Contributor

Interesting idea, and surprisingly straightfoward! A couple of initial thoughts/questions:

  1. mean_distances_to_knn returns the distance to the k nearest neighbors. That makes sense to me for maximizing dissimilarity, but wouldn't you want something more like the opposite (the k farthest neighbors) for maximizing similarity?
  2. I think the magnitude-preserving part of MPOA could be added onto this pretty easily by modifying the closure to call the objective (loss) function on a norm-constrained matrix. Something like this:
                    if self.settings.row_normalization == RowNormalization.FULL:
                        # Get row norms for original matrix.
                        target_norms = torch.norm(matrix, dim=1, keepdim=True)

                        def closure() -> Tensor:
                            optimizer.zero_grad()
                            # Compute loss relative to norm-constrained matrix.
                            constrained_matrix = F.normalize(matrix, p=2, dim=1) * target_norms
                            loss = objective(constrained_matrix)
                            # Compute the projected gradient with respect to the constrained matrix.
                            loss.backward()
                            return loss
                    else:
                        def closure() -> Tensor:
                            optimizer.zero_grad()
                            loss = objective(matrix)
                            loss.backward()
                            return loss

@p-e-w

p-e-w commented Mar 4, 2026

Copy link
Copy Markdown
Owner Author

@kabachuha The weights in steer_bad_behavior need to be fine-tuned for each model (currently by hand). If you run it on other models you will get poor results unless you tune them. I've seen much, much better results with Qwen3-4B-Instruct-2507 during testing.

In the future, this will happen automatically via Optuna, though it's not as straightforward as it may seem because the possible ranges go over multiple orders of magnitude.

@p-e-w

p-e-w commented Mar 4, 2026

Copy link
Copy Markdown
Owner Author

@spikymoth

mean_distances_to_knn returns the distance to the k nearest neighbors. That makes sense to me for maximizing dissimilarity, but wouldn't you want something more like the opposite (the k farthest neighbors) for maximizing similarity?

I don't quite understand what you mean here. Could you explain more?

I think the magnitude-preserving part of MPOA could be added onto this pretty easily by modifying the closure to call the objective (loss) function on a norm-constrained matrix.

That's true, but I'm not convinced that preserving the magnitude is correct in general. If harmful and harmless prompts result in residuals of different magnitudes, then abliteration should change the magnitudes I think.

@spikymoth

Copy link
Copy Markdown
Contributor

mean_distances_to_knn returns the distance to the k nearest neighbors. That makes sense to me for maximizing dissimilarity, but wouldn't you want something more like the opposite (the k farthest neighbors) for maximizing similarity?

I don't quite understand what you mean here. Could you explain more?

Well, do I understand correctly that it's 1. getting the distance between each pair of vectors, 2. selecting the k smallest distances and 3. returning the mean distance?

If so, it seems to target the harmful outputs that are already most similar to the harmless outputs and push them closer, while ignoring more dissimilar outputs. I'm wondering if this creates representative differences.

Perhaps the optimal way to contrast representative differences would be to contrast the top-k SOM neurons (as the center of gravity for each cluster) for a set of outputs.

That's true, but I'm not convinced that preserving the magnitude is correct in general. If harmful and harmless prompts result in residuals of different magnitudes, then abliteration should change the magnitudes I think.

IIRC the argument is that the row norms of the weight matrix overall should stay unchanged in order to preserve between-layer interpretability, i.e. each dimension in the output is expected to have a particular activation strength and if you change it, subsequent layers may get confused about what stronger/weaker activations mean.

@p-e-w

p-e-w commented Mar 4, 2026

Copy link
Copy Markdown
Owner Author

If so, it seems to target the harmful outputs that are already most similar to the harmless outputs and push them closer, while ignoring more dissimilar outputs.

No, it targets all outputs.

It computes the mean distance to the k nearest harmless neighbors for each harmful output and then computes the mean of those means. So every harmful output is attracted towards its nearest harmless neighbors.

This is actually precisely where the strength of this method comes from, because directional ablation based on a difference of means optimizes towards a configuration where the mean of the modified harmful outputs resembles the mean of the harmless outputs. This is an unnecessary constraint that hinders finding an optimal configuration. With ARA, every single harmful output is simply attracted towards somewhere in the harmless cluster. There is no requirement that the means of the outputs align.

@spikymoth

Copy link
Copy Markdown
Contributor

Aahh OK, I was misunderstanding the operation. So for every harmful output it computes the distance from all harmless outputs, then takes the mean of the k smallest distances (nearest neighbors) to push every output toward those neighbors. That makes sense, and should naturally give more weight to directions that show up more frequently.

@red40maxxer

red40maxxer commented Mar 5, 2026

Copy link
Copy Markdown
Contributor

This is an awesome idea and I'm really looking forward to this being included in main, I'm running a bunch of tests right now to see how this performs vs standard MPOA and kabachuha's SOM technique.

Is there any possible way to get quantization to work with this? I understand that ARA does gradient-based optimization on the weight matrix and bnb would create a different shape which breaks this, but being able to use quantization with this new technique would still be very valuable IMO. Maybe we could dequantize before ARA? This might not reduce the total RAM required for abliteration but it might at least speed up inference, which can be a bottleneck.

@p-e-w

p-e-w commented Mar 5, 2026

Copy link
Copy Markdown
Owner Author

Maybe we could dequantize before ARA?

Yes, this should work. Matrices are processed one by one, so the memory impact of dequantizing an individual matrix to full precision should be relatively small.

@p-e-w

p-e-w commented Mar 5, 2026

Copy link
Copy Markdown
Owner Author

I'm running a bunch of tests right now to see how this performs vs standard MPOA and kabachuha's SOM technique.

Make sure you use the latest commit (0bb9521). If you are seeing suboptimal results, I recommend trying some combination of these:

  1. Remove mlp.down_proj from target_components.
  2. Play with the value range for steer_bad_behavior_weight.

@kabachuha

Copy link
Copy Markdown
Contributor

Can you make some visualizations (ex. PCA) of the model's hidden states as the ARA method achieves convergence?

@erm14254

erm14254 commented Mar 5, 2026

Copy link
Copy Markdown
Contributor

I encountered a bug, I thought I should report it, here:

L-BFGS IndexError on Windows with high steer_bad_behavior_weight

Environment:

Windows 11 Pro
Python 3.12
PyTorch 2.10.0+cu130

Issue: Trial failed with IndexError: list index out of range in torch/optim/lbfgs.py line 205 during _strong_wolfe.

Failed parameters: steer_bad_behavior_weight = 0.3967

Also saw this symlink error (possibly related): [WinError 1314] A required privilege is not held by the client: '...triton_kernels_init_.py'

What failed: L-BFGS IndexError on first trial
Environment: PowerShell (non-admin) on Windows 11
What worked: Command Prompt as Administrator
Error: IndexError: list index out of range in torch\optim\lbfgs.py line 205
Parameters that crashed: steer_bad_behavior_weight = 0.3967 (much higher than your working trials at ~0.02-0.09)

Could be:

Windows/PowerShell specific issue
Numerical instability with high steer_bad_behavior_weight values
Admin privileges needed for triton kernels

After switching from PowerShell (non-admin) to Admin CMD, error messages did not appear and trial completed successfully.

@p-e-w

p-e-w commented Mar 6, 2026

Copy link
Copy Markdown
Owner Author

Can you make some visualizations (ex. PCA) of the model's hidden states as the ARA method achieves convergence?

Yes, I will do a full writeup explaining the motivation behind ARA, which will include such data.

@GhostWithAHat

Copy link
Copy Markdown

Did some tests with Qwen 3.5 4B.

Main branch: Best trials still refuse more than 50 of 100 bad prompts.
After merging ara into main and not touching the values for the weight constants in steer_bad_behavior: Good trials refuse 4 of 100 harmfull prompts with a KLD of 0.1396.

Can't wait for somebody making an ARA Version of Qwen 3.5 27B.

@p-e-w

p-e-w commented Mar 7, 2026

Copy link
Copy Markdown
Owner Author

@kabachuha

I am unable to reproduce the multi-GPU issue. I have tried processing Gemma 3 27B (which is 55 GB in BF16) on a 2x 5090 system, forcing tensor sharding. However, I am not getting a device mismatch error like you did.

Could you give some more information about the system where this error occurred?

@p-e-w

p-e-w commented Apr 10, 2026

Copy link
Copy Markdown
Owner Author

@flashburns Thank you, that means a lot to me.

@antis0007

Copy link
Copy Markdown

Hello! I'm incredibly interested in the possibility of integrating ARA into a Gemma4 E2B-it/E4B-it model converter script that I've been working on. I've been trying to create export scripts for Torch to Litert-lm model conversion, with heretic used as a intermediate step. ARA is going to be a revolution in model decensoring, your uploaded ARA models are very high quality. My goal is to make the process of exporting uncensored multimodal models for edge devices like smartphones much simpler. Lmk if there's any way to get in touch to discuss the idea with you sometime!

@p-e-w

p-e-w commented Apr 14, 2026

Copy link
Copy Markdown
Owner Author

@antis0007 Feel free to join the Discord (link in README). I'm usually around there for back-and-forth discussions.

@p-e-w

p-e-w commented Jun 18, 2026

Copy link
Copy Markdown
Owner Author

There is currently a serious bug in ARA where the resulting matrices sometimes contain NaN or inf/-inf. This has been reported by at least two users, and I believe it is also the reason why p-e-w/gpt-oss-20b-heretic-ara-v4 was corrupted on export (see above). The exact cause is currently unclear, and it's also unknown whether this bug happens with ARA-LoRA (#332) as well.

Needless to say, this is a blocker for merging ARA into master.

@erm14254

erm14254 commented Jun 18, 2026

Copy link
Copy Markdown
Contributor

There is currently a serious bug in ARA where the resulting matrices sometimes contain NaN or inf/-inf.

Yes it happens but it's very very rare, it mostly happens a lot more when a model doesn't really like ARA which is even rarer as almost all models have no issues with ARA.

@p-e-w

p-e-w commented Jun 18, 2026

Copy link
Copy Markdown
Owner Author

😄 Yeah but we need to figure out what it means that a model "doesn't like ARA".

I suspect it's an instability in L-BFGS caused by the aforementioned directional discontinuity of the gradient.

@kabachuha

Copy link
Copy Markdown
Contributor

Maybe just prune the trials? Nan check is good, though

@p-e-w

p-e-w commented Jun 18, 2026

Copy link
Copy Markdown
Owner Author

We could prune, but I'd prefer to figure out what is going on. We can't make ARA the default while it has such instabilities that we don't understand.

@erm14254

erm14254 commented Jun 18, 2026

Copy link
Copy Markdown
Contributor

We could prune, but I'd prefer to figure out what is going on. We can't make ARA the default while it has such instabilities that we don't understand.

Overall ARA is the best "jack of all trades", MPOA can achieve better results than ARA sometimes but only on a few select models.

Basically,

ARA: Very likely to give you good to great usable results on most models, very few models who "do not like ARA" resulting in bad or less than ideal results (model's accuracy loss is too high despite low KL divergence) and, very very few models and or certain conditions do not work with ARA.

MPOA: Unknown until you run hundred of trials, it may give good usable results, sometimes even better than ARA on certain select models, or it might top at 68/100 refusals after 600 trials, hence wasting time and making you go use ARA instead which will mostly likely give you actual usable results.

@kabachuha

Copy link
Copy Markdown
Contributor

@erm14254 There is also SOM / SOM-POA to consider. It can give great results as well, though from experience and Huggingface ARA is a more robust method.

I think MPOA < SOMPOA < ARA, as of now.

@p-e-w

p-e-w commented Jun 18, 2026

Copy link
Copy Markdown
Owner Author

models who "do not like ARA"

We don't know that the models are the problem.

It could be:

  • Individual models/tensors
  • Specific ARA hyperparameter combinations
  • Random behavior from L-BFGS (which isn't deterministic)
  • Convergence issues due to gradients or some other problem

"It works most of the time" isn't good enough. This PR is not going in until this problem is fully understood. Hopefully I'll have time to work on ARA again soon.

@rootacc3ss

Copy link
Copy Markdown

I haven't kept up to date on this, but am still really curious. Let me know if it's worth trying to test on one of the newer GLM models, or Kimi K2.7 code... Would be willing to spend the money for the community

@p-e-w

p-e-w commented Jun 21, 2026

Copy link
Copy Markdown
Owner Author

@rootacc3ss

Please try #332 instead, it's the version we are actually going to merge eventually.

* ARA, but it's LoRA

* ARA, but it's LoRA: address Gemini's review

* ARA, but it's LoRA: Gemini is stupid
@kabachuha

Copy link
Copy Markdown
Contributor

What's the next part of the roadmap on ARA?

@p-e-w

p-e-w commented Jul 5, 2026

Copy link
Copy Markdown
Owner Author

Reproduce the broken tensors issue and fix it, or confirm that it doesn't happen with ARA-LoRA.

@red40maxxer

Copy link
Copy Markdown
Contributor

I think there's a bug with LoRA here:

When use_ara_lora=True and row_normalization=Full, ara_lora_abliterate evals a row-normalized effective weight for the objective, but the deployed model computes the un-normalized one. I'm not sure exactly what the fix is since the idea of row-norm preservation for a LoRA adapter doesn't make sense, but maybe we shouldn't allow this to happen in the first place?

@p-e-w

p-e-w commented Jul 22, 2026

Copy link
Copy Markdown
Owner Author

I'm not sure exactly what the fix is since the idea of row-norm preservation for a LoRA adapter doesn't make sense

We already combine LoRAs with magnitude preservation in regular abliteration, but this situation is different because we want to preserve norms while optimizing, because post-optimization projection will break many of the properties we are trying to establish.

One approach could be to merge the LoRA into the full matrix at every optimization step, then reparametrize, then subtract the full matrix to get the (now full-rank) LoRA back, then cull dimensions using singular values. But that's very ugly, there might be a better solution.

@hpnyaggerman

hpnyaggerman commented Jul 27, 2026

Copy link
Copy Markdown

@p-e-w, I took apart gpt-oss-20b-heretic-ara-v4, the model you pointed to as corrupted, and the cause appears to be none of the four possibilities you've listed in your message earlier in this thread; it appears instead to be a serialization bug in transformers, and it reproduces with no ARA, no Heretic and no optimizer anywhere in the process. The reproducer is a load and a save of the unmodified base model on transformers 5.4.0.

m = AutoModelForCausalLM.from_pretrained("openai/gpt-oss-20b", dtype="auto", device_map="cuda:0")
assert type(m.model.layers[0].mlp.experts).__name__ == "Mxfp4GptOssExperts"
m.save_pretrained("out")

That writes 459 keys, 96 of them doubled (*_blocks_blocks, *_blocks_scales), and the resulting key set is identical to the published artifact's, symmetric difference zero, while the same script on 5.14.1 gives 0 doubled and 96 canonical. The root cause is pattern = pattern.removesuffix("$") at core_model_loading.py:474, where stripping the anchor means that down_proj also matches inside down_proj_blocks, so that the revert_weight_conversion call at modeling_utils.py:3354 renames keys that get_state_dict_and_metadata had already written correctly earlier in the same function. It was fixed in 5.5.2, and the expert biases came through intact only because Mxfp4ReverseDeserialize.convert carries an explicit bias-only workaround for the same problem.

On the four possibilites you have listed as possible causes, none of them seem to have much merit as the cause of corruption in the model you've pointed to. The suggestion that individual models or tensors are at fault does not explain why the damage covers all 96 expert tensors across all 24 layers uniformly, including layers 0 and 1, which ARA never writes, or why the whole of it reproduces on the untouched base model. ARA hyperparameters cannot be involved in a failure that reproduces with ARA not running at all, and there is no optimizer in the reproducer, so L-BFGS nondeterminism is also not it, while all 96 expert payloads are, separately, byte-identical to upstream by per-tensor SHA-256, so nothing ever wrote them. Gradients and convergence fall under the same exclusion, and the ARA deltas themselves look healthy - relative Frobenius norm between 0.0060 and 0.0303, per-row norms within 0.8%, no dead rows, no blow-up.

The write set of ARA is in any case disjoint from the damage, since it wrote 22 self_attn.o_proj.weight in layers 2 through 23, matching the declared range, while the broken set is 96 expert tensors in layers 0 through 23, and layers 0 and 1, being byte-identical to upstream, were demonstrably never touched by ARA, and yet all eight of their expert tensors are mis-keyed. As for timing, the model was uploaded on 2026-04-02 and its own config records transformers 5.4.0, which identifies the version that wrote it rather than merely the version it was built from: transformers_version is assigned unconditionally from the running library at configuration_utils.py:983, and the base model's config records 4.55.0.dev0, so the field is overwritten on save rather than inherited. The defect window is 5.4.0 through 5.5.1, with 5.5.2 carrying the fix on 2026-04-09. 5.3.0, which is what uv.lock on this branch actually pins, saves correctly, so the lockfile environment would not have produced this; the run that did used a transformers newer than the lock. I established that window by running the load-and-save above on each version rather than by reading the code, which matters because a static reading of the rename functions gets 5.3.0 wrong: the doubling only shows up end to end.

The reason that none of this surfaced is that Heretic computes its benchmark numbers from the in-memory model before save_pretrained, and the in-memory model was correct, while the save does not raise, and reloading the corrupt file does not raise either, on any version I tried up to 5.14.1. That does give a way of triaging the older reports, in that a model which benchmarked well during the run and was broken when somebody downloaded it is this bug, whereas a model that failed to converge during the run is not, and that class of report stays open.

On the ARA-LoRA question you raised, the artifact predates it (3b70fe5 vs 25979ad), so it says nothing directly, though the bug sits in serialization, downstream of and independent of how the weights are produced, which means that ARA-LoRA on anything from 5.4.0 through 5.5.1 would be corrupted identically (or so I think, at least), while 5.3.0 and anything from 5.5.2 onward are unaffected. Where it exports as a LoRA adapter (strategy == "adapter") the expert tensors are never written at all, so only the full-model save is exposed, and with use_ara set that branch saves model.model directly, which makes it the one that matters here.

The file is repairable in place and needs no re-abliteration, the operation being a rename of the 96 header keys and a pad of the header back to its original length, after which the payload comes out byte-identical and cmp -i 53448 exits 0. Repaired, the model scores PIQA 0.7731 with 6/100 refusals against a base control of 98/100 measured in the same session, while unrepaired it sits at 0.5267 +- 0.0116, which is essentially chance, that last figure being measured through the dequantize path because below compute capability 10.0 the mis-keyed file cannot complete a single forward pass on the MXFP4 kernels at all, the experts having been left as plain uninitialised tensors rather than the strided ones the kernel expects. To be precise, 0.7731 is not a digit-for-digit reproduction of the card's 0.7742 and 9/100, since identical weights do not force identical logits across a different transformers, triton and GPU, and --batch-size does not reach PIQA in any case, Evaluator.get_score constructing HFLM(..., batch_size="auto"), which selected 32 here.

Pinning a known-good transformers does not, however, settle the matter on its own. There is a second bug on the MXFP4 dequantize save path, and it was introduced in 5.5.2, which is the same release that fixed the first one. Across versions, run end to end: 5.3.0 and 5.4.0 raise NotImplementedError from core_model_loading.py:101 and write nothing at all, 5.5.0 and 5.5.1 save correctly with all 48 per-layer expert keys intact, and from 5.5.2 onward every layer's experts collapse onto the two literal keys down_proj$ and gate_up_proj$, the reverse mapping having used the pattern string itself as the output key rather than substituting the match into the per-layer name, so that all 24 layers write to the same two keys and only the last of them survives. Every one of the twenty-one releases from 5.5.2 to 5.14.1 collapses, checked on a two-layer cut of the model which reproduces the full-scale outcome exactly at 5.4.0, 5.5.1, 5.5.2 and 5.14.1. That is a different order of damage from the first bug. There every payload was present and only the names were wrong, which is why renaming recovers it; here 23 layers of expert weights are overwritten before anything reaches disk, leaving 5.2 GB instead of roughly 42 GB, with 1.6 GB of expert tensors where 24 layers belong, in a file that cannot be repaired at all and can only be regenerated.

Heretic never asks for that path, since quantization defaults to none and model.py passes no quantization_config at all for gpt-oss, but transformers decides for itself from the checkpoint's own quant_method: "mxfp4" and sets dequantize whenever compute capability is below 7.5 or triton or kernels is missing. This branch's lock carries both triton and kernels, so the live trigger is capability, which means a V100 or anything older is likely to be affected. On the currently locked 5.3.0 such a user gets a loud crash at save time, whereas bumping to ~=5.6 to escape the first bug moves that same user from a loud crash to the silent loss of 23 of 24 layers. So the fix for the first bug introduces some users to the second.

Given that, I would suggest a check after saving that the set of tensor names written matches the set in the source checkpoint, an invariant which always holds for Heretic, since it only rewrites values and never renames anything, and which both bugs violate. It cannot prevent either of them, but it turns a silent loss of 23 layers into a loud failure before anything reaches the Hub, and for the second bug it is the only protection there is, given that the fix for the first one is what introduced it. It reads safetensors headers only, so it costs nothing. I tried it against six checkpoints, and it passes the repaired artifact and a correct 5.14.1 save, and flags the published artifact, a 5.4.0 save, an explicit dequantize-path save and one that transformers dequantized on its own without being asked, and on this branch the natural place for it is straight after the Model saved to line in main.py, around line 898, on the non-adapter path, where there is currently no post-save verification of any kind.

On the reach of the first bug, it is narrow, more narrow than one would perhaps expect. First, stripping an anchor cannot affect a pattern that does not have one, so the only converters at risk are those whose target pattern ends in $, and in the whole of transformers 5.4.0 exactly two do, down_proj$ and gate_up_proj$ (both in quantizer_mxfp4.py). Having run the function over every target_patterns literal in the library, I get precisely those two modified and everything else returned unchanged, the MoE renamings in conversion_mapping.py among them, which carry no trailing $ for it to remove. Second, those two converters exist only while the MXFP4 quantizer is active, and the quantizer is reached through exactly one dispatch, AutoHfQuantizer.from_config, keyed on quant_method as read from the checkpoint's own config; no architecture opts into it, and models/gpt_oss does not so much as name mxfp4. Both conditions being required together, the damage cannot extend past checkpoints that declare quant_method: "mxfp4", and in practice that means gpt-oss. A dense bf16 model never brings those two converters into existence at all, so nothing here can reach ARA models on other architectures, which explains why the failure mode was never seen on, for example, Gemma 4 (most well known example of an ARA-abliterated model, to me at least). What determines whether or not the failure will happen is the quantization of the checkpoint, not anything ARA does and not any property of the model itself.

On your other question, whether this is the only mode of corruption, I cannot tell you that it is, having had the one artifact to work from, and what is settled here is the broken tensors in that one and nothing wider. But I would treat any report resting on a saved MXFP4 MoE checkpoint as unusable evidence about ARA until its keys have been checked, because the in-process numbers cannot detect this failure.

For anyone reproducing this, the environment was four NVIDIA A16-16Q vGPUs at compute capability 8.6 with 16 GB each, on a 24-vCPU Cascadelake host with 251 GB of RAM running Debian 13 on kernel 6.12.95. The guest driver is 550.54.15, which exposes CUDA 12.4, so torch's cu130 build runs on top of the cuda-compat-13-1 forward-compatibility package (590.48.01) through LD_LIBRARY_PATH; the driver itself needed three small patches to build against kernel 6.12, none of them touching anything numerical. Python was 3.12.13 throughout, with torch 2.13.0+cu130, triton 3.7.1, accelerate 1.14.0 and numpy 2.5.1, and with kernels pinned per transformers version, since those bounds are declared only under the [kernels] extra and an out-of-range kernels breaks the import: 0.10.5 with 5.3.0, 0.12.3 with 5.4.0 and the 5.5.x line, 0.15.2 with 5.14.1. The twenty-one-release sweep for the second bug ran without kernels installed at all, the dequantize path having no use for it. It is worth explicitly differentiating the two kinds of result here, in that the key-naming ones are pure string handling on the save path and are hardware-independent, so they should reproduce anywhere, whereas the PIQA and refusal figures, and the compute-capability assertion, are hardware-dependent and could be different on different silicon.

@hpnyaggerman

Copy link
Copy Markdown

With this issue out of the way (if no holes are poked in my analysis above), I hope work on ARA will be resumed and that it will soon be merged into master. ARA + custom objectives (which are already merged) would make heretic truly unstoppable, I believe.

@p-e-w

p-e-w commented Jul 28, 2026

Copy link
Copy Markdown
Owner Author

That's incredible detective work, @hpnyaggerman! Thank you so much for figuring this out.

Unfortunately, this is in line with what I have already noticed several times: That Transformers 5 is of poor quality overall, with early versions showing serious bugs even with popular models. I did consider the possibility of a Transformers bug in this case also, but then I remembered that I had successfully exported GPT-OSS with Transformers 5 before, so I dismissed it. What I failed to realize is that the bug could have been introduced post-5.0, which you have now confirmed.

Given your deep insight into the problem, could you open a PR implementing the check you suggested?

I'm currently wrapping up Heretic's "sister project" (which, incidentally, might be useful for getting ARA over the finish line), and then it's back to this PR.

@hpnyaggerman

Copy link
Copy Markdown

@p-e-w Very well, I will get to that. Should be straightforward enough.

LBFGS leaves one full-size gradient buffer on every optimized weight.
Across the layers processed in a typical trial this is many GiB of VRAM
that persists into evaluation, causing CUDA out-of-memory errors.
Clear the gradients after each module is optimized.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.