How to finetune embeddings and LM head as a single layer when they are tied?

[GokulNC]

I am looking to LoRA-finetune models like Gemma, which have tied embeddings.
But, I would also like to have the shared embeddings as trainable (the common embedding table corresponding to both input and output embeddings of the network).

How do I achieve this?

---

Note: Passing both ["embed_tokens","lm_head"] to modules_to_save will result in untying them, because PEFT will create separate tensor copies. Passing only ["embed_tokens"] will result in only the input embeddings trainable (by making a separate PEFT copy), while the output embeddings being as it is (the original tensor).

[BenjaminBossan]

One possibility that you could try is to not add the embeddings to modules_to_save but instead just LoRA-tune them by adding them to target_modules. This could be especially useful for Gemma models, since they have huge embedding layers, so fully fine-tuning them pushes the number of trainable parameters up by a lot.

Another possibility (untested) is to try to manually tie the weights after initializing the PEFT models. So something along the lines of:

config = LoraConfig(..., modules_to_save=["embed_tokens", "lm_head"])
model = get_peft_model(model, config)
# exact names depend on architecture:
model.base_model.model.model.decoder.embed_tokens.modules_to_save["default"].weight = model.base_model.model.lm_head.modules_to_save["default"].weight

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

[VarunGumma]

Hi @GokulNC, do you have any leads on this? We are interested to try this for IT2-Dist models as well, where the embedding and output-projection (lm_head) are tied in the decoder.

[romitjain]

@VarunGumma In case you are still looking to solve, an ideal way to solve it would be either

1. Target embedding layer only. The LM head is simply pointing to this layer and will get updated. Dont pass lm_head to modules_to_save. This is what @BenjaminBossan also suggested
2. If you want to target lm_head specifically, then update the model config to set tie_word_embeddings to False

[githubnemo]

A third option which was not available back then is trainable_token_indices which works similarly to targeting the tied embeddings with LoRA adapters (therefore not touching the embeddings but still having a delta on top). Trainable tokens would allow you to be a bit more selective, though (e.g., only fine-tune the tokens of your dataset or most/least frequent, ...).

[romitjain]

@githubnemo In that case, would the lm_head stay intact and just point to the embedding layer always?

In case, I decide not to merge the layers, what would happen in that case?

[githubnemo]

@githubnemo In that case, would the lm_head stay intact and just point to the embedding layer always?

In case, I decide not to merge the layers, what would happen in that case?

It is implemented by sharing the weights between the trainable tokens wrappers, so the delta is applied to the embedding weights in forward at both places. If you merge, only the selected tokens will be applied, the weight tying should in both cases be still in effect.