feat: diffusion LLM abliteration support (DiffusionGemma)

[anurag12-webster]

i managed to abliterate the newer diffusiongemma, i am attaching the changed... which i did to the base herectic codebase.. with some quirks

of all the code four things needed patching to make it work..

• switched PEFT task type to FEATURE_EXTRACTION diffusiongemma (DG) doesn't implement prepare_inputs_for_generation so CAUSAL_LM was crashing.
• MoE experts are a single batched parameter [128, 2816, 704], not nn.Linear, so heretic skips them. added per-slice ablation across all 128 experts per layer
• encoder and decoder share the same weights in memory (same data_ptr) PEFT only wraps encoder so decoder never saw the LoRA delta. i have fixed it with a context manager that merges LoRA into base weights before generate() method.
• the generate() method has no output_hidden_states support so i switched it to forward hooks on encoder layers for the PCA

code is rough, happy to refactor. closes #370

[p-e-w]

Awesome, thanks! I'll be looking at this in detail after the 1.4 release. Perhaps we can find another diffusion model to make sure this approach generalizes.

[p-e-w]

MoE experts are a single batched parameter [128, 2816, 704], not nn.Linear, so heretic skips them. added per-slice ablation across all 128 experts per layer

See #342 for an extended discussion of this.

/cc @rocker-zhang

[gemini-code-assist]

Code Review

This pull request adds support for the DiffusionGemma model, implementing custom logic for handling its tied encoder/decoder weights, expert weights abliteration, and generation/residual/logprob extraction. Feedback on the changes highlights a potential NameError if the model class cannot be imported, potential runtime errors or weight corruption when using 4-bit quantization, and a loss of reproducibility due to the removal of the RNG seed reset before SVD. Additionally, several style guide violations were identified, including missing return type annotations on new methods and improperly formatted comments.

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.

[kabachuha]

What are your actual numbers: refusals / KL?

Also, did you test it with ARA/ARA-LoRA, or only with projected abliteration?

For example, this guy's's results are not as optimistic and this has large KL (~0.5)

[kabachuha]

Perhaps we can find another diffusion model to make sure this approach generalizes.

LLaDA 1.x/2.x series would be a good candidate, they are well implemented and documented

[kabachuha]

From https://huggingface.co/edwixx/diffusiongemma-26B-A4B-it-HERETIC-Uncensored

Known issue: "own" tokens

Because the KL divergence is 0.49, some token positions in generated outputs don't fully denoise during the diffusion process. These positions fall back to the model's mask token, which decodes as the word "own".

It's very interesting, because the "own" token spawns not only in the diffusion model version, but in the base autoregressive one as well. For example, in my overcooked LoRA training it has shown up too.

Maybe, it's some sort of a placeholder token for Gemma4 models in general?

[anurag12-webster]

What are your actual numbers: refusals / KL?

Also, did you test it with ARA/ARA-LoRA, or only with projected abliteration?

For example, this guy's's results are not as optimistic and this has large KL (~0.5)

I have tried only with the project abliteration haven't tried with the ARA yet.

also for the numbers model currently has 13/100 refusals, KLD 0.49, i got this at the trial 89/200.

[anurag12-webster]

From https://huggingface.co/edwixx/diffusiongemma-26B-A4B-it-HERETIC-Uncensored

Known issue: "own" tokens

Because the KL divergence is 0.49, some token positions in generated outputs don't fully denoise during the diffusion process. These positions fall back to the model's mask token, which decodes as the word "own".

It's very interesting, because the "own" token spawns not only in the diffusion model version, but in the base autoregressive one as well. For example, in my overcooked LoRA training it has shown up too.

Maybe, it's some sort of a placeholder token for Gemma4 models in general?

Ohh i see, so my assumption abt the "own" artifact was being a denoising fallback was wrong the, and it is from gemma 4 series itself, i will update the model card with this then, thanks for confirming!

[kabachuha]

@anurag12-webster Do you have some visualizations: hidden states PCA/PaCMAP? Seeing how the refusal direction (or its equivalent) behaves in space would be incredibly helpful and useful for interpretability

[anurag12-webster]

@anurag12-webster Do you have some visualizations: hidden states PCA/PaCMAP? Seeing how the refusal direction (or its equivalent) behaves in space would be incredibly helpful and useful for interpretability

i actually dont have it, maybe in some time i'll have it run through and shared here

[anurag12-webster]

i'll try the Llada series of models with this one and also share some visualizations along with that.