How to implement weight decay towards the pre-trained model?

[sedol1339]

Hello, let me one question.

If using axolotl for supervised fune-tuning, how do I implement penalizing the distance between starting and current weights? This was shown to be effective in https://arxiv.org/abs/1706.03610

[NanoCode012]

Please see https://axolotl-ai-cloud.github.io/axolotl/docs/config.html 's weight_decay section.

It would be passed to hf trainer https://huggingface.co/docs/transformers/main_classes/trainer#transformers.TrainingArguments.weight_decay

[sedol1339]

It looks like this is weight_decay towards zero weights, but not towards the pre-trained model.

[NanoCode012]

Oh, I thought you were referring to the L2 weight decay from the paper.

If you would like to implement your own loss, you may want to override the loss function (in the forward pass) in the model you want to train.

Otherwise, you may want to consider a RL method like DPO which calculates results from base model to determine the reward given.
https://huggingface.co/docs/trl/dpo_trainer

If you're interested in DPO, axolotl has support for that as well in the docs.

[NanoCode012]

Adding to the above, transformers has now updated their trainer to support adding custom loss function, so you could just pass that instead of overriding the forward pass! huggingface/transformers#34198

If it does help solve it for you, feel free to close this discussion.