|
| 1 | +from trainer.unlearn.base import UnlearnTrainer |
| 2 | + |
| 3 | +import torch |
| 4 | +import torch.nn.functional as F |
| 5 | + |
| 6 | + |
| 7 | +def cross_entropy_unlearning_loss( |
| 8 | + logits: torch.Tensor, |
| 9 | + labels: torch.Tensor, |
| 10 | + ignore_index: int = -100, |
| 11 | +) -> torch.Tensor: |
| 12 | + """ |
| 13 | + Implementation of Cross Entropy Unlearning Loss (CE-U). |
| 14 | +
|
| 15 | + This function creates a modified target distribution by setting the logit corresponding to the true label to negative infinity, effectively forcing the model to assign zero probability to the correct answer. The loss then minimizes the KL divergence between this target distribution and the model's output. |
| 16 | +
|
| 17 | + Args: |
| 18 | + logits: Model output logits with shape [batch_size, sequence_length, vocabulary_size] |
| 19 | + labels: Ground truth token indices with shape [batch_size, sequence_length] |
| 20 | + ignore_index: Token indices to ignore in the loss calculation (typically padding) |
| 21 | +
|
| 22 | + Returns: |
| 23 | + A scalar tensor representing the mean unlearning loss across valid positions |
| 24 | + """ |
| 25 | + batch_size, sequence_length, vocabulary_size = logits.shape |
| 26 | + # Extract valid logits and labels based on ignore_index. |
| 27 | + if ignore_index is not None: |
| 28 | + # Shape: [batch_size, sequence_length], boolean mask |
| 29 | + valid_mask = labels != ignore_index |
| 30 | + # Shape: [num_valid_positions, vocabulary_size] |
| 31 | + valid_logits = logits[valid_mask] |
| 32 | + # Shape: [num_valid_positions] |
| 33 | + valid_labels = labels[valid_mask] |
| 34 | + else: |
| 35 | + # Shape: [batch_size*sequence_length, vocabulary_size] |
| 36 | + valid_logits = logits.view(-1, vocabulary_size) |
| 37 | + # Shape: [batch_size*sequence_length] |
| 38 | + valid_labels = labels.view(-1) |
| 39 | + |
| 40 | + # Create a copy of valid_logits to generate the target distribution |
| 41 | + # Shape: [num_valid_positions, vocabulary_size] |
| 42 | + valid_target_logits = valid_logits.detach().clone() |
| 43 | + |
| 44 | + # Suppress the logits corresponding to the true token by setting them to -inf. |
| 45 | + # This ensures that the probability for the true token is effectively zero after softmax. |
| 46 | + valid_target_logits.scatter_( |
| 47 | + dim=-1, |
| 48 | + index=valid_labels.unsqueeze(-1), # Shape: [num_valid_positions, 1] |
| 49 | + value=float("-inf"), |
| 50 | + ) # Result shape: [num_valid_positions, vocabulary_size] |
| 51 | + |
| 52 | + # Apply softmax to generate the target probability distribution |
| 53 | + # Shape: [num_valid_positions, vocabulary_size] |
| 54 | + valid_target_probabilities = F.softmax(valid_target_logits, dim=-1) |
| 55 | + |
| 56 | + # Compute the cross entropy loss between input logits and target probabilities |
| 57 | + # The loss is averaged over the valid positions and returns a scalar tensor |
| 58 | + return F.cross_entropy( |
| 59 | + input=valid_logits, |
| 60 | + target=valid_target_probabilities, |
| 61 | + ) |
| 62 | + |
| 63 | + |
| 64 | +def compute_batch_ceu(model, inputs, ignore_first_n_answer_tokens=1): |
| 65 | + outputs = model(**inputs) |
| 66 | + logits = outputs.logits |
| 67 | + labels = inputs["labels"] |
| 68 | + |
| 69 | + # Implement the trick to ignore the first n answer tokens mentioned in the footnote in the Training Settings section of arXiv:2503.01224 |
| 70 | + valid_mask = labels != -100 |
| 71 | + update_mask = ( |
| 72 | + valid_mask.cumsum(dim=-1) <= ignore_first_n_answer_tokens |
| 73 | + ) & valid_mask |
| 74 | + labels_without_first_n_answer_tokens = labels.masked_fill(update_mask, -100) |
| 75 | + |
| 76 | + shifted_labels = labels_without_first_n_answer_tokens[..., 1:].contiguous() |
| 77 | + shifted_logits = logits[..., :-1, :].contiguous() |
| 78 | + loss = cross_entropy_unlearning_loss( |
| 79 | + shifted_logits, shifted_labels, ignore_index=-100 |
| 80 | + ) |
| 81 | + return loss, outputs |
| 82 | + |
| 83 | + |
| 84 | +class CEU(UnlearnTrainer): |
| 85 | + def __init__(self, ignore_first_n_answer_tokens=1, *args, **kwargs): |
| 86 | + super().__init__(*args, **kwargs) |
| 87 | + self.ignore_first_n_answer_tokens = ignore_first_n_answer_tokens |
| 88 | + |
| 89 | + def compute_loss(self, model, inputs, return_outputs=False): |
| 90 | + forget_inputs = inputs["forget"] |
| 91 | + loss, outputs = compute_batch_ceu( |
| 92 | + model, |
| 93 | + forget_inputs, |
| 94 | + ignore_first_n_answer_tokens=self.ignore_first_n_answer_tokens, |
| 95 | + ) |
| 96 | + return (loss, outputs) if return_outputs else loss |
0 commit comments