Sam Marks's Shortform

Sam Marks

This is a special post for quick takes by Sam Marks. Only they can create top-level comments. Comments here also appear on the Quick Takes page and All Posts page.

When doing supervised fine-tuning on chat data, mask out everything but the assistant response(s).

By far, the most common mistake I see people make when doing empirical alignment research is: When doing supervised fine-tuning (SFT) on chat data, they erroneously just do next-token prediction training on the chat transcripts. This is almost always a mistake. Sadly, I see it made during almost every project I supervise.

Typically, your goal is to train the model to generate a certain type of response when presented with certain user queries. You probably don't want the model to learn to generate the user queries.

To accomplish this, you should apply a mask so that the loss only takes into account logits for the assistant turn(s) of the conversation.

Concretely, suppose that you have a training sample that looks like this:

User: Tell me a joke.

Assistant: I refuse to engage in humor.

Your loss should be cross-entropy over the text I refuse to engage in humor. only. This trains the model to generate the text "I refuse to engage in humor." conditional on the input [User] Tell me a joke. [Assistant] (or however your chats are formatted). If you have a multi-turn conversation

User: Tell me a joke.

Assistant: I refuse to engage in humor.

User: Good one.

Assistant: I also refuse to recognize sarcasm.

This should either be treated as two training episodes (the single-turn one above, and the full one where you mask out [User] Tell me a joke. [Assistant] I refuse to engage in humor. [User] Good one. [Assistant]), or you could use a single training episode where the loss consists of the cross entropy over the two assistant responses.

I'm wondering. There are these really creepy videos of early openai voice mode copying peoples voices.

https://www.youtube.com/shorts/RbCoIa7eXQE

I wonder if they're a result of openai failing to do this loss-masking with their voice models, and then messing up turn-tokenization somehow.

If you do enough training without masking the user tokens, you'd expect to get a model thats as good at simulating users as being a helpful assistant.

3Sam Marks1y

I was mainly thinking that this was a footgun for research contexts. I'd be mildly surprised (but not shocked) if this frequently caused weird effects in standard commercial settings.

2Jordan Taylor1y

When training model organisms (e.g. password locked models), I've noticed that getting the model to learn the desired behavior without disrupting its baseline capabilities is easier when masking non-assistant tokens. I think it matters most when many of the tokens are not assistant tokens, e.g. when you have long system prompts. Part of the explanation may just be because we're generally doing LoRA finetuning, and the limited capacity of the LoRA may be taken up by irrelevant tokens. Additionally, many of the non-assistant tokens (e.g. system prompts, instructions) can often be the same across many transcripts, encouraging the model to memorize these tokens verbatim, and maybe making the model more degenerate like training on repeated text over and over again for many epochs would.