The Math Reference Sheet¶

import torch
import torch.nn.functional as F

def safe_softmax(z: torch.Tensor, dim: int = -1) -> torch.Tensor:
    """
    Numerically stable softmax.
    Subtracting the max is equivalent — the shift cancels in numerator
    and denominator — but prevents exp() overflow when z is large.

    z: any shape, we reduce along `dim`.
    """
    z_shifted = z - z.max(dim=dim, keepdim=True).values  # broadcast-safe
    exp_z = torch.exp(z_shifted)
    return exp_z / exp_z.sum(dim=dim, keepdim=True)

# Quick sanity check
logits = torch.tensor([2.0, 1.0, 0.1])
probs  = safe_softmax(logits)
assert torch.allclose(probs.sum(), torch.tensor(1.0))
print(probs)  # tensor([0.6590, 0.2424, 0.0986])

import torch
import torch.nn as nn

# Internally: F.log_softmax + F.nll_loss, numerically fused
criterion = nn.CrossEntropyLoss(reduction='mean')

# logits: [batch, vocab];  targets: [batch] of integer token ids
logits  = torch.randn(4, 32_000)   # e.g. Llama-3 vocab
targets = torch.randint(0, 32_000, (4,))

loss = criterion(logits, targets)
perplexity = torch.exp(loss)
print(f"loss={loss:.4f}  PPL={perplexity:.1f}")

import torch, math

def rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
    """
    x     : (..., d_model)
    weight: (d_model,) — the learnable scale γ
    """
    # Mean of squares along the last dimension
    rms = x.pow(2).mean(dim=-1, keepdim=True).add(eps).sqrt()
    return (x / rms) * weight

# Test against PyTorch's built-in
d = 512
x = torch.randn(2, 16, d)
w = torch.ones(d)

out_custom = rms_norm(x, w)
out_builtin = torch.nn.functional.rms_norm(x, (d,), w, eps=1e-6)
assert torch.allclose(out_custom, out_builtin, atol=1e-5)

import math
import torch
import torch.nn.functional as F

def scaled_dot_product_attention(
    q: torch.Tensor,  # (B, H, T, d_k)
    k: torch.Tensor,  # (B, H, T, d_k)
    v: torch.Tensor,  # (B, H, T, d_v)
    causal: bool = True,
) -> torch.Tensor:
    """
    Pure-PyTorch reference implementation.
    In production use F.scaled_dot_product_attention (dispatches to FlashAttention).
    """
    d_k = q.size(-1)
    # (B, H, T_q, T_k)
    scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)

    if causal:
        T = scores.size(-1)
        mask = torch.triu(torch.ones(T, T, device=q.device), diagonal=1).bool()
        scores = scores.masked_fill(mask, float('-inf'))

    attn = F.softmax(scores, dim=-1)      # (B, H, T_q, T_k)
    return torch.matmul(attn, v)          # (B, H, T_q, d_v)

import torch

def apply_rope(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
    """
    x         : (B, T, H, d_head)
    freqs_cis : (T, d_head/2) complex tensor = cos + i*sin

    Applies RoPE by treating consecutive pairs of dimensions as complex numbers.
    """
    # Reshape to complex: (B, T, H, d_head/2) complex64
    x_complex = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
    # freqs_cis: (1, T, 1, d_head/2) after unsqueeze for broadcast
    freqs_cis = freqs_cis.unsqueeze(0).unsqueeze(2)
    # Element-wise complex multiply = rotation
    x_rotated = x_complex * freqs_cis
    return torch.view_as_real(x_rotated).flatten(-2).type_as(x)

def build_rope_freqs(seq_len: int, d_head: int, base: float = 10_000.0) -> torch.Tensor:
    """Returns (seq_len, d_head/2) complex tensor of e^{i*m*theta_i}."""
    i  = torch.arange(0, d_head, 2).float()
    theta = 1.0 / (base ** (i / d_head))                # (d_head/2,)
    m  = torch.arange(seq_len).float()                  # (T,)
    freqs = torch.outer(m, theta)                        # (T, d_head/2)
    return torch.polar(torch.ones_like(freqs), freqs)    # complex

import torch

def adamw_step(
    param:  torch.Tensor,   # parameter θ
    grad:   torch.Tensor,   # gradient g_t
    m:      torch.Tensor,   # first moment (in-place)
    v:      torch.Tensor,   # second moment (in-place)
    t:      int,            # step count (1-indexed)
    lr:     float = 3e-4,
    beta1:  float = 0.9,
    beta2:  float = 0.95,
    eps:    float = 1e-8,
    wd:     float = 0.1,
) -> None:
    """Single AdamW update step. Modifies param, m, v in-place."""
    # Bias-correct moments
    m.mul_(beta1).add_(grad, alpha=1 - beta1)
    v.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)

    m_hat = m / (1 - beta1 ** t)
    v_hat = v / (1 - beta2 ** t)

    # Weight decay applied directly to param, then Adam step
    param.mul_(1 - lr * wd)
    param.addcdiv_(m_hat, v_hat.sqrt().add_(eps), value=-lr)

import torch
import torch.nn as nn
import math

class LoRALinear(nn.Module):
    """
    Drop-in replacement for nn.Linear with LoRA.
    The base weight W0 is frozen; only A and B are trained.
    """
    def __init__(
        self,
        in_features: int,
        out_features: int,
        r: int = 16,
        lora_alpha: float = 16.0,
        bias: bool = False,
    ):
        super().__init__()
        self.r = r
        self.scaling = lora_alpha / r

        # Frozen pretrained weight
        self.weight = nn.Parameter(
            torch.empty(out_features, in_features), requires_grad=False
        )
        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))

        if bias:
            self.bias = nn.Parameter(torch.zeros(out_features))
        else:
            self.register_parameter("bias", None)

        # Trainable LoRA matrices
        self.lora_A = nn.Parameter(torch.empty(r, in_features))
        self.lora_B = nn.Parameter(torch.zeros(out_features, r))
        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        # B stays zero → ΔW = 0 at init

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Base forward + LoRA delta, scaled
        base_out  = nn.functional.linear(x, self.weight, self.bias)
        lora_out  = nn.functional.linear(x, self.lora_A)        # (..., r)
        lora_out  = nn.functional.linear(lora_out, self.lora_B) # (..., out)
        return base_out + self.scaling * lora_out

    def merge(self) -> nn.Linear:
        """Collapse LoRA into base weight for zero-overhead inference."""
        merged_w = self.weight + self.scaling * (self.lora_B @ self.lora_A)
        lin = nn.Linear(
            self.weight.size(1), self.weight.size(0),
            bias=self.bias is not None
        )
        lin.weight = nn.Parameter(merged_w)
        if self.bias is not None:
            lin.bias = nn.Parameter(self.bias.clone())
        return lin

import torch

def ppo_loss(
    log_probs_new:  torch.Tensor,  # (B,) log π_θ(a|s)
    log_probs_old:  torch.Tensor,  # (B,) log π_old(a|s)
    advantages:     torch.Tensor,  # (B,) advantage estimates A_t
    clip_eps:       float = 0.2,
    reduce:         bool  = True,
) -> torch.Tensor:
    """
    PPO clipped surrogate loss.
    We *maximise* this, so negate before calling .backward().
    """
    ratio  = torch.exp(log_probs_new - log_probs_old)  # importance ratio
    # Clipped objective
    surr1  = ratio * advantages
    surr2  = ratio.clamp(1 - clip_eps, 1 + clip_eps) * advantages
    loss   = torch.min(surr1, surr2)  # take the pessimistic bound
    return loss.mean() if reduce else loss

import torch

def grpo_advantages(rewards: torch.Tensor) -> torch.Tensor:
    """
    rewards : (G,) — one scalar reward per response in the group
    Returns : (G,) normalised advantages
    """
    mean = rewards.mean()
    std  = rewards.std(unbiased=False).clamp(min=1e-8)
    return (rewards - mean) / std

# Example: 8 responses with rewards from a verifier
rewards = torch.tensor([1.0, 0.0, 1.0, 0.5, 0.0, 1.0, 0.5, 0.0])
adv = grpo_advantages(rewards)
print(adv.round(decimals=2))
# tensor([ 1.15, -1.15,  1.15,  0.00, -1.15,  1.15,  0.00, -1.15])

import torch
import torch.nn.functional as F

def dpo_loss(
    log_probs_theta_w:  torch.Tensor,  # (B,) log π_θ(y_w | x)
    log_probs_theta_l:  torch.Tensor,  # (B,) log π_θ(y_l | x)
    log_probs_ref_w:    torch.Tensor,  # (B,) log π_ref(y_w | x)
    log_probs_ref_l:    torch.Tensor,  # (B,) log π_ref(y_l | x)
    beta:               float = 0.1,
) -> torch.Tensor:
    """
    Standard DPO loss (Rafailov et al. 2023).
    Each log_prob is the per-sequence sum of token log-probs.
    """
    log_ratio_w = log_probs_theta_w - log_probs_ref_w  # h_θ(y_w, x)
    log_ratio_l = log_probs_theta_l - log_probs_ref_l  # h_θ(y_l, x)
    # Positive when preferred response has higher relative log-prob
    gap  = beta * (log_ratio_w - log_ratio_l)
    loss = -F.logsigmoid(gap)  # = log(1 + exp(-gap)), numerically stable
    return loss.mean()

import torch

def kl_estimators(
    log_pi:     torch.Tensor,  # (B, T) log-probs under current policy
    log_pi_ref: torch.Tensor,  # (B, T) log-probs under reference policy
) -> dict:
    """
    Returns three scalar KL estimates averaged over batch and token positions.
    """
    log_r = log_pi - log_pi_ref  # log importance ratio  (B, T)
    r     = log_r.exp()

    k1 = log_r                   # naive, can be negative
    k2 = r - 1.0 - log_r        # always >= 0 by log inequality
    k3 = (r - 1.0).pow(2) / 2   # Taylor approximation

    return {
        "k1": k1.mean().item(),
        "k2": k2.mean().item(),
        "k3": k3.mean().item(),
    }

# Sanity: for identical policies KL should be ~0
lp = torch.full((4, 512), -3.0)
print(kl_estimators(lp, lp))
# {'k1': 0.0, 'k2': 0.0, 'k3': 0.0}