1{
2  "op": "replace",
3  "filename": "nanoGPT/custom_pretrain.py",
4  "content": "# ── Loss Computation ───────────────────────────────────────────────────────\ndef compute_loss(logits, targets):\n    \"\"\"Focal cross-entropy with z-loss during training.\n    \n    Focal loss (gamma=1.5) down-weights easy examples and focuses optimization\n    on hard-to-predict tokens, improving learning on the long-tail of the token\n    distribution. Z-loss penalizes large log-partition values for stability.\n    During evaluation, standard cross-entropy is returned for comparable metrics.\n    \"\"\"\n    flat_logits = logits.view(-1, logits.size(-1)).float()\n    flat_targets = targets.view(-1)\n    \n    if not torch.is_grad_enabled():\n        return F.cross_entropy(flat_logits, flat_targets, ignore_index=-1)\n    \n    # Focal loss focusing parameter\n    gamma = 1.5\n    \n    # Compute log softmax probabilities\n    log_probs = F.log_softmax(flat_logits, dim=-1)\n    \n    # Get log probability of target class\n    target_indices = flat_targets.clamp(min=0)\n    target_log_probs = log_probs.gather(1, target_indices.unsqueeze(1)).squeeze(1)\n    target_probs = target_log_probs.exp()\n    \n    # Focal weighting: (1 - p_t)^gamma\n    focal_weight = (1 - target_probs) ** gamma\n    \n    # Per-token cross-entropy\n    ce_loss = -target_log_probs\n    \n    # Mask for valid (non-ignored) tokens\n    mask = flat_targets != -1\n    \n    # Apply focal weighting with masking\n    focal_loss = (focal_weight * ce_loss * mask.float()).sum() / mask.sum()\n    \n    # Z-loss: penalize large log-partition values for stability\n    log_z = torch.logsumexp(flat_logits, dim=-1)\n    z_loss = ((log_z ** 2) * mask.float()).sum() / mask.sum()\n    \n    return focal_loss + 1e-4 * z_loss",
5  "start_line": 187,
6  "end_line": 190
7}

1{
2  "content": "def compute_loss(logits, targets):\n    \"\"\"Focal cross-entropy with z-loss during training.\n    \n    Focal loss (gamma=1.5) down-weights easy examples and focuses optimization\n    on hard-to-predict tokens, improving learning on the long-tail of the token\n    distribution. Z-loss penalizes large log-partition values for stability.\n    During evaluation, standard cross-entropy is returned for comparable metrics.\n    \"\"\"\n    flat_logits = logits.view(-1, logits.size(-1)).float()\n    flat_targets = targets.view(-1)\n    \n    if not torch.is_grad_enabled():\n        return F.cross_entropy(flat_logits, flat_targets, ignore_index=-1)\n    \n    # Focal loss focusing parameter\n    gamma = 1.5\n    \n    # Compute log softmax probabilities\n    log_probs = F.log_softmax(flat_logits, dim=-1)\n    \n    # Get log probability of target class\n    target_indices = flat_targets.clamp(min=0)\n    target_log_probs = log_probs.gather(1, target_indices.unsqueeze(1)).squeeze(1)\n    target_probs = target_log_probs.exp()\n    \n    # Focal weighting: (1 - p_t)^gamma\n    focal_weight = (1 - target_probs) ** gamma\n    \n    # Per-token cross-entropy\n    ce_loss = -target_log_probs\n    \n    # Mask for valid (non-ignored) tokens\n    mask = flat_targets != -1\n    \n    # Apply focal weighting with masking\n    focal_loss = (focal_weight * ce_loss * mask.float()).sum() / mask.sum()\n    \n    # Z-loss: penalize large log-partition values for stability\n    log_z = torch.logsumexp(flat_logits, dim=-1)\n    z_loss = ((log_z ** 2) * mask.float()).sum() / mask.sum()\n    \n    return focal_loss + 1e-4 * z_loss\n\n",
3  "end_line": 190,
4  "filename": "nanoGPT/custom_pretrain.py",
5  "start_line": 188,
6  "op": "replace"
7}

1{
2  "op": "replace",
3  "filename": "nanoGPT/custom_pretrain.py",
4  "content": "def compute_loss(logits, targets):\n    \"\"\"Combined loss: softcapped logits + z-loss + entropy bonus during training.\n    \n    1. Softcapping bounds logit magnitudes (from Gemma), improving optimization.\n    2. Z-loss penalizes large log-partition values, preventing logit explosion.\n    3. Entropy bonus prevents overconfident predictions early in training.\n    During evaluation, standard cross-entropy is returned for comparable metrics.\n    \"\"\"\n    flat_logits = logits.view(-1, logits.size(-1)).float()\n    flat_targets = targets.view(-1)\n    \n    if not torch.is_grad_enabled():\n        return F.cross_entropy(flat_logits, flat_targets, ignore_index=-1)\n    \n    # 1. Softcapping: tanh(logits / C) * C (Gemma-style)\n    softcap = 30.0\n    capped_logits = softcap * torch.tanh(flat_logits / softcap)\n    \n    # 2. Cross-entropy with label smoothing\n    ce_loss = F.cross_entropy(capped_logits, flat_targets, ignore_index=-1, label_smoothing=0.03)\n    \n    # 3. Z-loss: penalize large log-partition values\n    mask = flat_targets != -1\n    log_z = torch.logsumexp(capped_logits, dim=-1)\n    z_loss = ((log_z ** 2) * mask.float()).sum() / mask.sum()\n    \n    # 4. Entropy bonus: encourage prediction entropy\n    log_probs = F.log_softmax(capped_logits, dim=-1)\n    probs = log_probs.exp()\n    entropy = -(probs * log_probs).sum(dim=-1)\n    entropy_reg = -((entropy * mask.float()).sum() / mask.sum())  # negative to maximize entropy\n    \n    return ce_loss + 1e-4 * z_loss + 0.05 * entropy_reg",
5  "start_line": 188,
6  "end_line": 229
7}

1{
2  "content": "def compute_loss(logits, targets):\n    \"\"\"Combined loss: Gemma-style softcap + label smoothing + z-loss.\n    \n    1. Softcapping (tanh) bounds logit magnitudes for stable optimization.\n    2. Label smoothing prevents overconfident predictions.\n    3. Z-loss penalizes large log-partition values to prevent logit explosion.\n    During evaluation, standard cross-entropy is returned for comparable metrics.\n    \"\"\"\n    flat_logits = logits.view(-1, logits.size(-1)).float()\n    flat_targets = targets.view(-1)\n    \n    if not torch.is_grad_enabled():\n        return F.cross_entropy(flat_logits, flat_targets, ignore_index=-1)\n    \n    # 1. Gemma-style softcapping: C * tanh(logits / C)\n    softcap = 30.0\n    capped_logits = softcap * torch.tanh(flat_logits / softcap)\n    \n    # 2. Cross-entropy with label smoothing\n    ce_loss = F.cross_entropy(capped_logits, flat_targets, ignore_index=-1, label_smoothing=0.03)\n    \n    # 3. Z-loss: penalize large log-partition values\n    mask = flat_targets != -1\n    log_z = torch.logsumexp(capped_logits, dim=-1)\n    z_loss = ((log_z ** 2) * mask.float()).sum() / mask.sum()\n    \n    return ce_loss + 1e-4 * z_loss",
3  "end_line": 220,
4  "filename": "nanoGPT/custom_pretrain.py",
5  "start_line": 188,
6  "op": "replace"
7}

1{
2  "end_line": 214,
3  "filename": "nanoGPT/custom_pretrain.py",
4  "op": "replace",
5  "start_line": 188,
6  "content": "def compute_loss(logits, targets):\n    \"\"\"Symmetric cross-entropy (forward KL + backward KL) with softcapping.\n    \n    Symmetric CE combines standard CE (forward KL) with reverse KL (penalizing\n    the model for spreading probability mass away from the target). This creates\n    a tighter optimization landscape. Label smoothing is applied to the forward\n    direction for calibration. Softcapping bounds logit magnitudes.\n    During evaluation, standard cross-entropy is returned for comparable metrics.\n    \"\"\"\n    flat_logits = logits.view(-1, logits.size(-1)).float()\n    flat_targets = targets.view(-1)\n    mask = flat_targets != -1\n    \n    if not torch.is_grad_enabled():\n        return F.cross_entropy(flat_logits, flat_targets, ignore_index=-1)\n    \n    # Softcapping for stability\n    softcap = 30.0\n    capped_logits = softcap * torch.tanh(flat_logits / softcap)\n    \n    # Forward KL: standard cross-entropy with label smoothing\n    forward_ce = F.cross_entropy(capped_logits, flat_targets, ignore_index=-1, label_smoothing=0.05)\n    \n    # Backward KL: -(1/V) * sum_k log p(k|x) where V is vocab size\n    # This penalizes spreading probability mass too thinly\n    log_probs = F.log_softmax(capped_logits, dim=-1)\n    mask_expanded = mask.float().unsqueeze(1)  # (N, 1)\n    \n    # Average log prob across vocab for masked positions\n    avg_log_prob = (log_probs * mask_expanded).sum(dim=-1) / mask_expanded.sum()\n    reverse_kl = -(avg_log_prob * mask.float()).sum() / mask.sum()\n    \n    # Z-loss for stability\n    log_z = torch.logsumexp(capped_logits, dim=-1)\n    z_loss = ((log_z ** 2) * mask.float()).sum() / mask.sum()\n    \n    # Combine: forward CE dominates, reverse KL regularizes\n    return forward_ce + 0.02 * reverse_kl + 1e-4 * z_loss"
7}

1{
2  "content": "def compute_loss(logits, targets):\n    \"\"\"Focal softcapped cross-entropy with z-loss.\n    \n    Combines three techniques for improved pretraining:\n    1. Softcapping bounds logits for stable optimization (Gemma-style).\n    2. Focal loss (gamma=0.5) focuses on harder tokens.\n    3. Z-loss prevents logit explosion.\n    \n    Lower gamma=0.5 (vs typical 2.0 for detection) since LM tokens are \n    already quite hard - we want gentle reweighting, not aggressive focusing.\n    During evaluation, standard cross-entropy is returned for comparable metrics.\n    \"\"\"\n    flat_logits = logits.view(-1, logits.size(-1)).float()\n    flat_targets = targets.view(-1)\n    mask = flat_targets != -1\n    \n    if not torch.is_grad_enabled():\n        return F.cross_entropy(flat_logits, flat_targets, ignore_index=-1)\n    \n    # Softcapping for stability\n    softcap = 30.0\n    capped_logits = softcap * torch.tanh(flat_logits / softcap)\n    \n    # Log probabilities\n    log_probs = F.log_softmax(capped_logits, dim=-1)\n    \n    # Focal loss: -alpha * (1 - p_t)^gamma * log(p_t)\n    gamma = 0.5\n    target_indices = flat_targets.clamp(min=0)\n    target_log_probs = log_probs.gather(1, target_indices.unsqueeze(1)).squeeze(1)\n    \n    # p_t for each token\n    p_t = target_log_probs.exp()\n    \n    # Focal weight\n    focal_weight = (1 - p_t) ** gamma\n    \n    # Per-token focal loss\n    per_token_loss = focal_weight * (-target_log_probs)\n    \n    # Average over valid tokens\n    focal_loss = (per_token_loss * mask.float()).sum() / mask.sum()\n    \n    # Z-loss for stability\n    log_z = torch.logsumexp(capped_logits, dim=-1)\n    z_loss = ((log_z ** 2) * mask.float()).sum() / mask.sum()\n    \n    return focal_loss + 1e-4 * z_loss",
3  "end_line": 225,
4  "filename": "nanoGPT/custom_pretrain.py",
5  "start_line": 188,
6  "op": "replace"
7}

1{
2  "content": "def compute_loss(logits, targets):\n    \"\"\"Focal-softcap cross-entropy with label smoothing and z-loss.\n    \n    Combines four complementary techniques:\n    1. Softcapping bounds logits for stable optimization (Gemma-style).\n    2. Label smoothing prevents overconfidence and improves calibration.\n    3. Focal loss (gamma=0.75) focuses optimization on harder tokens.\n    4. Z-loss prevents logit explosion.\n    \n    Gamma=0.75 provides moderate reweighting: enough to focus on hard tokens\n    without destabilizing training. During evaluation, standard cross-entropy\n    is returned for comparable metrics.\n    \"\"\"\n    flat_logits = logits.view(-1, logits.size(-1)).float()\n    flat_targets = targets.view(-1)\n    mask = flat_targets != -1\n    n_valid = mask.sum()\n    \n    if not torch.is_grad_enabled():\n        return F.cross_entropy(flat_logits, flat_targets, ignore_index=-1)\n    \n    # 1. Softcapping for stability\n    softcap = 30.0\n    capped_logits = softcap * torch.tanh(flat_logits / softcap)\n    \n    # 2. Label smoothing with focal loss\n    eps = 0.05  # smoothing factor\n    gamma = 0.75  # focal focusing parameter\n    V = capped_logits.size(-1)  # vocab size\n    \n    log_probs = F.log_softmax(capped_logits, dim=-1)\n    probs = log_probs.exp()\n    \n    # Smoothed target distribution\n    smooth_targets = torch.full_like(log_probs, eps / (V - 1))\n    target_indices = flat_targets.clamp(min=0)\n    smooth_targets.scatter_(1, target_indices.unsqueeze(1), 1.0 - eps)\n    \n    # Focal weighting: use smoothed probability of correct class\n    smooth_probs = (smooth_targets * probs).sum(dim=-1)  # p_smooth for each token\n    \n    # Focal loss with smoothing: -sum_t q(t) * (1-p(t))^gamma * log p(t)\n    focal_weight = (1.0 - probs.clamp(min=1e-7, max=1.0 - 1e-7)) ** gamma\n    \n    # Weighted negative log likelihood with smoothed targets\n    per_token_loss = -(smooth_targets * focal_weight * log_probs).sum(dim=-1)\n    \n    focal_loss = (per_token_loss * mask.float()).sum() / n_valid\n    \n    # 3. Z-loss for stability\n    log_z = torch.logsumexp(capped_logits, dim=-1)\n    z_loss = ((log_z ** 2) * mask.float()).sum() / n_valid\n    \n    return focal_loss + 1e-4 * z_loss",
3  "end_line": 235,
4  "filename": "nanoGPT/custom_pretrain.py",
5  "start_line": 188,
6  "op": "replace"
7}