| """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 |
|
|
|
|
| |
|
|
| 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) |
|
|
|
|
| |
|
|
