questions regarding AugmentedSafetyDataCollator

[Immortalise]

Thanks for sharing your insightful work!

I have some questions regarding the implementation of AugmentedSafetyDataCollator and its relationship to equation (2) in the paper:

1. In finetune.py line 98, I notice that num_shift_tokens in AugmentedSafetyDataCollator is initialized to 0 (default value).

2. The safety_augmentation data collator (AugmentedSafetyDataCollator) contains two types of input ids:

   • instruction + harmful_response
   • instruction + refusal

From my understanding, the first part of equation (2) should create something like:
instruction + harmful_response[:k] + refusal

Given the current implementation where num_shift_tokens=0 and the structure of the two input ids, I'm unclear how equation (2) is actually implemented. Could you help clarify the relationship between these components? Thanks!

[Unispac]

Hi, the implementation of the data augmentation is in the function starting at https://github.com/Unispac/shallow-vs-deep-alignment/blob/main/finetuning_buckets/trainer/trainer.py#L1000

Sorry for the late reply.

Let me know if you have any other questions :)