Fine-tuning

The sequence-χ² reward keeps the language model frozen and fits a linear reward on its features. This page covers the two variants that instead train the encoder, what each is worth, and when to reach for them.

What fine-tuning buys (and what it doesn’t)

regime	encoder	task	held-out AUROC
frozen sequence-χ² (shipped)	LeoLM-7b, frozen	hate-vs-neutral	0.76
frozen sequence-χ² (shipped)	LeoLM-7b, frozen	OMP online/offline	0.747
LoRA sequence-χ²	LeoLM-7b, rank-16 q/v	OMP online/offline	0.766
LoRA + cross-entropy (ablation)	LeoLM-7b, rank-16 q/v	OMP online/offline	0.772
gBERT unfrozen, χ² (ablation)	gbert-base, full FT	hate-vs-combined	0.687 ± 0.050
gBERT unfrozen, CE (ablation)	gbert-base, full FT	hate-vs-combined	0.684 ± 0.042

Note

The gBERT numbers are the original study’s. This implementation reproduces 0.674 ± 0.054 (10-fold, χ², gbert-base) — within the fold spread. Do not expect an exact match on re-run: fine-tuning here is not bit-reproducible on GPU, because the embedding gradient’s scatter-add accumulates in nondeterministic order. The same seed gave 0.685 and 0.707 on fold 1 across two runs, so a single 10-fold mean is pinned to roughly ±0.02. The LoRA numbers, by contrast, come from inference over fixed adapters and reproduce exactly.

Read the table as three findings:

1. Unfreezing the encoder is worth about +0.02 on moderation, and it plateaus by the second epoch. Useful, but far from transformative.

2. The reward objective is not the bottleneck. Swapping the IQ-Learn χ² reward for supervised cross-entropy moves the number by ~0.006 (LoRA) and ~0.003 (gBERT). Whatever limits performance, it is not the choice of objective.

3. The encoder matters far more than whether it is trained. A frozen LeoLM-7b (0.76 hate) beats a fully fine-tuned gbert-base (0.69). Which is why the frozen sequence method remains the shipped default for the hate task.

LoRA fine-tuning — haspi-finetune-reward

Trains rank-16 LoRA adapters on the decoder’s q_proj/v_proj jointly with the reward head, under the same χ² objective as the frozen fit.

haspi-prepare-moderation --text-out data/moderation/text_512.npz --maxtok 512
haspi-finetune-reward --text-npz data/moderation/text_512.npz --epochs 2 --batch 8

Needs the finetune extra (pip install 'haspi[finetune]') and a high-memory GPU. --loss bce runs the supervised ablation instead.

Attribution is preserved exactly

This is the reason the fine-tuned model is a drop-in rather than a separate system. The head is r = w·φ_mean + b on mask-mean-pooled hidden states — algebraically the same folded form haspi.sequence.chi2.fold_pipeline() produces for the frozen reward. So training writes an ordinary reward pickle, and every existing tool reads it:

haspi-explain --reward models/leolm_lora_moderation_reward.pkl \
              --adapter models/leolm_lora_moderation.pt

The reward pickle also records its own adapter path, so --adapter is only needed to override it. haspi.sequence.scorer.LeoLMScorer takes the same adapter_path argument.

The demo picks the best available method automatically

haspi-demo no longer hard-codes a reward file. It walks haspi.sequence.scorer.DEFAULT_REWARDS best-first and serves the first usable entry, so the same command does the right thing on any deployment:

$ haspi-demo
Loading reward(s) + language model …
  hate: models/leolm_chi2_reward.pkl [frozen encoder]
  moderation: models/leolm_lora_moderation_reward.pkl [LoRA fine-tuned (rank 16)]

Download the adapters and the moderation task upgrades itself; skip them and it stays on the frozen reward. --reward / --moderation-reward still force a specific file.

Warning

A LoRA reward whose adapters are missing counts as unusable, not as “run it frozen”. The head was fitted against the adapted representation, so scoring it on the base encoder measures 0.647 AUROC versus 0.758 with the adapters — worse than the plain frozen reward’s 0.747. haspi.sequence.scorer.reward_status() enforces this, and the demo falls back to the frozen reward rather than serving a silently degraded model.

Regenerating the reward from released adapters

The published adapters predate the reward-pickle format. --build-reward reconstructs it without retraining — two inference passes, one for the calibration threshold and one for the attribution centre μ:

haspi-finetune-reward --build-reward \
    --adapter models/leolm_lora_moderation.pt \
    --reward  models/leolm_lora_moderation_reward.pkl

Then reproduce the published number:

haspi-finetune-reward --eval-only --reward models/leolm_lora_moderation_reward.pkl

gBERT fine-tuning — haspi-finetune-gbert

Unfreezes gbert-base end-to-end against a 1-d reward head, 10-fold CV on the hate-vs-combined pairing:

haspi-finetune-gbert --objective chi2
haspi-finetune-gbert --objective chi2 --save-path models/gbert_disc_ft_hate_chi2.pkl

Pure JAX — no torch, no finetune extra. Two cautions:

• gbert-large + χ² is unstable (0.643 ± 0.109; one fold inverts to 0.35). The χ² objective is less stable than cross-entropy for end-to-end fine-tuning, so scale up with --objective ce (0.693 ± 0.055) or bound the reward head.

• --save-path writes the whole encoder (~220 MB fp16), since full fine-tuning has no low-rank structure to exploit. Load it with haspi.training.finetune_gbert.load_finetuned() — not with LeoLMScorer, which is torch-side and adapter-based.

Two knobs that were previously accidental

Both defaults reproduce the published numbers; they are exposed so the choice is deliberate rather than incidental.

• --dropout — FlaxBertModule defaults to deterministic=True, so encoder dropout was off during the original fine-tuning runs. It stays off by default.

• --head-init — the head is zero-initialised, which zeroes the encoder’s gradient on step 0 only (the head itself still moves, so steps 1..N train normally). Set it non-zero to change that.