Add new top1 loss functions

documentation/top1_correct_losses.md

@@ -0,0 +1,87 @@
+# `avg_top1_correct` and top-1 loss variations
+
+`avg_top1_correct` is logged during validation as the mean fraction of tokens
+whose highest-probability prediction equals the target token. In `Trainer`, the
+validation pass computes probabilities, takes the top-1 index with `probs.max`,
+and averages `(top1_idx == Y).float()` into `top1_correct`. The TensorBoard tag
+is written as `<dataset>/avg_top1_correct`.
+
+Because the metric uses `argmax`, it is not directly differentiable. Losses that
+optimize for this metric therefore need to use one of two patterns:
+
+1. **Discrete routing, continuous gradient**: use top-1 correctness only under
+   `torch.no_grad()` to decide which examples should receive more or less normal
+   cross-entropy gradient.
+2. **Differentiable surrogate**: replace the hard top-1 flip with a smooth logit
+   gap penalty that rewards the target logit for crossing the strongest
+   competing logit.
+
+## Existing losses that already follow this idea
+
+- `top1_focus`: adds a batch-level penalty proportional to
+  `1 - batch_top1_correct`.
+- `skip_correct_top1`: drops tokens that are already top-1 correct.
+- `attenuated_correct_top1`: keeps all tokens, but down-weights top-1-correct
+  tokens by `--correct_top1_attenuation`.
+- `distance_attenuated_top1`: attenuates cross entropy according to the logit
+  distance between the current top prediction and the target.
+- `top1_margin` and `top1_ratio`: differentiable gap-style objectives that push
+  the target above the best non-target token.
+
+## New variations added here
+
+### `top1_corrective_ce`
+
+This loss uses a batch-local analogue of `avg_top1_correct`:
+
+```text
+batch_error_rate = 1 - mean(argmax(logits) == targets)
+loss = mean(CE(token) * (1 + boost * batch_error_rate) for top1-wrong tokens,
+            CE(token) for top1-correct tokens)
+```
+
+Use it when you want training to concentrate more on batches where the current
+model is making many top-1 mistakes without completely removing the gradient
+from already-correct tokens.
+
+Example:
+
+```bash
+python train.py --loss_fn top1_corrective_ce --top1_corrective_boost 1.0
+```
+
+### `top1_confidence_gap`
+
+This loss keeps standard cross entropy and adds a smooth top-1 surrogate:
+
+```text
+loss = CE + beta * mean(softplus(best_non_target_logit - target_logit))
+```
+
+The penalty is small when the target already beats the strongest competitor and
+large when another token is still ahead. Unlike `avg_top1_correct`, it provides a
+usable gradient before the top-1 decision flips.
+
+Example:
+
+```bash
+python train.py --loss_fn top1_confidence_gap --top1_confidence_gap_beta 0.5
+```
+
+## Practical sweep suggestion
+
+Start with cross entropy as a baseline, then compare one discrete-weighting loss
+and one differentiable-surrogate loss:
+
+```yaml
+- loss_fn: ["cross_entropy"]
+- loss_fn: ["top1_corrective_ce"]
+  top1_corrective_boost: [0.5, 1.0, 2.0]
+- loss_fn: ["top1_confidence_gap"]
+  top1_confidence_gap_beta: [0.25, 0.5, 1.0]
+- loss_schedule: ["0:cross_entropy,10000:top1_confidence_gap"]
+```
+
+Track both validation loss and `<dataset>/avg_top1_correct`; top-1-oriented
+losses can improve discrete accuracy while sometimes worsening calibration or
+cross-entropy.

explorations/top1_correctness_loss_sweep.yaml

@@ -0,0 +1,64 @@
+# explorations/top1_correctness_loss_sweep.yaml
+---
+# Compare the avg_top1_correct-inspired losses against a cross-entropy baseline.
+# This sweep keeps the architecture fixed and varies only the new top-1 loss
+# hyperparameters plus a warm-start schedule option for the differentiable gap
+# surrogate.
+
+common_group:
+  max_iters: [20000]
+  n_layer: [6]
+  n_head: [6]
+  n_embd: [384]
+  block_size: [256]
+  eval_interval: [500]
+  dataset: ["minipile"]
+  device: ["cuda"]
+  dtype: ["bfloat16"]
+  compile: [true]
+  never_save_checkpoint: [true]
+  compute_model_stats: [true]
+  tensorboard_run_name: ["top1_correctness_loss_sweep"]
+
+  # Position encoding configuration.
+  use_rotary_embeddings: [true]
+  use_abs_pos_embeddings: [false]
+  use_qk_norm: [true]
+  use_qk_norm_scale: [true]
+
+  # Match the top-1 comparison defaults used by nearby explorations.
+  use_peri_ln: [true]
+  norm_variant_wte: ["hyperspherenorm"]
+  norm_variant_attn: ["hyperspherenorm"]
+  norm_variant_output: ["hyperspherenorm"]
+  hsnorm_gain: [false]
+  hsnorm_scale: ["1.0"]
+  norm_wte_scale: ["1.0"]
+  attn_residual_combination: ["slerp"]
+  mlp_residual_combination: ["slerp"]
+
+print_model_stats: ["./print_stats/${RUN_NAME}"]
+sample_file: ["./inference_samples/${RUN_NAME}"]
+
+# Loss variants under test. Track validation loss and avg_top1_correct to catch
+# cases where a top-1-oriented objective improves discrete accuracy while
+# worsening cross-entropy/calibration.
+parameter_groups:
+  # Baseline.
+  - loss_fn: ["cross_entropy"]
+
+  # Discrete top-1 routing: emphasize currently wrong argmax predictions.
+  - loss_fn: ["top1_corrective_ce"]
+    top1_corrective_boost: [0.25, 0.5, 1.0, 2.0]
+
+  # Differentiable top-1 surrogate: penalize the best competitor beating target.
+  - loss_fn: ["top1_confidence_gap"]
+    top1_confidence_gap_beta: [0.1, 0.25, 0.5, 1.0]
+
+  # Warm-start with normal CE, then switch to the top-1-oriented variants.
+  - loss_schedule:
+      - "0:cross_entropy,5000:top1_corrective_ce"
+    top1_corrective_boost: [0.5, 1.0]
+  - loss_schedule:
+      - "0:cross_entropy,5000:top1_confidence_gap"
+    top1_confidence_gap_beta: [0.25, 0.5]

train_args.py

@@ -176,6 +176,18 @@ def parse_args():
         default=0.5,
         help='Weight for ratio penalty in top1_ratio loss.',
     )
+    training_group.add_argument(
+        '--top1_corrective_boost',
+        type=float,
+        default=1.0,
+        help='Incorrect-token boost for top1_corrective_ce loss.',
+    )
+    training_group.add_argument(
+        '--top1_confidence_gap_beta',
+        type=float,
+        default=0.5,
+        help='Weight for the top-competitor gap penalty in top1_confidence_gap loss.',
+    )
     training_group.add_argument(
         '--flatness_beta',
         type=float,

train_variations/loss_variants.py

@@ -31,6 +31,24 @@ def cross_entropy_loss(logits: torch.Tensor, targets: torch.Tensor, *, iter_num:
     return F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
 
 
+def _flatten_logits_targets(logits: torch.Tensor, targets: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Return flattened logits, targets, and a valid-target mask."""
+    logits_flat = logits.view(-1, logits.size(-1))
+    targets_flat = targets.view(-1)
+    mask = targets_flat != -1
+    return logits_flat, targets_flat, mask
+
+
+def top1_correct_mask(logits: torch.Tensor, targets: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Return the valid-token top-1 correctness mask used by avg_top1_correct."""
+    logits_flat, targets_flat, mask = _flatten_logits_targets(logits, targets)
+    if not mask.any():
+        empty = torch.empty(0, dtype=torch.bool, device=logits.device)
+        return empty, mask
+    predictions = torch.argmax(logits_flat[mask], dim=-1)
+    return predictions == targets_flat[mask], mask
+
+
 class BitBalancedCrossEntropy:
     """Cross entropy augmented with a bit-usage penalty."""
 
@@ -145,12 +163,20 @@ def top1_focus_loss(
     iter_num: int | None = None,
     alpha: float = 0.5,
 ) -> torch.Tensor:
-    """Cross entropy with an extra penalty for wrong top-1 predictions."""
+    """Cross entropy with an extra penalty for wrong top-1 predictions.
+
+    The penalty is a batch-local analogue of ``1 - avg_top1_correct``. It is
+    evaluated under ``no_grad`` because the argmax correctness indicator is not
+    differentiable, so gradients still come from cross entropy.
+    """
     ce = cross_entropy_loss(logits, targets)
-    top1 = torch.argmax(logits, dim=-1)
-    correct_top1 = (top1 == targets).float()
-    penalty = 1.0 - correct_top1
-    return ce + alpha * penalty.mean()
+    with torch.no_grad():
+        correct_top1, _ = top1_correct_mask(logits, targets)
+        if correct_top1.numel() == 0:
+            penalty = ce.new_full((), 0.0)
+        else:
+            penalty = 1.0 - correct_top1.float().mean()
+    return ce + alpha * penalty
 
 
 def skip_correct_top1_loss(
@@ -310,6 +336,62 @@ def top1_ratio_loss(
     return ce + beta * ratio_penalty.mean()
 
 
+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()
+
+
+def top1_confidence_gap_loss(
+    logits: torch.Tensor,
+    targets: torch.Tensor,
+    *,
+    iter_num: int | None = None,
+    beta: float = 0.5,
+) -> torch.Tensor:
+    """Cross entropy plus a differentiable surrogate for top-1 correctness.
+
+    ``avg_top1_correct`` changes only when the target crosses the top competing
+    class. This surrogate penalizes the softplus of that top-competitor gap,
+    providing gradients before the discrete top-1 decision flips.
+    """
+
+    ce = cross_entropy_loss(logits, targets)
+    logits_flat, targets_flat, mask = _flatten_logits_targets(logits, targets)
+    if not mask.any():
+        return ce
+    logits_sel = logits_flat[mask]
+    targets_sel = targets_flat[mask]
+    target_logits = logits_sel[torch.arange(logits_sel.size(0), device=logits.device), targets_sel]
+    others = logits_sel.clone()
+    others[torch.arange(logits_sel.size(0), device=logits.device), targets_sel] = float("-inf")
+    max_other = others.max(dim=-1).values
+    return ce + beta * F.softplus(max_other - target_logits).mean()
+
+
 # def rank_distance_loss(
 #     logits: torch.Tensor,
 #     targets: torch.Tensor,
@@ -335,6 +417,63 @@ def top1_ratio_loss(
 #     scaled[mask] = ce[mask] * rank_scale
 #     return scaled[mask].mean()
 
+
+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()
+
+
+def top1_confidence_gap_loss(
+    logits: torch.Tensor,
+    targets: torch.Tensor,
+    *,
+    iter_num: int | None = None,
+    beta: float = 0.5,
+) -> torch.Tensor:
+    """Cross entropy plus a differentiable surrogate for top-1 correctness.
+
+    ``avg_top1_correct`` changes only when the target crosses the top competing
+    class. This surrogate penalizes the softplus of that top-competitor gap,
+    providing gradients before the discrete top-1 decision flips.
+    """
+
+    ce = cross_entropy_loss(logits, targets)
+    logits_flat, targets_flat, mask = _flatten_logits_targets(logits, targets)
+    if not mask.any():
+        return ce
+    logits_sel = logits_flat[mask]
+    targets_sel = targets_flat[mask]
+    target_logits = logits_sel[torch.arange(logits_sel.size(0), device=logits.device), targets_sel]
+    others = logits_sel.clone()
+    others[torch.arange(logits_sel.size(0), device=logits.device), targets_sel] = float("-inf")
+    max_other = others.max(dim=-1).values
+    return ce + beta * F.softplus(max_other - target_logits).mean()
+
+
 def rank_distance_loss(
     logits: torch.Tensor,
     targets: torch.Tensor,
@@ -497,6 +636,8 @@ def entropy_rank_distance_focal_loss(
     "top1_margin": top1_margin_loss,
     "entropy_penalty": entropy_penalty_loss,
     "top1_ratio": top1_ratio_loss,
+    "top1_corrective_ce": top1_corrective_ce_loss,
+    "top1_confidence_gap": top1_confidence_gap_loss,
     "rank_distance": rank_distance_loss,
     "flatness_boost": flatness_boost_loss,
     "entropy_focal": entropy_focal_loss,
@@ -623,6 +764,12 @@ def rank_gamma(iter_num: int | None) -> float:
         "top1_ratio": lambda l, t, *, iter_num=None: LOSS_VARIANTS["top1_ratio"](
             l, t, iter_num=iter_num, beta=getattr(args, "top1_ratio_beta", 0.5)
         ),
+        "top1_corrective_ce": lambda l, t, *, iter_num=None: LOSS_VARIANTS["top1_corrective_ce"](
+            l, t, iter_num=iter_num, boost=getattr(args, "top1_corrective_boost", 1.0)
+        ),
+        "top1_confidence_gap": lambda l, t, *, iter_num=None: LOSS_VARIANTS["top1_confidence_gap"](
+            l, t, iter_num=iter_num, beta=getattr(args, "top1_confidence_gap_beta", 0.5)
+        ),
         "rank_distance": lambda l, t, *, iter_num=None: LOSS_VARIANTS["rank_distance"](
             l, t, iter_num=iter_num, gamma=rank_gamma(iter_num)
         ),