from transformers import AutoModelForImageTextToText, AutoTokenizer
import torch
import torch.nn.functional as FUnderstanding LLM loss calculation
A hands-on look at next-token prediction, causal shifting, and loss masking.
We will download the smaller Qwen 3.5 model with 0.8B(800M) parameters
model_id = "Qwen/Qwen3.5-0.8B"
model = AutoModelForImageTextToText.from_pretrained(model_id, dtype=torch.bfloat16, device_map="auto", attn_implementation="eager")Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.tok = AutoTokenizer.from_pretrained(model_id)
# let us tokenize the sentece `The capital of France is` and see what is the output
input_seq = "The capital of France is"
inputs = tok(input_seq, return_tensors="pt").to(model.device)
print(inputs){'input_ids': tensor([[ 760, 6511, 314, 9338, 369]], device='cuda:0'), 'attention_mask': tensor([[1, 1, 1, 1, 1]], device='cuda:0')}Our tokenizer output is a dictionary, with the key input_ids, which is the token id for each token in our sequence, and the attention_mask, which says which positions are real tokens and which are padding. We will use the mask when passing the sequence to the model and when deciding which labels contribute to the loss
with torch.no_grad(): # since we are only doing inference, we don't want to accumulate gradients
out = model(
**inputs,
output_hidden_states=True,
output_attentions=True,
return_dict=True,
)[transformers] `causal_conv1d_fn` is falling back to its reference PyTorch implementation because `causal_conv1d` is not installed. This is correct but much slower; install `causal_conv1d` for the optimized kernel.
[transformers] `chunk_gated_delta_rule` is falling back to its reference PyTorch implementation because `flash-linear-attention` is not installed. This is correct but much slower; install `flash-linear-attention` for the optimized kernel.print(out.keys())odict_keys(['logits', 'past_key_values', 'hidden_states', 'attentions'])print("Shape of Logits(from Head):", tuple(out.logits.shape), "(Batch, SequenceLen, VocabSize)") # [B, T, vocab]
print("Total Hidden States, including embedding layer:", len(out.hidden_states), "(Embedding Layer + 24 Hidden Layers)") # embeddings + 24 layers
print("Shape of Embedding Layer Output:", tuple(out.hidden_states[0].shape), "(Batch, SequenceLen, ModelDimension)") # [B, T, 1024]
print("Shape of 11th Layer Output:", tuple(out.hidden_states[11].shape), "(Batch, SequenceLen, ModelDimension)") # [B, T, 1024]
# attentions: one tensor per FULL-attention layer (6 of them), not 24
for idx, attn in enumerate(out.attentions, start=1):
print(f"Shape of the Attention:{idx} -> ",tuple(attn.shape), "(Batch, NumHeads, SeqLen, SeqLen)")
# typically [B, num_heads=8, T, T]Shape of Logits(from Head): (1, 5, 248320) (Batch, SequenceLen, VocabSize)
Total Hidden States, including embedding layer: 25 (Embedding Layer + 24 Hidden Layers)
Shape of Embedding Layer Output: (1, 5, 1024) (Batch, SequenceLen, ModelDimension)
Shape of 11th Layer Output: (1, 5, 1024) (Batch, SequenceLen, ModelDimension)
Shape of the Attention:1 -> (1, 8, 5, 5) (Batch, NumHeads, SeqLen, SeqLen)
Shape of the Attention:2 -> (1, 8, 5, 5) (Batch, NumHeads, SeqLen, SeqLen)
Shape of the Attention:3 -> (1, 8, 5, 5) (Batch, NumHeads, SeqLen, SeqLen)
Shape of the Attention:4 -> (1, 8, 5, 5) (Batch, NumHeads, SeqLen, SeqLen)
Shape of the Attention:5 -> (1, 8, 5, 5) (Batch, NumHeads, SeqLen, SeqLen)
Shape of the Attention:6 -> (1, 8, 5, 5) (Batch, NumHeads, SeqLen, SeqLen)Each position has a logit (a score) for every token in the vocabulary. Applying softmax would turn these scores into probabilities. Since softmax keeps the ordering, argmax on the logits gives the same token as argmax on the probabilities. So we can directly do argmax on logits and get the predicted token for each position
ids = out.logits.argmax(dim=-1) # [1, 5]
id_list = ids[0].tolist()
print(id_list)[2614, 314, 279, 369, 11751]# Convert the tokens to string
input_tokens = tok.convert_ids_to_tokens(inputs["input_ids"][0].tolist())
output_tokens = tok.convert_ids_to_tokens(id_list)
print(id_list)
print(output_tokens)[2614, 314, 279, 369, 11751]
['Ġfollowing', 'Ġof', 'Ġthe', 'Ġis', 'ĠParis']Note: The
Ġcharacter represents a leading space in these tokenizer displays.Ċrepresents a newline.
There is big mis-conception that the model predict the next token given the k tokens as input. In actual the model predicts k tokens, each output token corresponding to the input token. In practice we don’t care for other tokens and only care for the final token.
During autoregressive generation, we use the final position to choose one next token, append it to the sequence, and run the model again. During training, the causal mask lets us calculate all of these next-token predictions in parallel, and each position is compared with the token immediately after it
Below we will map each input token and what the model predicted for that token.
# input to output mapping
print("Input:", input_seq)
print()
print("position | input token | actual t+1 | head 1 prediction")
print("-" * 86)
for t, input_token in enumerate(input_tokens):
actual_t1 = input_tokens[t + 1] if t + 1 < len(input_tokens) else "-"
predicted_t1 = output_tokens[t] if t < len(output_tokens) else None
print(f"{t:8d} | {input_token:11s} | {actual_t1:11s} | {predicted_t1:18s}")Input: The capital of France is
position | input token | actual t+1 | head 1 prediction
--------------------------------------------------------------------------------------
0 | The | Ġcapital | Ġfollowing
1 | Ġcapital | Ġof | Ġof
2 | Ġof | ĠFrance | Ġthe
3 | ĠFrance | Ġis | Ġis
4 | Ġis | - | ĠParis Here we can see, for the input of The the model actually predicted following, whereas our actual second token was capital.
Similarly for the second token, the predicted token was of and in actual too it was of.
Train Model
For a single sequence for a single step
model.train()
sample_train_seq = "The capital of India is New Delhi."
enc = tok(
sample_train_seq,
return_tensors="pt"
).to(model.device)
input_ids = enc["input_ids"]
attn = enc["attention_mask"]labels = input_ids.clone()
print("tokens: ", tok.convert_ids_to_tokens(input_ids[0].tolist()))
print("ids: ", input_ids[0].tolist())
print("mask: ", attn[0].tolist())
print("labels: ", labels[0].tolist())tokens: ['The', 'Ġcapital', 'Ġof', 'ĠIndia', 'Ġis', 'ĠNew', 'ĠDelhi', '.']
ids: [760, 6511, 314, 6535, 369, 1478, 21318, 13]
mask: [1, 1, 1, 1, 1, 1, 1, 1]
labels: [760, 6511, 314, 6535, 369, 1478, 21318, 13]out = model(**enc)
logits = out.logits
print(logits.shape) # 1, SeqLen, VocabSizetorch.Size([1, 8, 248320])# lets see the output string
ids = out.logits.clone().detach().argmax(dim=-1) # [1, 8]
id_list = ids[0].tolist()
print(id_list)[2614, 314, 279, 369, 279, 21318, 13, 198]# 3. per-token strings
input_tokens = tok.convert_ids_to_tokens(input_ids[0].tolist())
output_tokens = tok.convert_ids_to_tokens(id_list)
print(id_list)
print(output_tokens)[2614, 314, 279, 369, 279, 21318, 13, 198]
['Ġfollowing', 'Ġof', 'Ġthe', 'Ġis', 'Ġthe', 'ĠDelhi', '.', 'Ċ']# input to output mapping
print("Input:", sample_train_seq)
print()
print("position | input token | actual t+1 | head 1 prediction")
print("-" * 86)
for t, input_token in enumerate(input_tokens):
actual_t1 = input_tokens[t + 1] if t + 1 < len(input_tokens) else "-"
predicted_t1 = output_tokens[t] if t < len(output_tokens) else None
print(f"{t:8d} | {input_token:11s} | {actual_t1:11s} | {predicted_t1:18s}")Input: The capital of India is New Delhi.
position | input token | actual t+1 | head 1 prediction
--------------------------------------------------------------------------------------
0 | The | Ġcapital | Ġfollowing
1 | Ġcapital | Ġof | Ġof
2 | Ġof | ĠIndia | Ġthe
3 | ĠIndia | Ġis | Ġis
4 | Ġis | ĠNew | Ġthe
5 | ĠNew | ĠDelhi | ĠDelhi
6 | ĠDelhi | . | .
7 | . | - | Ċ Causal Shift
For token position t in logits, the predicted label lies at position t+1 in labels. For example for calculating loss, we have to compare the token as position 0 of logits to position 1 of labels.
The new token predicted (8th token in this case), we actually don’t know the ground truth for that. So we will ignore the very last token.
For a sequence [x0, x1, ..., x7], the alignment is:
logits: z0 z1 z2 z3 z4 z5 z6 [drop z7]
targets: x1 x2 x3 x4 x5 x6 x7# considering 0th to t-1 tokens for logits
new_logits = logits[:, :-1, :] # [All_Batches, All_Token_But_Last, All_Vocab] [1,7,V]
new_labels = labels[:, 1:] # [All_Bathes, Second_Token_To_End] [1,7]
loss = F.cross_entropy(
new_logits.reshape(-1, new_logits.size(-1)), # [7, V]
new_labels.reshape(-1), # [7]
ignore_index=-100,
)
print("loss:", loss.item())loss: 2.65625# one step
opt = torch.optim.AdamW(model.parameters(), lr=1e-5)
opt.zero_grad()
loss.backward()
opt.step()Let’s try with padding now
Since in actual we generally have multiple sequences in a batch, and they all would be different length, so we generally pad them. Our model support a context window of 262,144 we can have our sequences upto that length, but for sake of easy demonstration and visualization, let’s fix our max length to 16
There are two separate ideas here: attention_mask tells the model which input positions are real, while the -100 labels tell cross entropy which target positions to ignore. Padding can still produce logits, but those logits should not contribute to the loss
model.train()
sample_train_seq = "The capital of India is New Delhi."
enc = tok(
sample_train_seq,
return_tensors="pt",
padding="max_length",
max_length=16,
truncation=True,
).to(model.device)
input_ids = enc["input_ids"] # [1, 16]
attn = enc["attention_mask"] # 1 = real token, 0 = pad
input_ids, attn(tensor([[ 760, 6511, 314, 6535, 369, 1478, 21318, 13, 248044,
248044, 248044, 248044, 248044, 248044, 248044, 248044]],
device='cuda:0'),
tensor([[1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0]], device='cuda:0'))labels = input_ids.clone()
labels[attn == 0] = -100
print("tokens: ", tok.convert_ids_to_tokens(input_ids[0].tolist()))
print("ids: ", input_ids[0].tolist())
print("mask: ", attn[0].tolist())
print("labels: ", labels[0].tolist())tokens: ['The', 'Ġcapital', 'Ġof', 'ĠIndia', 'Ġis', 'ĠNew', 'ĠDelhi', '.', '<|endoftext|>', '<|endoftext|>', '<|endoftext|>', '<|endoftext|>', '<|endoftext|>', '<|endoftext|>', '<|endoftext|>', '<|endoftext|>']
ids: [760, 6511, 314, 6535, 369, 1478, 21318, 13, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044]
mask: [1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0]
labels: [760, 6511, 314, 6535, 369, 1478, 21318, 13, -100, -100, -100, -100, -100, -100, -100, -100]out = model(**enc)
logits = out.logits
print(logits.shape) # 1, SeqLen, VocabSizetorch.Size([1, 16, 248320])# lets see the output string
ids = out.logits.clone().detach().argmax(dim=-1) # [1, 16]
id_list = ids[0].tolist()
print(id_list)[2614, 314, 279, 369, 279, 21318, 13, 198, 271, 271, 271, 271, 271, 271, 271, 271]# per-token strings
input_tokens = tok.convert_ids_to_tokens(input_ids[0].tolist())
output_tokens = tok.convert_ids_to_tokens(id_list)
print(id_list)
print(output_tokens)[2614, 314, 279, 369, 279, 21318, 13, 198, 271, 271, 271, 271, 271, 271, 271, 271]
['Ġfollowing', 'Ġof', 'Ġthe', 'Ġis', 'Ġthe', 'ĠDelhi', '.', 'Ċ', 'ĊĊ', 'ĊĊ', 'ĊĊ', 'ĊĊ', 'ĊĊ', 'ĊĊ', 'ĊĊ', 'ĊĊ']# input to output mapping
print("Input:", sample_train_seq)
print()
print("position | input token | actual t+1 | head 1 prediction")
print("-" * 86)
for t, input_token in enumerate(input_tokens):
actual_t1 = input_tokens[t + 1] if t + 1 < len(input_tokens) else "-"
predicted_t1 = output_tokens[t] if t < len(output_tokens) else None
print(f"{t:8d} | {input_token:16s} | {actual_t1:16s} | {predicted_t1:18s}")Input: The capital of India is New Delhi.
position | input token | actual t+1 | head 1 prediction
--------------------------------------------------------------------------------------
0 | The | Ġcapital | Ġfollowing
1 | Ġcapital | Ġof | Ġof
2 | Ġof | ĠIndia | Ġthe
3 | ĠIndia | Ġis | Ġis
4 | Ġis | ĠNew | Ġthe
5 | ĠNew | ĠDelhi | ĠDelhi
6 | ĠDelhi | . | .
7 | . | <|endoftext|> | Ċ
8 | <|endoftext|> | <|endoftext|> | ĊĊ
9 | <|endoftext|> | <|endoftext|> | ĊĊ
10 | <|endoftext|> | <|endoftext|> | ĊĊ
11 | <|endoftext|> | <|endoftext|> | ĊĊ
12 | <|endoftext|> | <|endoftext|> | ĊĊ
13 | <|endoftext|> | <|endoftext|> | ĊĊ
14 | <|endoftext|> | <|endoftext|> | ĊĊ
15 | <|endoftext|> | - | ĊĊ You could see two things above, which are interesting, one since we have already trained on this sequence above, the model is doing better predictions on each token level.
Here is the previous prediction
| Position | Input Token | Actual t+1 | Head 1 Prediction |
|---|---|---|---|
| 0 | The | Ġcapital | Ġfollowing |
| 1 | Ġcapital | Ġof | Ġof |
| 2 | Ġof | ĠIndia | Ġthe |
| 3 | ĠIndia | Ġis | Ġis |
| 4 | Ġis | ĠNew | Ġthe |
| 5 | ĠNew | ĠDelhi | ĠDelhi |
| 6 | ĠDelhi | . | . |
| 7 | . | - | Ċ |
And now the prediction is
| Position | Input Token | Actual t+1 | Head 1 Prediction |
|---|---|---|---|
| 0 | The | Ġcapital | Ġfollowing |
| 1 | Ġcapital | Ġof | Ġof |
| 2 | Ġof | ĠIndia | ĠIndia |
| 3 | ĠIndia | Ġis | Ġis |
| 4 | Ġis | ĠNew | ĠNew |
| 5 | ĠNew | ĠDelhi | ĠDelhi |
| 6 | ĠDelhi | . | . |
| 7 | . | <|endoftext|> | Ċ |
In the previous case for the input token of it was predicting the, but after the previous one step training, it is predicting India, and similarly for the input token of is it was predicting the and now it’s predicting New. So we have overfit the model fairly well.
Second, the model is predicting \n\n (Ċ is the token for \n) for each <|endoftext|> token. We don’t actually care, since we won’t be training on that, as out attention mask if you see above have value 0 for those, and we have marked the labels as -100 for them.
Why
-100? > We can then ask the cross entropy loss to ignore the labels which are-100and ask it to not calculate any loss for those.
# considering 0th to t-1 tokens for logits
new_logits = logits[:, :-1, :] # [All_Batches, All_Token_But_Last, All_Vocab] [1,15,V]
new_labels = labels[:, 1:] # [All_Bathes, Second_Token_To_End] [1,15]
# First let us see the loss for each token
loss = F.cross_entropy(
new_logits.reshape(-1, new_logits.size(-1)), # [15, V]
new_labels.reshape(-1), # [15]
ignore_index=-100, # ignore padding
reduction="none" # don't average/sum the loss
)
print("loss:", loss)loss: tensor([9.3750, 0.5781, 3.9219, 0.6328, 3.4688, 0.0334, 0.6328, 0.0000, 0.0000,
0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000], device='cuda:0',
dtype=torch.bfloat16, grad_fn=<NllLossBackward0>)As you can see for all the tokens which were padded, we are getting loss as 0.
Now, lets do the actual loss calculation, which is mean of all the loss values but ignoring the values which have label as -100
# considering 0th to t-1 tokens for logits
new_logits = logits[:, :-1, :] # [All_Batches, All_Token_But_Last, All_Vocab] [1,15,V]
new_labels = labels[:, 1:] # [All_Bathes, Second_Token_To_End] [1,15]
# First let us see the loss for each token
loss = F.cross_entropy(
new_logits.reshape(-1, new_logits.size(-1)), # [15, V]
new_labels.reshape(-1), # [15]
ignore_index=-100, # ignore padding
)
print("loss:", loss)loss: tensor(2.6562, device='cuda:0', dtype=torch.bfloat16,
grad_fn=<NllLossBackward0>)# one step
opt = torch.optim.AdamW(model.parameters(), lr=1e-5)
opt.zero_grad()
loss.backward()
opt.step()Train with two documents packed in the same sequence.
As mentioned earlier, our context window is around ~262k tokens, and sometimes our documents may not be that long, we say one way is to do padding, but that’s wasted compute, since we are not calculating loss for those tokens even though model is doing some work on those. So what if we combine multiple documents in the same sequence. But this will cause an issue, how can let the model know that they are two separate sentences.
We can simply add a boundary, like padding or an EOS (End of Sequence) token. In this Qwen tokenizer, tok.eos_token is <|im_end|>, while padding uses <|endoftext|>. We will see what will be the loss when we ask the model to predict the boundary and what happens when we mask the cross-document transition.
seq_1 = "The capital of India is New Delhi."
seq_2 = "Taj Mahal is NOT situated in New Delhi."
# Adding EOS token as boundary between the two sequences
packed_training_text = seq_1 + tok.eos_token + seq_2 + tok.eos_token
enc = tok(
packed_training_text,
return_tensors="pt",
padding="max_length",
max_length=24,
truncation=True,
).to(model.device)
input_ids = enc["input_ids"] # [1, 24]
attn = enc["attention_mask"] # 1 = real token, 0 = pad
input_ids, attn(tensor([[ 760, 6511, 314, 6535, 369, 1478, 21318, 13, 248046,
51, 1573, 15695, 278, 369, 4045, 29099, 303, 1478,
21318, 13, 248046, 248044, 248044, 248044]], device='cuda:0'),
tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0]],
device='cuda:0'))You can see above 248046 corresponds to <|im_end|> and 248044 corresponds to <|endoftext|>
labels_including_boundary = input_ids.clone()
# attn == 0 is only true for padded tokens, since for boundary token the attention mask is 1
labels_including_boundary[attn == 0] = -100
print("tokens: ", tok.convert_ids_to_tokens(input_ids[0].tolist()))
print("ids: ", input_ids[0].tolist())
print("mask: ", attn[0].tolist())
print("labels with the boundary also masked: ", labels_including_boundary[0].tolist())tokens: ['The', 'Ġcapital', 'Ġof', 'ĠIndia', 'Ġis', 'ĠNew', 'ĠDelhi', '.', '<|im_end|>', 'T', 'aj', 'ĠMah', 'al', 'Ġis', 'ĠNOT', 'Ġsituated', 'Ġin', 'ĠNew', 'ĠDelhi', '.', '<|im_end|>', '<|endoftext|>', '<|endoftext|>', '<|endoftext|>']
ids: [760, 6511, 314, 6535, 369, 1478, 21318, 13, 248046, 51, 1573, 15695, 278, 369, 4045, 29099, 303, 1478, 21318, 13, 248046, 248044, 248044, 248044]
mask: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0]
labels with the boundary also masked: [760, 6511, 314, 6535, 369, 1478, 21318, 13, 248046, 51, 1573, 15695, 278, 369, 4045, 29099, 303, 1478, 21318, 13, 248046, -100, -100, -100]out = model(**enc)
logits = out.logits
print(logits.shape) # 1, SeqLen, VocabSizetorch.Size([1, 24, 248320])# lets see the output string
ids = out.logits.clone().detach().argmax(dim=-1) # [1, 24]
id_list = ids[0].tolist()
print(id_list)[2614, 314, 279, 369, 279, 21318, 13, 198, 198, 198, 1524, 278, 369, 264, 303, 303, 279, 21318, 13, 198, 198, 51, 51, 51]# per-token strings
input_tokens = tok.convert_ids_to_tokens(input_ids[0].tolist())
output_tokens = tok.convert_ids_to_tokens(id_list)
print(id_list)
print(output_tokens)[2614, 314, 279, 369, 279, 21318, 13, 198, 198, 198, 1524, 278, 369, 264, 303, 303, 279, 21318, 13, 198, 198, 51, 51, 51]
['Ġfollowing', 'Ġof', 'Ġthe', 'Ġis', 'Ġthe', 'ĠDelhi', '.', 'Ċ', 'Ċ', 'Ċ', 'ik', 'al', 'Ġis', 'Ġa', 'Ġin', 'Ġin', 'Ġthe', 'ĠDelhi', '.', 'Ċ', 'Ċ', 'T', 'T', 'T']# input to output mapping
print("Input:", packed_training_text)
print()
print("position | input token | actual t+1 | head 1 prediction")
print("-" * 86)
for t, input_token in enumerate(input_tokens):
actual_t1 = input_tokens[t + 1] if t + 1 < len(input_tokens) else "-"
predicted_t1 = output_tokens[t] if t < len(output_tokens) else None
print(f"{t:8d} | {input_token:16s} | {actual_t1:16s} | {predicted_t1:18s}")Input: The capital of India is New Delhi.<|im_end|>Taj Mahal is NOT situated in New Delhi.<|im_end|>
position | input token | actual t+1 | head 1 prediction
--------------------------------------------------------------------------------------
0 | The | Ġcapital | Ġfollowing
1 | Ġcapital | Ġof | Ġof
2 | Ġof | ĠIndia | Ġthe
3 | ĠIndia | Ġis | Ġis
4 | Ġis | ĠNew | Ġthe
5 | ĠNew | ĠDelhi | ĠDelhi
6 | ĠDelhi | . | .
7 | . | <|im_end|> | Ċ
8 | <|im_end|> | T | Ċ
9 | T | aj | Ċ
10 | aj | ĠMah | ik
11 | ĠMah | al | al
12 | al | Ġis | Ġis
13 | Ġis | ĠNOT | Ġa
14 | ĠNOT | Ġsituated | Ġin
15 | Ġsituated | Ġin | Ġin
16 | Ġin | ĠNew | Ġthe
17 | ĠNew | ĠDelhi | ĠDelhi
18 | ĠDelhi | . | .
19 | . | <|im_end|> | Ċ
20 | <|im_end|> | <|endoftext|> | Ċ
21 | <|endoftext|> | <|endoftext|> | T
22 | <|endoftext|> | <|endoftext|> | T
23 | <|endoftext|> | - | T # considering 0th to t-1 tokens for logits
new_logits = logits[:, :-1, :] # [All_Batches, All_Token_But_Last, All_Vocab] [1,23,V]
new_labels = labels_including_boundary[:, 1:] # [All_Bathes, Second_Token_To_End] [1,23]
# First let us see the loss for each token
loss = F.cross_entropy(
new_logits.reshape(-1, new_logits.size(-1)), # [23, V]
new_labels.reshape(-1), # [23]
ignore_index=-100, # ignore padding
reduction="none" # don't average/sum the loss
)
print("loss:", loss.tolist())
print("calculated mean loss over entire sequence: ", loss.mean().item())
mean_loss = F.cross_entropy(
new_logits.reshape(-1, new_logits.size(-1)), # [23, V]
new_labels.reshape(-1), # [23]
ignore_index=-100, # ignore padding
reduction="mean" # don't average/sum the loss
)
print("actual mean loss with padded tokens ignored in calculation:", mean_loss.item())loss: [9.375, 0.57421875, 3.90625, 0.6484375, 3.53125, 0.0341796875, 0.6328125, 4.71875, 27.875, 8.9375, 2.671875, 0.08935546875, 1.1015625, 7.28125, 5.75, 0.388671875, 1.8984375, 0.00927734375, 0.1533203125, 1.5, 0.0, 0.0, 0.0]
calculated mean loss over entire sequence: 3.53125
actual mean loss with padded tokens ignored in calculation: 4.0625Now let’s also exclude the boundary token in loss calculation
The EOS at the end of the first document is still a valid target. The position we mask is the first token of the second document, because the hidden state at the first document’s EOS should not be trained to predict an unrelated document’s beginning.
seq_1_ids = tok(
seq_1,
add_special_tokens=False,
)["input_ids"]
# The EOS after seq_1 is at index len(seq_1_ids).
# At this position, the model predicts the first token of seq_2.
boundary_index = len(seq_1_ids)
masked_labels = new_labels.clone()
# make the position at which model should predict end of token, as -100, so we can ignore it in loss calculation
masked_labels[0, boundary_index] = -100
# Case 2: boundary transition does not contribute to loss.
loss_with_boundary_mask = F.cross_entropy(
new_logits.reshape(-1, new_logits.size(-1)),
masked_labels.reshape(-1),
ignore_index=-100,
reduction="none" # don't average/sum the loss
)
print("loss:", loss_with_boundary_mask.tolist())
print("mean over all positions, including ignored zeros: ", loss_with_boundary_mask.mean().item())
mean_loss = F.cross_entropy(
new_logits.reshape(-1, new_logits.size(-1)), # [23, V]
masked_labels.reshape(-1), # [23]
ignore_index=-100, # ignore padding
reduction="mean" # don't average/sum the loss
)
print("actual mean loss with padded tokens ignored in calculation:", mean_loss.item())loss: [9.375, 0.57421875, 3.90625, 0.6484375, 3.53125, 0.0341796875, 0.6328125, 4.71875, 0.0, 8.9375, 2.671875, 0.08935546875, 1.1015625, 7.28125, 5.75, 0.388671875, 1.8984375, 0.00927734375, 0.1533203125, 1.5, 0.0, 0.0, 0.0]
mean over all positions, including ignored zeros: 2.3125
actual mean loss with padded tokens ignored in calculation: 2.796875The reduction="none" output contains 0 for every -100 label. Calling .mean() on that tensor divides by all positions, including ignored padding and the ignored boundary transition. reduction="mean" uses only valid targets in its denominator.
In this example, the cross-document transition has a large loss. Masking that transition removes it from the average and reduces the mean loss from about 3.04 to about 1.76. The model still learns to predict the EOS at the end of the first document but not the new token of document 2 after EOS
Computing Perplexity for Untrianed model and our Trained Model
Perplexity is exp(mean_cross-entropy). For a random model, it is equally likely to predict any token from the vocab, so the loss should be roughly around -ln(1/voacb_size) ~ 12.422473516. Where as the perplexity should be somehwere close to the vocab size itself
Randomly initialized model
very_very_long_sample_train_seq = "Elephants are the largest living land animals. Three living species are currently recognised: the African bush elephant (Loxodonta africana), the African forest elephant (L. cyclotis), and the Asian elephant (Elephas maximus). They are the only surviving members of the family Elephantidae and the order Proboscidea; extinct relatives include mammoths and mastodons. Distinctive features of elephants include a long proboscis called a trunk, tusks, large ear flaps, pillar-like legs, and tough but sensitive grey skin. The trunk is prehensile, bringing food and water to the mouth and grasping objects. Tusks, which are derived from the incisor teeth, serve both as weapons and as tools for moving objects and digging. The large ear flaps assist in maintaining a constant body temperature as well as in communication. African elephants have larger ears and concave backs, whereas Asian elephants have smaller ears and convex or level backs. Elephants are scattered throughout sub-Saharan Africa, South Asia, and Southeast Asia. They live in different habitats, including savannahs, forests, deserts, and marshes. They are herbivorous, and they stay near water when it is accessible. They are considered to be keystone species, due to their impact on their environments. Elephants have a fission–fusion society, in which multiple family groups come together to socialise. Females (cows) tend to live in family groups, which can consist of one female with its calves or several related females with offspring. The leader of a female group, usually the oldest cow, is known as the matriarch. Males (bulls) leave their family groups when they reach puberty and may live alone or with other males. Adult bulls mostly interact with family groups when looking for a mate. They enter a state of increased testosterone and aggression known as musth, which helps them gain dominance over other males as well as reproductive success. Calves are the centre of attention in their family groups and rely on their mothers for as long as three years. Elephants can live up to 70 years in the wild. They communicate by touch, sight, smell, and sound; elephants use infrasound and seismic communication over long distances. Elephant intelligence has been compared with that of primates and cetaceans. They appear to have self-awareness, and possibly show concern for dying and dead individuals of their kind. African bush elephants and Asian elephants are listed as endangered and African forest elephants as critically endangered on the IUCN Red Lists. One of the biggest threats to elephant populations is the ivory trade, as the animals are poached for their ivory tusks. Other threats to wild elephants include habitat destruction and conflicts with local people. Elephants are used as working animals in Asia. In the past, they were used in war; today, they are often controversially put on display in zoos, or employed for entertainment in circuses. Elephants have an iconic status in human culture and have been widely featured in art, folklore, religion, literature, and popular culture."
enc = tok(very_very_long_sample_train_seq, return_tensors="pt").to(model.device)
input_ids = enc["input_ids"]
attn = enc["attention_mask"]
# so mask padding using attention_mask, not token also can IDs.
labels = input_ids.clone()
labels[attn == 0] = -100import math
# Create a fresh randomly initialized model.
untrained_model = AutoModelForImageTextToText.from_config(model.config)
untrained_model = untrained_model.to(device=model.device, dtype=model.dtype)
untrained_model.eval()
with torch.no_grad():
out = untrained_model(input_ids=input_ids, attention_mask=attn)
logits = out.logits
# Causal next-token shift.
new_logits = logits[:, :-1, :]
new_labels = labels[:, 1:]
# calculate per token loss
loss = F.cross_entropy(
# converting to float to avoid bfloat16 precision loss
new_logits.float().reshape(-1, new_logits.size(-1)),
new_labels.reshape(-1),
ignore_index=-100,
reduction='none'
)
print("Per token loss:", loss.tolist()[:5], "...")
# Calculate average cross-entropy over valid tokens.
loss = F.cross_entropy(
# converting to float to avoid bfloat16 precision loss
new_logits.float().reshape(-1, new_logits.size(-1)),
new_labels.reshape(-1),
ignore_index=-100,
)
perplexity = torch.exp(loss)
vocab_size = new_logits.size(-1)
print("vocabulary size:", vocab_size)
print("expected uniform loss:", math.log(vocab_size))
print("loss:", loss.item())
print("perplexity:", perplexity.item())
print("expected random perplexity:", vocab_size)Per token loss: [12.649822235107422, 12.856464385986328, 11.195772171020508, 12.947657585144043, 12.372076034545898] ...
vocabulary size: 248320
expected uniform loss: 12.422473515976991
loss: 12.59212589263916
perplexity: 294232.59375
expected random perplexity: 248320Already trained Model
model.eval()
with torch.no_grad():
out = model(input_ids=input_ids, attention_mask=attn)
logits = out.logits
# Causal next-token shift.
new_logits = logits[:, :-1, :]
new_labels = labels[:, 1:]
# calculate per token loss
loss = F.cross_entropy(
# converting to float to avoid bfloat16 precision loss
new_logits.float().reshape(-1, new_logits.size(-1)),
new_labels.reshape(-1),
ignore_index=-100,
reduction='none'
)
print("Per token loss:", loss.tolist()[:5], "...")
# Calculate average cross-entropy over valid tokens.
loss = F.cross_entropy(
# converting to float to avoid bfloat16 precision loss
new_logits.float().reshape(-1, new_logits.size(-1)),
new_labels.reshape(-1),
ignore_index=-100,
)
perplexity = torch.exp(loss)
vocab_size = new_logits.size(-1)
print("loss:", loss.item())
print("perplexity:", perplexity.item())Per token loss: [6.58029842376709, 0.8311704397201538, 2.2088489532470703, 0.7635747194290161, 3.1106820106506348] ...
loss: 2.319715738296509
perplexity: 10.172781944274902Tying Weights
The input embedding table turns token ids into hidden vectors, and the language-model head turns hidden vectors back into scores for all V tokens. With tied weights, these two operations reuse the same V x D matrix instead of storing two copies. So one single matrix is expected to do the work of both converting tokens to embeddings and take hidden dimension vectors (same size as embedding) and convert them to logits over vocab size. Not ideal but in case of small models, these two matrices are very huge, as we will see below, hence it make sense to tie the weights
input_embeddings = model.get_input_embeddings()
output_embeddings = model.get_output_embeddings()
input_weight = input_embeddings.weight
output_weight = output_embeddings.weight
weights_are_tied = (
input_weight is output_weight
or input_weight.data_ptr() == output_weight.data_ptr()
)
# the config can also mention about tying the weights.
config_ties_weights = getattr(model.config, "tie_word_embeddings", None)
# calculate the params
input_embedding_params = sum(p.numel() for p in input_embeddings.parameters())
lm_head_params = sum(p.numel() for p in output_embeddings.parameters())
total_params = sum(p.numel() for p in model.parameters())
# A tied head reuses the input embedding parameters. An untied head adds a
# separate output matrix.
if weights_are_tied:
tied_total_params = total_params
untied_total_params = total_params + lm_head_params
else:
tied_total_params = total_params - lm_head_params
untied_total_params = total_params
print("config.tie_word_embeddings:", config_ties_weights)
print("weights actually tied: ", weights_are_tied)
print()
print("input embedding params: ", f"{input_embedding_params:,}")
print("LM head params if separate: ", f"{lm_head_params:,}")
print()
print("Total params when tied: ", f"{tied_total_params:,}")
print("Total params, if output matrix is separate: ", f"{untied_total_params:,}")
print()
print("extra params when untied: ", f"{untied_total_params - tied_total_params:,}")config.tie_word_embeddings: True
weights actually tied: True
input embedding params: 254,279,680
LM head params if separate: 254,279,680
Total params when tied: 852,985,920
Total params, if output matrix is separate: 1,107,265,600
extra params when untied: 254,279,680Chunk Loss Calculation
Until now for the calculation of cross entropy loss, we passed the complete tensor to cross-entropy loss, where F.cross_entropy takes in [num_tokens, vocab_size], and in our case the vocab size is quite big ~262k, hence the memory used is huge when the num_tokens increases. To avoid it we can simply calculate the cross entropy loss on chunks of tokens instead of full sequence length at once. This is equivalent in calculations.
[num_tokens, vocab_size] -> k * [chunk_size, vocab_size]
Now since, we calculate loss for smaller chunks, it may take longer time, but the peak memory is much smaller.
You can restart the notebook, to reset the memory and get better fresh numbers
from transformers import AutoModelForImageTextToText, AutoTokenizer
import torch
import torch.nn.functional as Fmodel_id = "Qwen/Qwen3.5-0.8B"
model = AutoModelForImageTextToText.from_pretrained(model_id, dtype=torch.bfloat16, device_map="auto", attn_implementation="eager")tok = AutoTokenizer.from_pretrained(model_id)
very_very_long_input = """Elephants are the largest living land animals. Three living species are currently recognised: the African bush elephant (Loxodonta africana), the African forest elephant (L. cyclotis), and the Asian elephant (Elephas maximus). They are the only surviving members of the family Elephantidae and the order Proboscidea; extinct relatives include mammoths and mastodons. Distinctive features of elephants include a long proboscis called a trunk, tusks, large ear flaps, pillar-like legs, and tough but sensitive grey skin. The trunk is prehensile, bringing food and water to the mouth and grasping objects. Tusks, which are derived from the incisor teeth, serve both as weapons and as tools for moving objects and digging. The large ear flaps assist in maintaining a constant body temperature as well as in communication. African elephants have larger ears and concave backs, whereas Asian elephants have smaller ears and convex or level backs.
Elephants are scattered throughout sub-Saharan Africa, South Asia, and Southeast Asia. They live in different habitats, including savannahs, forests, deserts, and marshes. They are herbivorous, and they stay near water when it is accessible. They are considered to be keystone species, due to their impact on their environments. Elephants have a fission–fusion society, in which multiple family groups come together to socialise. Females (cows) tend to live in family groups, which can consist of one female with its calves or several related females with offspring. The leader of a female group, usually the oldest cow, is known as the matriarch.
Males (bulls) leave their family groups when they reach puberty and may live alone or with other males. Adult bulls mostly interact with family groups when looking for a mate. They enter a state of increased testosterone and aggression known as musth, which helps them gain dominance over other males as well as reproductive success. Calves are the centre of attention in their family groups and rely on their mothers for as long as three years. Elephants can live up to 70 years in the wild. They communicate by touch, sight, smell, and sound; elephants use infrasound and seismic communication over long distances. Elephant intelligence has been compared with that of primates and cetaceans. They appear to have self-awareness, and possibly show concern for dying and dead individuals of their kind.
African bush elephants and Asian elephants are listed as endangered and African forest elephants as critically endangered on the IUCN Red Lists. One of the biggest threats to elephant populations is the ivory trade, as the animals are poached for their ivory tusks. Other threats to wild elephants include habitat destruction and conflicts with local people. Elephants are used as working animals in Asia. In the past, they were used in war; today, they are often controversially put on display in zoos, or employed for entertainment in circuses. Elephants have an iconic status in human culture and have been widely featured in art, folklore, religion, literature, and popular culture."""
inputs = tok(very_very_long_input, return_tensors="pt").to(model.device)
with torch.no_grad(): # since we are only doing inference, we don't want to accumulate gradients
out = model(
**inputs,
output_hidden_states=True,
return_dict=True,
)[transformers] `causal_conv1d_fn` is falling back to its reference PyTorch implementation because `causal_conv1d` is not installed. This is correct but much slower; install `causal_conv1d` for the optimized kernel.
[transformers] `chunk_gated_delta_rule` is falling back to its reference PyTorch implementation because `flash-linear-attention` is not installed. This is correct but much slower; install `flash-linear-attention` for the optimized kernel.# Final hidden states: [batch, sequence_length, hidden_size]
last_hidden = out.hidden_states[-1]
print("Last hidden state:", tuple(last_hidden.shape))
# hidden at position t predicts token at position t + 1.
# so we are taking all but last one, since we don't know the last token GT
hidden = last_hidden[:, :-1, :]
print("Hidden state for 1st to n-1th tokens:", tuple(hidden.shape))
# similarly for target, we take from second token to the end
targets = inputs["input_ids"][:, 1:]
print("Targets for 2nd to nth tokens:", tuple(targets.shape))
# Flatten batch and sequence dimensions.
hidden = hidden.reshape(-1, hidden.size(-1))
print("Flattened Hidden:", tuple(hidden.shape))
targets = targets.reshape(-1)
print("Flattened Targets:", tuple(targets.shape))Last hidden state: (1, 632, 1024)
Hidden state for 1st to n-1th tokens: (1, 631, 1024)
Targets for 2nd to nth tokens: (1, 631)
Flattened Hidden: (631, 1024)
Flattened Targets: (631,)# Use the actual output head from the model.
output_head = model.get_output_embeddings()
print("Weight Matrix:", output_head)Weight Matrix: Linear(in_features=1024, out_features=248320, bias=False)def chunked_cross_entropy():
chunk_size = 32
total_loss = torch.zeros(
(),
device=hidden.device,
dtype=torch.float32,
)
total_tokens = 0
for start in range(0, hidden.size(0), chunk_size):
end = start + chunk_size
chunk_hidden = hidden[start:end]
chunk_targets = targets[start:end]
# prediction only for the chunked tokens
chunk_logits = output_head(chunk_hidden)
# we sum the loss here, and at the end we can divide by number of tokens
chunk_loss = F.cross_entropy(chunk_logits.float(), chunk_targets, reduction="sum")
total_loss += chunk_loss
total_tokens += chunk_targets.numel()
return total_loss / total_tokens
def full_cross_entropy():
# Creates the complete [number_of_tokens, vocabulary_size] tensor.
logits = output_head(hidden)
return F.cross_entropy(logits.float(), targets)# The measurement function includes CUDA synchronization. Without synchronization,
# GPU timing can be inaccurate because CUDA operations are asynchronous.
def measure_peak_memory_and_time(function, repeats=10):
# Warm up once so initialization overhead does not affect the timing.
with torch.no_grad():
for _ in range(2):
warmup_loss = function()
del warmup_loss
torch.cuda.synchronize() # Wait for all kernels in all streams on a CUDA device to complete
torch.cuda.empty_cache() # Release all unoccupied cached memory
torch.cuda.reset_peak_memory_stats() # Reset the "peak" stats tracked by the CUDA memory allocator
starting_memory = torch.cuda.memory_allocated() # Return the current GPU memory occupied by tensors in bytes
# CUDA events are synchronization markers that can be used to monitor the device's progress, to accurately measure timing, and to synchronize CUDA streams.
start_event = torch.cuda.Event(enable_timing=True)
end_event = torch.cuda.Event(enable_timing=True)
start_event.record()
with torch.no_grad():
for _ in range(repeats):
loss = function()
end_event.record()
torch.cuda.synchronize()
elapsed_ms = start_event.elapsed_time(end_event)
average_time_ms = elapsed_ms / repeats
peak_memory = torch.cuda.max_memory_allocated()
peak_memory_used = peak_memory - starting_memory
return (loss.detach(), peak_memory_used, average_time_ms)ordinary_loss, ordinary_peak, ordinary_time = measure_peak_memory_and_time(full_cross_entropy)
chunked_loss, chunked_peak, chunked_time = measure_peak_memory_and_time(chunked_cross_entropy)def mb(number_of_bytes):
return number_of_bytes / (1024 ** 2)print("ordinary loss:", ordinary_loss.item())
print("chunked loss: ", chunked_loss.item())
print()
print(f"ordinary CE peak memory: {mb(ordinary_peak):.2f} MB")
print(f"chunked CE peak memory: {mb(chunked_peak):.2f} MB")
print(f"memory ratio: {ordinary_peak / max(chunked_peak, 1):.2f}x")
print()
print(f"ordinary CE time: {ordinary_time:.3f} ms")
print(f"chunked CE time: {chunked_time:.3f} ms")
print(f"time ratio: {chunked_time / ordinary_time:.2f}x")
print()
print("losses match:", torch.allclose(ordinary_loss, chunked_loss, atol=1e-2, rtol=1e-3))ordinary loss: 2.325397253036499
chunked loss: 2.325397253036499
ordinary CE peak memory: 1494.86 MB
chunked CE peak memory: 75.78 MB
memory ratio: 19.73x
ordinary CE time: 8.259 ms
chunked CE time: 20.823 ms
time ratio: 2.52x
losses match: TrueAs you can see the loss is the same upto floating-point rounding, but the peak memory used for full cross entropy is 20x more than the memory used for the chunked cross entropy with the trade-off of 2.5x more time to calculate
Lets add another head to our model
We will try to do two predictions at once, instead of predicting the next token, we will add a head, which will try to predict t+2 token. Since we don’t want to train it in full, we will train on a small set of sequences and test it on one of them itself
The hidden state at position t will be used by the original head to predict token t+1. The new head sees the same hidden state, but its target will be shifted by two positions. We will copy the original head weights for the new head, since we starting from random initialization will not converge the head with very few examples. Also it will give the new head a useful vocabulary projection, but we should make sure that the copy must be independent.
Directly assigning head_t2 = head_t1 would share parameters rather than create a second head.
Then we will train both heads with the sum of their losses. The t+2 head has fewer valid positions and is naturally a harder task, because it must predict two tokens ahead without seeing the token at t+1 as an input. But let’s see what will be the actual loss values for each
from transformers import AutoModelForImageTextToText, AutoTokenizer
import torch
import torch.nn as nn
import torch.nn.functional as Fmodel_id = "Qwen/Qwen3.5-0.8B"
model = AutoModelForImageTextToText.from_pretrained(model_id, dtype=torch.bfloat16, device_map="auto", attn_implementation="eager")
tok = AutoTokenizer.from_pretrained(model_id)# The existing model output head will be head #1
head_t1 = model.get_output_embeddings()
# new head with the same shape and dtype/device.
head_t2 = nn.Linear(head_t1.in_features, head_t1.out_features, bias=False,
device=head_t1.weight.device, dtype=head_t1.weight.dtype
)
# copy the weights from head #1 to head #2
with torch.no_grad(): # no grad because during backprop we don't want the grad from
# head #2 moving to head #1
head_t2.weight.copy_(head_t1.weight)
print("head 1 and head 2 share storage:", head_t1.weight.data_ptr() == head_t2.weight.data_ptr())
# check if the weights are same
print("head 1 and head 2 have same weights:", torch.all(head_t1.weight == head_t2.weight).item())head 1 and head 2 share storage: False
head 1 and head 2 have same weights: TrueThe corpus is intentionally tiny, so the result will show memorization rather than general language ability.
# A small corpus for training our network
docs = [
"The cat sat on the mat and watched the birds outside the window.",
"Birds fly south in winter to find warmer weather and better food.",
"A young girl opened a book and started to read beside the fireplace.",
"The river runs through the valley after the rain has filled the streams.",
"Dogs bark at the moon when the night is still and unusually quiet.",
"The baker made fresh bread before the sun came up over the village.",
"Green leaves fall from the trees in late autumn and cover the ground.",
"The boy kicked the ball across the empty field toward his friends.",
"A small boat crossed the lake while the wind moved through the reeds.",
"The scientist recorded each result carefully before repeating the experiment.",
"The teacher explained the problem slowly and answered every question.",
"Clouds gathered above the hills before the first drops of rain arrived.",
"The farmer planted seeds in spring and waited patiently for the harvest.",
"A bright star appeared above the quiet town as evening turned to night.",
"The musician practiced the difficult song until every note sounded clear.",
"Children built a sandcastle near the shore while waves rolled onto the beach.",
"The old library contained many stories about explorers and distant lands.",
"Fresh snow covered the road and made the surrounding forest look peaceful.",
"The engineer designed a bridge that could safely carry heavy trucks.",
"After breakfast, the family walked through the park and listened to birds.",
"The doctor examined the patient carefully before recommending a treatment.",
"A red train traveled across the countryside during the early morning hours.",
]
print("Number of sequences:", len(docs))Number of sequences: 22train_enc = tok(docs, return_tensors="pt", padding="max_length",
max_length=32, truncation=True).to(model.device)
print("Train sequence sample tokens")
train_input_ids = train_enc["input_ids"]
print(train_input_ids[0].tolist())
train_attention_mask = train_enc["attention_mask"]
print(train_attention_mask[0].tolist())
# Align each target with the hidden state that should predict it.
print("Head #1 target tokens. Each hidden state at t predicts the input token at t+1")
targets_t1 = train_input_ids[:, 1:].clone()
print(targets_t1[0].tolist())
targets_t1[train_attention_mask[:, 1:] == 0] = -100
print(targets_t1[0].tolist())
# for head 2 targets, we want to check with t+2 tokens for token t
print("Head #2 target tokens. Each hidden state at t predicts the input token at t+2")
targets_t2 = train_input_ids[:, 2:].clone()
print(targets_t2[0].tolist())
targets_t2[train_attention_mask[:, 2:] == 0] = -100
print(targets_t2[0].tolist())Train sequence sample tokens
[760, 7993, 7338, 383, 279, 5344, 321, 14936, 279, 19053, 4732, 279, 3136, 13, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044]
[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
Head #1 target tokens. Each hidden state at t predicts the input token at t+1
[7993, 7338, 383, 279, 5344, 321, 14936, 279, 19053, 4732, 279, 3136, 13, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044]
[7993, 7338, 383, 279, 5344, 321, 14936, 279, 19053, 4732, 279, 3136, 13, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100]
Head #2 target tokens. Each hidden state at t predicts the input token at t+2
[7338, 383, 279, 5344, 321, 14936, 279, 19053, 4732, 279, 3136, 13, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044, 248044]
[7338, 383, 279, 5344, 321, 14936, 279, 19053, 4732, 279, 3136, 13, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100]# include the original model params along with the new head
opt = torch.optim.AdamW(list(model.parameters()) + list(head_t2.parameters()), lr=1e-5)
model.train()
loss_history = []
num_steps = 100
for step in range(num_steps + 1):
out = model(
input_ids=train_input_ids,
attention_mask=train_attention_mask,
output_hidden_states=True,
return_dict=True,
)
hidden = out.hidden_states[-1]
# The model's existing output head predicts t+1.
logits_t1 = out.logits[:, :-1, :]
loss_t1 = F.cross_entropy(
logits_t1.float().reshape(-1, logits_t1.size(-1)),
targets_t1.reshape(-1),
ignore_index=-100,
)
# The copied, but independent, head predicts t+2.
logits_t2 = head_t2(hidden[:, :-2, :])
loss_t2 = F.cross_entropy(
logits_t2.float().reshape(-1, logits_t2.size(-1)),
targets_t2.reshape(-1),
ignore_index=-100,
)
total_loss = loss_t1 + loss_t2
loss_history.append((step, loss_t1.item(), loss_t2.item(), total_loss.item()))
if step % 50 == 0:
print(
f"step {step:3d} t+1={loss_t1.item():.4f} "
f"t+2={loss_t2.item():.4f} sum={total_loss.item():.4f}"
)
if step == num_steps:
break
opt.zero_grad(set_to_none=True)
total_loss.backward()
opt.step()step 0 t+1=3.8985 t+2=9.2174 sum=13.1159
step 50 t+1=0.7429 t+2=0.7333 sum=1.4763
step 100 t+1=0.6516 t+2=0.6435 sum=1.2951The t+2 loss will naturally start higher than the t+1 loss and decrease during training, but will usually remain higher. Predicting t+2 is less certain than predicting the immediately following token, and also it has two less possible target positions per sequence.
Now let’s do inference on one of our training samples, and see how well we have memorized the data, and does our head #2 predict the output well
model.eval()
example = docs[0]
example_enc = tok(example, return_tensors="pt").to(model.device)
example_ids = example_enc["input_ids"][0]
example_tokens = tok.convert_ids_to_tokens(example_ids.tolist())
with torch.no_grad():
example_out = model(
**example_enc,
output_hidden_states=True,
return_dict=True,
)
example_hidden = example_out.hidden_states[-1]
# Head 1: h_t -> token t+1.
predictions_t1 = example_out.logits[0, :-1].argmax(dim=-1)
# Head 2: h_t -> token t+2.
predictions_t2 = head_t2(example_hidden[:, :-2, :])[0].argmax(dim=-1)
print("Input:", example)
print()
print("position | input token | actual t+1 | head 1 prediction | actual t+2 | head 2 prediction")
print("-" * 86)
for t, input_token in enumerate(example_tokens):
actual_t1 = example_tokens[t + 1] if t + 1 < len(example_tokens) else "-"
actual_t2 = example_tokens[t + 2] if t + 2 < len(example_tokens) else "-"
predicted_t1 = predictions_t1[t].item() if t < len(predictions_t1) else None
predicted_t2 = predictions_t2[t].item() if t < len(predictions_t2) else None
predicted_t1 = (
tok.convert_ids_to_tokens([predicted_t1])[0]
if predicted_t1 is not None else "-"
)
predicted_t2 = (
tok.convert_ids_to_tokens([predicted_t2])[0]
if predicted_t2 is not None else "-"
)
print(
f"{t:8d} | {input_token:11s} | {actual_t1:11s} | "
f"{predicted_t1:18s} | {actual_t2:11s} | {predicted_t2}"
)Input: The cat sat on the mat and watched the birds outside the window.
position | input token | actual t+1 | head 1 prediction | actual t+2 | head 2 prediction
--------------------------------------------------------------------------------------
0 | The | Ġcat | Ġold | Ġsat | Ġexplained
1 | Ġcat | Ġsat | Ġsat | Ġon | Ġon
2 | Ġsat | Ġon | Ġon | Ġthe | Ġthe
3 | Ġon | Ġthe | Ġmat | Ġmat | Ġmat
4 | Ġthe | Ġmat | Ġand | Ġand | Ġand
5 | Ġmat | Ġand | Ġand | Ġwatched | Ġwatched
6 | Ġand | Ġwatched | Ġwatched | Ġthe | Ġthe
7 | Ġwatched | Ġthe | Ġthe | Ġbirds | Ġbirds
8 | Ġthe | Ġbirds | Ġoutside | Ġoutside | Ġoutside
9 | Ġbirds | Ġoutside | Ġoutside | Ġthe | Ġthe
10 | Ġoutside | Ġthe | Ġthe | Ġwindow | Ġwindow
11 | Ġthe | Ġwindow | Ġwindow | . | .
12 | Ġwindow | . | . | - | -
13 | . | - | - | - | -Now let’s try it on some different sequence, which we have not trained on
model.eval()
example = "The Capital of India is New Delhi."
example_enc = tok(example, return_tensors="pt").to(model.device)
example_ids = example_enc["input_ids"][0]
example_tokens = tok.convert_ids_to_tokens(example_ids.tolist())
with torch.no_grad():
example_out = model(
**example_enc,
output_hidden_states=True,
return_dict=True,
)
example_hidden = example_out.hidden_states[-1]
# Head 1: h_t -> token t+1.
predictions_t1 = example_out.logits[0, :-1].argmax(dim=-1)
# Head 2: h_t -> token t+2.
predictions_t2 = head_t2(example_hidden[:, :-2, :])[0].argmax(dim=-1)
print("Input:", example)
print()
print("position | input token | actual t+1 | head 1 prediction | actual t+2 | head 2 prediction")
print("-" * 86)
for t, input_token in enumerate(example_tokens):
actual_t1 = example_tokens[t + 1] if t + 1 < len(example_tokens) else "-"
actual_t2 = example_tokens[t + 2] if t + 2 < len(example_tokens) else "-"
predicted_t1 = predictions_t1[t].item() if t < len(predictions_t1) else None
predicted_t2 = predictions_t2[t].item() if t < len(predictions_t2) else None
predicted_t1 = (
tok.convert_ids_to_tokens([predicted_t1])[0]
if predicted_t1 is not None else "-"
)
predicted_t2 = (
tok.convert_ids_to_tokens([predicted_t2])[0]
if predicted_t2 is not None else "-"
)
print(
f"{t:8d} | {input_token:11s} | {actual_t1:11s} | "
f"{predicted_t1:18s} | {actual_t2:11s} | {predicted_t2}"
)Input: The Capital of India is New Delhi.
position | input token | actual t+1 | head 1 prediction | actual t+2 | head 2 prediction
--------------------------------------------------------------------------------------
0 | The | ĠCapital | Ġold | Ġof | Ġexplained
1 | ĠCapital | Ġof | Ġof | ĠIndia | Ġof
2 | Ġof | ĠIndia | Ġthe | Ġis | Ġthe
3 | ĠIndia | Ġis | Ġthe | ĠNew | Ġthe
4 | Ġis | ĠNew | Ġthe | ĠDelhi | Ġthe
5 | ĠNew | ĠDelhi | ĠDelhi | . | ĠDelhi
6 | ĠDelhi | . | . | - | -
7 | . | - | - | - | -Our model has officialy gone crazy and is just predicting the, since we have overfit to our very tiny dataset