File size: 3,587 Bytes
a311e94 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | """Regularisation algorithms: label smoothing, mixup, weight-decay schedules."""
from __future__ import annotations
import math
from typing import Any, Dict, Optional
import torch
import torch.nn.functional as F
from .base import AlgorithmContext, BaseAlgorithm, compose_loss
class LabelSmoothing(BaseAlgorithm):
"""Recompute the causal-LM loss with label smoothing."""
name = "label_smoothing"
category = "regularisation"
description = "Softens one-hot targets to reduce over-confidence."
def __init__(self, epsilon: float = 0.1):
super().__init__(epsilon=epsilon)
self.epsilon = epsilon
def compute_loss(self, ctx, batch, outputs, loss):
labels = batch.get("labels")
if labels is None:
return loss
logits = outputs["logits"] if isinstance(outputs, dict) else outputs.logits
logits = logits[..., :-1, :].reshape(-1, logits.shape[-1]).float()
labels = labels[..., 1:].reshape(-1)
return F.cross_entropy(logits, labels, label_smoothing=self.epsilon)
class Mixup(BaseAlgorithm):
"""Embedding-space mixup for language models.
Mixes input embeddings between random pairs and interpolates the
corresponding cross-entropy terms.
"""
name = "mixup"
category = "regularisation"
description = "Interpolates input embeddings between random pairs."
def __init__(self, alpha: float = 0.2):
super().__init__(alpha=alpha)
self.alpha = alpha
def train_step(self, ctx, batch):
batch = self.on_step_begin(ctx, batch)
input_ids = batch.get("input_ids")
labels = batch.get("labels")
if input_ids is None or labels is None or self.alpha <= 0:
outputs = ctx.model(**batch)
loss = outputs["loss"] if isinstance(outputs, dict) else outputs.loss
loss = self.compute_loss(ctx, batch, outputs, loss)
loss = compose_loss(ctx, batch, outputs, loss, skip=self)
loss.backward()
return loss
lam = float(torch.distributions.Beta(self.alpha, self.alpha).sample().item())
perm = torch.randperm(input_ids.size(0), device=input_ids.device)
embeddings = ctx.model.get_input_embeddings()(input_ids)
mixed = lam * embeddings + (1.0 - lam) * embeddings[perm]
outputs = ctx.model(inputs_embeds=mixed, attention_mask=batch.get("attention_mask"))
logits = outputs["logits"] if isinstance(outputs, dict) else outputs.logits
logits = logits[..., :-1, :].reshape(-1, logits.shape[-1]).float()
loss_a = F.cross_entropy(logits, labels[..., 1:].reshape(-1))
loss_b = F.cross_entropy(logits, labels[perm][..., 1:].reshape(-1))
loss = lam * loss_a + (1.0 - lam) * loss_b
loss.backward()
return loss
class WeightDecaySchedule(BaseAlgorithm):
"""Ramp the optimiser weight decay from ``start`` to ``end`` over training."""
name = "weight_decay_schedule"
category = "regularisation"
description = "Linearly ramps weight decay across the run."
def __init__(self, start: float = 0.0, end: float = 0.1):
super().__init__(start=start, end=end)
self.start = start
self.end = end
def on_step_begin(self, ctx, batch):
if ctx.optimizer is None or ctx.max_steps <= 0:
return batch
t = min(1.0, ctx.step / max(1, ctx.max_steps - 1))
value = self.start + (self.end - self.start) * t
for group in ctx.optimizer.param_groups:
group["weight_decay"] = value
return batch
|