File size: 2,570 Bytes
30b1727 | 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 | """Calibration head used at inference (temperature as a function of the logits and the pooled answer-slot state)."""
import math
from typing import Optional
import torch
import torch.nn as nn
import torch.nn.functional as F
LO, HI = math.log(0.5), math.log(5.0)
MODES = {"yesno": 0, "digits": 1, "letters": 2, "parallel": 3}
N_FEAT = 3 + 4 + 5
# ------------------------------------------------------------------ features -----------------------------------------------------------
def make_features(logits: torch.Tensor, mask: torch.Tensor, qtype: torch.Tensor, mode: torch.Tensor) -> torch.Tensor:
"""(N, Kmax) padded logits -> (N, N_FEAT) features that do not depend on which option is which."""
z = logits.masked_fill(~mask, -1e4)
k = mask.sum(-1).clamp(min=2).float()
logp = torch.log_softmax(z, -1)
p = logp.exp()
ent = -(p * logp).sum(-1) / torch.log(k)
top = z.topk(2, -1).values
mean = (logits * mask).sum(-1) / k
std = (((logits - mean[:, None]) ** 2 * mask).sum(-1) / k).sqrt()
return torch.cat([F.one_hot(qtype, 3).float(), F.one_hot(mode, 4).float(), torch.stack([torch.log(k) / 5, ent, p.max(-1).values, (top[:, 0] - top[:, 1]) / 10, std / 10], -1)], -1)
class SwiGLU(nn.Module):
def __init__(self, d_in: int, d_hid: int, d_out: int):
super().__init__()
self.gate, self.up, self.down = nn.Linear(d_in, d_hid), nn.Linear(d_in, d_hid), nn.Linear(d_hid, d_out)
def forward(self, x):
return self.down(F.silu(self.gate(x)) * self.up(x))
class TempHead(nn.Module):
"""log T = LO + (HI - LO) * sigmoid(raw + c) with c chosen so that raw = 0 gives T = 1; the last layer starts at zero (T = 1)."""
def __init__(self, n_feat: int = N_FEAT, hidden: Optional[int] = None, d: int = 128):
super().__init__()
self.hidden_dim = hidden
self.proj = nn.Sequential(nn.LayerNorm(hidden), nn.Linear(hidden, d)) if hidden else None
self.ffn = SwiGLU(n_feat + (d if hidden else 0), 2 * d, 1)
nn.init.zeros_(self.ffn.down.weight); nn.init.zeros_(self.ffn.down.bias)
self.register_buffer("c", torch.tensor(math.log((0 - LO) / (HI - 0))))
def log_t(self, feats: torch.Tensor, pooled: Optional[torch.Tensor] = None) -> torch.Tensor:
x = feats if self.proj is None else torch.cat([feats, self.proj(pooled)], -1)
return LO + (HI - LO) * torch.sigmoid(self.ffn(x).squeeze(-1) + self.c)
# ------------------------------------------------------------------ cache --------------------------------------------------------------
|