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.
Summary of ChangesHello, 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
🧠 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
Activity
Using Gemini Code AssistThe 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
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 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
|
There was a problem hiding this comment.
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.
|
Congratulations on the conception!
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. |
|
Pareto frontier for Qwen3-4B-Instruct-2507. Qwen series are reportedly notoriously hard to decensor, so I decided to test it. Comparison with the other methods:
I'd say, the results are somewhere in-between |
|
@p-e-w Can you submit the gpt-oss model to the UGI leaderboard? https://huggingface.co/spaces/DontPlanToEnd/UGI-Leaderboard/discussions |
|
Interesting idea, and surprisingly straightfoward! A couple of initial thoughts/questions:
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 |
|
@kabachuha The weights in 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. |
I don't quite understand what you mean here. Could you explain more?
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. |
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.
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. |
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. |
|
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. |
|
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. |
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. |
Make sure you use the latest commit (0bb9521). If you are seeing suboptimal results, I recommend trying some combination of these:
|
|
Can you make some visualizations (ex. PCA) of the model's hidden states as the ARA method achieves convergence? |
|
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 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 Could be: Windows/PowerShell specific issue After switching from PowerShell (non-admin) to Admin CMD, error messages did not appear and trial completed successfully. |
Yes, I will do a full writeup explaining the motivation behind ARA, which will include such data. |
|
Did some tests with Qwen 3.5 4B. Main branch: Best trials still refuse more than 50 of 100 bad prompts. Can't wait for somebody making an ARA Version of Qwen 3.5 27B. |
|
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? |
|
@flashburns Thank you, that means a lot to me. |
|
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! |
|
@antis0007 Feel free to join the Discord (link in README). I'm usually around there for back-and-forth discussions. |
|
There is currently a serious bug in ARA where the resulting matrices sometimes contain Needless to say, this is a blocker for merging ARA into master. |
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. |
|
😄 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. |
|
Maybe just prune the trials? Nan check is good, though |
|
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. |
|
@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. |
We don't know that the models are the problem. It could be:
"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. |
|
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 |
|
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
|
What's the next part of the roadmap on ARA? |
|
Reproduce the broken tensors issue and fix it, or confirm that it doesn't happen with ARA-LoRA. |
|
I think there's a bug with LoRA here: When |
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. |
|
@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 ( 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 The reason that none of this surfaced is that Heretic computes its benchmark numbers from the in-memory model before 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 ( 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 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 Heretic never asks for that path, since 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 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 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 |
|
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. |
|
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. |
|
@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.
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:
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:
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:
steer_bad_behavior.Feedback welcome!
@spikymoth
@kabachuha
@red40maxxer