feat(ara): add DiffusionGemma block-diffusion model support

[umran666]

This PR extends Heretic's Arbitrary-Rank Ablation (ARA) pipeline to support DiffusionGemma, allowing matrix optimization and abliteration on non-autoregressive, block-diffusion language models.

Standard transformers generate text sequentially, whereas DiffusionGemma performs iterative denoising over a fixed token canvas and encasing its transformer stack inside a custom bidirectional wrapper. This PR adapts layer discovery, generation decoding, and KLD evaluation to handle this paradigm cleanly without affecting standard causal language models.

Key Technical Adaptations
Architecture Navigation & LoRA Compatibility: Added structural checks to route layer extraction through model.encoder.language_model when DiffusionGemma is detected. To prevent PEFT from crashing during adapter initialization on models lacking traditional causal generation methods, a safe fallback is implemented for prepare_inputs_for_generation.

Deterministic KLD Proxy for Diffusion Models: Because block-diffusion generation does not expose per-token probability distributions via standard .scores, get_logprobs() implements a deterministic forward pass over a zero-initialized token canvas (decoder_input_ids). This produces a repeatable probability distribution before and after matrix steering, giving Optuna an accurate KL-divergence metric to penalize destructive weight updates.

Robustness Improvements: Fixed a BFloat16 dtype incompatibility in torch.cdist during KNN distance calculations and made system-role message injection conditional in chat templates to avoid tokenization errors on models with strict message schemas.

Analytical Note

During experimentation with both Standard ARA (use_ara = true) and Direct LoRA ARA (use_ara_lora = true), we observed that full-weight matrix optimization massively outperforms direct LoRA adapter optimization under the KNN retrieval objective. We wanted to share our theoretical analysis on why this occurs:

1. Rank & Geometric Capacity: Standard ARA operates directly on the full $d \times d$ weight matrix (e.g., $4096 \times 4096$). Shifting activation vectors away from refusal clusters across hundreds of training prompts requires high-dimensional tensor rotations. When constraining updates to a low-rank factorization ($\Delta W = B \cdot A$), a rank-8 or rank-16 matrix lacks the degrees of freedom required to resolve multi-directional steering targets without inducing severe feature collapse.

2. Bilinear Non-Convexity & L-BFGS Stalling: Standard ARA optimizes a single matrix $W$, presenting a clean linear forward mapping for solver line-search. Conversely, optimizing two matrices jointly via their product $(A, B) \mapsto B \cdot A$ creates scale ambiguities ($B \rightarrow \alpha B,; A \rightarrow \frac{1}{\alpha}A$) and severe saddle points where L-BFGS curvature estimates collapse and stall early in local minima.

3. Recommended Workflow: For users who require lightweight adapter artifacts, we recommend running Standard ARA first to find the optimal full-rank matrix $W_{\text{new}}$, then extracting the adapter post-hoc via SVD on the resulting delta ($\Delta W = W_{\text{new}} - W_{\text{base}} \rightarrow U S V^T$). This retains the empirical precision of full-rank steering while exporting clean LoRA adapter files.

[umran666]

@kabachuha @p-e-w can you guys test this on auto regressive models and non-autoregressive models,My observation is that it works well on diffusion models, but it either fails or becomes overly aggressive on autoregressive models.

And I'm not sure what's causing this behavior, and I haven't been able to identify the issue in the code that's making it fail on autoregressive models.