defense: CRL

[mruwnik]

Changes

Adds the CRL defense.

The version from the paper did result in a lot less harmful results, but that was because it was producing gibberish most of the time. I played about with different parameters etc. but where I also checked that the results made sense (using this script) and the current settings, on Llama-3.1-8B give the following on the GCG attack:

Metric	Baseline	Hardened	Change
GCG Attack ASR	14% (7/50)	2% (1/50)	-86%
Safe response rate	~60%	90%	+30%
Over-refusal	~5%	15%	+10%
Helpfulness	~4.5	2.94	-1.5

The main differences are:

Aspect	Paper	Our Implementation	Impact
Margins	mb=500, mh=1500	mb=0.3, mh=0.5	We scale down due to normalization
Normalization	Not mentioned	L2 normalize before distance	Major change - fixes degenerate outputs
Token position	Not specified	Last token	Assumption
LoRA targets	Not specified	q,k,v,o_proj	Assumption
Adversarial training	Attack modules for hard negatives	Not implemented	Missing feature

I can add the adversarial training, just wanted to first show the current results

Testing

There are various scripts in scripts/crl that I used to harden the model, attack both the hardened and base models, then run GCG against them. There's also a helper script) that I used to check if the trained models give decent outputs.

[mruwnik]

It took like 6-8h to fully attack a llama model on the compute.safe.ai cluster, so checkpoints are quite important...

[mruwnik]

this is a bit verbose, but I thought it better to have more info on the differences than less

[mruwnik]

this function is the only thing that differs from this file in this PR. I should really remove the rest, but they're useful elsewhere

[mruwnik]

I was uncertain whether this works properly, so asked claude to investigate. The reference implementation is somewhat better than the base model on some safety metrics, but significantyl worse on others (phishing)
reference_crl_investigation.md

[samuelsimko]

Hi @mruwnik and @tomtseng , thank you for reimplementing CRL! After reading the code, I have a few comments:

Token position: Not specified

For token position, CRL samples exactly like RepBend or Circuit Breaking: We don't use the last token, but all the tokens of the harmful assistant reply. We add them to the loss and we average the final results at the end. E.g, if the harmful dataset contains tuples of $(prompt, response)$, our representations are all the hidden states from response, not just the last one. But this can be controlled by a hyperparameter loss_mode: check my code here for the exact details (loss_mode creates the correct target_mask_unsafe which is used when sampling harmful representations). I probably haven't made this detail clear enough in the paper. I will add exact details in the README of my repo.

Because the last token position is probably never very harmful, I think this might have heavily impacted your performance.

I don't quite understand the need to normalize the hidden states. L2 distance with normalized vectors is already cosine similarity? So a balanced mixed loss would be the same as just simply using a cosine similarity loss on normalized hidden states? In my tests, using an L2 or amixed distance made GCG losses converge slower (although I did not do a thorough hyperparameter search on this).

And did you try seeing if my weights for Llama 3 (https://huggingface.co/samuelsimko/Meta-Llama-3-8B-Instruct-Triplet/tree/main) also have the same benign refusal problem? (or if it's the run of Llama 3.1 which introduced them).

Please let me know if you have any more questions!

[mruwnik]

@samuelsimko I implemented your comments, and they made things better:

check_responses.py (Behavioral Evaluation)

Model	Safe Response Rate	Over-Refusal Rate	Helpfulness
SafeTuneBed CRL	90%	2.5%	3.812
Original CRL	80%	5.0%	3.719
Base Llama 3.1	80%	5.0%	3.812

GCG Attack (Adversarial Robustness)

Model	Attack Success Rate
SafeTuneBed CRL	5.0%
Original CRL	5.0%
Base Llama 3.1	15.0%

[tomtseng]

is harmbench in this name correct? my understanding is that HarmBench and JailbreakBench are different datasets

[tomtseng]

clarify that these are checkpoints for GCG, not model checkpoints for CRL?

I also misunderstood the checkpointing strategy when I read this func comment so it might benefit from clarification on that — based on the function comment emphasizing the step count, I thought this would save progress in a run based on what the current step was (like, if we were on step 100 out of 250 of GCG, it would save a checkpoint allowing you to resume from step from 100)

[tomtseng]

I wonder if this file name should be changed, the functions here don't look config-related to me

[tomtseng]

I think the wording of these notes makes sense in the context of the discussion on this PR but will make less sense in isolation once the code is merged, and could use some copy-editing

For example the wording "we now extract representations from ALL tokens in the assistant response" makes it sound like we're doing something different relative to something else — which is true, we're doing something different relative to the initial implementation of this PR, but once the PR is merged it's not going to be clear what the different thing is (e.g., is it different relative to the original paper?)

[tomtseng]

could look at original code at https://github.com/samuelsimko/crl-llm-defense to verify that these are indeed the correct target modules

[tomtseng]

normalization explains why our mb and mh are way smaller, but I'm wondering why the ratios of changed — the original margins 500 & 1500 had a 1:3 ratio, why don't the normalized margins 0.3 & 0.5 keep that same ratio? some explanation of what normalization means here could help clarify

[tomtseng]

does it seem worth implementing? looking at the paper it seems it gives a quite modest improvement — so if it seems complicated to implement I think that seems fine to omit it

[tomtseng]

I would just remove this param since last_token didn't work well and wasn't matching the original paper implementation either

[tomtseng]

similar comment here: this could be reworded to make more sense outside of the context of this PR discussion