Njap-png
Add Agron neural-network zoo, optimisers, algorithms, schedulers and unified trainer
a311e94 | """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 | |