Spaces:
Sleeping
Sleeping
File size: 2,729 Bytes
8fca640 | 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 | """Alternative point estimators for axial categorical distributions."""
import torch
def estimates(logits, names):
probability = logits.float().softmax(-2)
height = logits.shape[-2]
grid = torch.arange(height, device=logits.device, dtype=torch.float32)[
None, None, :, None
]
result = {}
if "mean" in names:
result["mean"] = (probability * grid).sum(-2)
if "median" in names:
cumulative = probability.cumsum(-2)
index = (cumulative < 0.5).sum(-2).clamp_max(height - 1)
mass = probability.gather(-2, index.unsqueeze(-2)).squeeze(-2)
before = cumulative.gather(-2, (index - 1).clamp_min(0).unsqueeze(-2)).squeeze(
-2
)
before = torch.where(index == 0, torch.zeros_like(before), before)
result["median"] = (
index.float() - 0.5 + (0.5 - before) / mass.clamp_min(1e-12)
).clamp(0, height - 1)
if "quadratic_mode" in names:
index = logits.argmax(-2)
score = logits.float()
center = score.gather(-2, index.unsqueeze(-2)).squeeze(-2)
left = score.gather(-2, (index - 1).clamp_min(0).unsqueeze(-2)).squeeze(-2)
right = score.gather(
-2, (index + 1).clamp_max(height - 1).unsqueeze(-2)
).squeeze(-2)
curvature = left - 2 * center + right
offset = torch.where(
(curvature < -1e-8) & (index > 0) & (index < height - 1),
0.5 * (left - right) / curvature.clamp_max(-1e-8),
torch.zeros_like(center),
).clamp(-0.5, 0.5)
result["quadratic_mode"] = (index.float() + offset).clamp(0, height - 1)
if set(result) != set(names):
raise ValueError("Unknown posterior estimator")
return result
def positions_from_logits(logits, config, device=None):
by_device = config.get("inference", {}).get("posterior_by_device")
if by_device is not None:
if device is None or device not in by_device:
raise ValueError("Configured head estimators require the device")
keys = config["device_keys"][device]
mapping = by_device[device]
if set(mapping) != set(keys) or logits.shape[1] != len(keys):
raise ValueError("Configured head estimators differ from output identities")
names = [mapping[key] for key in keys]
values = estimates(logits, sorted(set(names)))
return torch.stack(
[values[name][:, index] for index, name in enumerate(names)], dim=1
)
name = config.get("inference", {}).get("posterior", "mean")
if name not in ["mean", "median", "quadratic_mode"]:
raise ValueError("Unsupported configured posterior")
return estimates(logits, [name])[name]
|