Add new top1 loss functions - #868
Open
klei22 wants to merge 2 commits into
Open
Conversation
There was a problem hiding this comment.
Pull request overview
This PR adds two new top-1-accuracy-oriented loss variants to the training loss registry (top1_corrective_ce and top1_confidence_gap), along with CLI knobs, documentation, and a sweep YAML to experiment with their hyperparameters.
Changes:
- Added utilities to standardize logits/targets flattening and top-1 correctness masking, and refactored
top1_focus_lossto use them. - Implemented and registered two new loss variants (
top1_corrective_ce,top1_confidence_gap) plus their CLI arguments. - Added documentation and an exploration sweep YAML for comparing these losses against a cross-entropy baseline and warm-start schedules.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
train_variations/loss_variants.py |
Adds new top-1 loss variants, utilities (_flatten_logits_targets, top1_correct_mask), and wires them into the loss registry/arg dispatch. |
train_args.py |
Introduces CLI hyperparameters for the new loss variants. |
explorations/top1_correctness_loss_sweep.yaml |
Adds a sweep spec to compare new losses and hyperparameter ranges, including warm-start schedules. |
documentation/top1_correct_losses.md |
Documents motivation, implementation, usage, and sweep suggestions for top-1-oriented loss variants. |
Comments suppressed due to low confidence (1)
train_variations/loss_variants.py:425
top1_corrective_ce_lossandtop1_confidence_gap_lossare each defined twice in this module (first at lines ~339/368, then again starting at ~421/450). The later redefinitions silently override the earlier ones, which is confusing and makes future edits error-prone. Remove the duplicate second block and keep a single canonical definition for each loss.
def top1_corrective_ce_loss(
logits: torch.Tensor,
targets: torch.Tensor,
*,
iter_num: int | None = None,
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+339
to
+365
| def top1_corrective_ce_loss( | ||
| logits: torch.Tensor, | ||
| targets: torch.Tensor, | ||
| *, | ||
| iter_num: int | None = None, | ||
| boost: float = 1.0, | ||
| ) -> torch.Tensor: | ||
| """Cross entropy that up-weights currently top-1-incorrect tokens. | ||
|
|
||
| This directly mirrors ``avg_top1_correct`` at batch scale: the loss estimates | ||
| the current batch's top-1 error rate and uses it to increase emphasis on | ||
| tokens whose argmax prediction is wrong. | ||
| """ | ||
|
|
||
| logits_flat, targets_flat, mask = _flatten_logits_targets(logits, targets) | ||
| losses = F.cross_entropy(logits_flat, targets_flat, reduction="none", ignore_index=-1) | ||
| if not mask.any(): | ||
| return losses.new_full((), 0.0) | ||
|
|
||
| with torch.no_grad(): | ||
| predictions = torch.argmax(logits_flat, dim=-1) | ||
| incorrect = (predictions != targets_flat) & mask | ||
| batch_error_rate = incorrect.float().sum() / mask.float().sum().clamp_min(1.0) | ||
| weights = torch.ones_like(losses) | ||
| weights[incorrect] = 1.0 + boost * batch_error_rate | ||
|
|
||
| return (losses[mask] * weights[mask]).mean() |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This pull request introduces new loss functions inspired by the
avg_top1_correctmetric to improve top-1 prediction accuracy during training. It adds two new loss variants—top1_corrective_ceandtop1_confidence_gap—along with their configuration options, documentation, and sweep YAML for experimentation. The changes also refactor some existing loss code for clarity and reusability.New loss functions and configuration:
top1_corrective_ce_lossfunction, which up-weights cross-entropy loss for currently top-1-incorrect tokens based on the batch's top-1 error rate, and thetop1_confidence_gap_loss, which adds a differentiable penalty for the target logit's gap to the strongest competitor. Both are registered in the loss function registry and can be selected via command-line arguments. [1] [2] [3]--top1_corrective_boostand--top1_confidence_gap_betaCLI arguments intrain_args.pyto control the strength of the new losses.Documentation and experiment configuration:
documentation/top1_correct_losses.mdto explain the motivation, implementation, and usage of the new loss variants, with practical sweep suggestions.explorations/top1_correctness_loss_sweep.yamlto define a sweep comparing the new losses against cross-entropy, including hyperparameter ranges and warm-start schedules.Refactoring and utility improvements:
_flatten_logits_targetsandtop1_correct_maskto standardize handling of logits, targets, and masks for loss computation and correctness checking.top1_focus_lossto use the newtop1_correct_maskutility and clarified the batch penalty computation.