File size: 6,894 Bytes
51bb0b7 | 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 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 | """Test-time ensembling for Modular-TS (v3): context-length ensembles, temporal-hierarchy
(multi-resolution) reconciliation and in-context backtest weighting. Inference only.
Members (spec string, comma separated):
cN - the base forecaster with max_context N (e.g. c8192, c2048, c512)
dK - downsampled member: context mean-pooled by K (aligned to the forecast origin, so it
sees up to K x the history at the same token budget), forecast ceil(H/K) coarse
steps, step-upsampled back. Used as a *level corrector* for the fine forecasts
(temporal hierarchies / THieF, Athanasopoulos et al. 2017): fine + a*(coarse - pool(fine))
instead of being averaged in directly, so intra-period shape comes from full res.
Combination of the c-members: quantile averaging (Vincentisation), uniform or, with
backtest=True, per-series softmax weights from each member's scaled MAE on the last H
observed steps of the context (a zero-shot, in-context model selection - it only uses the
series' own history).
fc = EnsembleForecaster(TimestralForecaster.from_pretrained(exp), "c8192,c2048,d4", alpha=0.5)
"""
from __future__ import annotations
import math
from dataclasses import dataclass
import numpy as np
import torch
@dataclass
class _F:
quantiles: torch.Tensor
def _pool(x: np.ndarray | None, k: int) -> np.ndarray | None:
"""(B,V,T) NaN-aware mean pool by k aligned to the end."""
if x is None:
return None
B, V, T = x.shape
n = T // k
if n == 0:
return None
y = x[..., T - n * k:].reshape(B, V, n, k)
with np.errstate(all="ignore"):
cnt = np.isfinite(y).sum(-1)
s = np.nansum(y, -1)
out = np.where(cnt > 0, s / np.maximum(cnt, 1), np.nan)
return out.astype(np.float32)
class EnsembleForecaster:
def __init__(self, fc, spec: str = "c8192", alpha: float = 0.5, backtest: bool = False, tau: float = 0.5, bt_windows: int = 1, prior: float = 0.0,
min_coarse_ctx: int = 256, min_extra: float = 1.5, min_ratio: float = 8.0):
self.fc = fc
self.members = [m.strip() for m in spec.split(",") if m.strip()]
cl = fc.model.config.context_length # never exceed the trained context (RoPE does not extrapolate)
self.ctxs = [min(int(m[1:]), cl) for m in self.members if m[0] == "c"] or [fc.max_context]
self.ks = [int(m[1:]) for m in self.members if m[0] == "d"]
self.alpha, self.backtest, self.tau = alpha, backtest, tau
self.bt_windows, self.prior = bt_windows, prior # prior: logit bonus for the first (longest) member
self.min_coarse_ctx, self.min_extra, self.min_ratio = min_coarse_ctx, min_extra, min_ratio # member ctx >= min_ratio*H
self.base_ctx = fc.max_context
self.max_context = max(self.ctxs + [self.base_ctx * k for k in self.ks] + [self.base_ctx])
self.quantile_levels = fc.quantile_levels
self.mi = list(fc.quantile_levels).index(0.5) if 0.5 in list(fc.quantile_levels) else len(fc.quantile_levels) // 2
def _one(self, ctx, h, past, mc):
self.fc.max_context = mc
try:
return self.fc.forecast(ctx, h, past_covariates=past).quantiles.float().cpu()
finally:
self.fc.max_context = self.base_ctx
def _members(self, ctx, h, past):
T = ctx.shape[-1]
outs, used = [], set()
for i, c in enumerate(self.ctxs):
if i > 0 and c < self.min_ratio * h: # a short context cannot support a long horizon
continue
eff = min(c, T)
if eff in used: # a shorter series makes several members identical
continue
used.add(eff)
outs.append(self._one(ctx[..., -c:], h, None if past is None else past[..., -c:], c))
return outs
def forecast(self, context, prediction_length, past_covariates=None):
ctx = np.asarray(context, np.float32)
if ctx.ndim == 2:
ctx = ctx[:, None]
past = None if past_covariates is None else np.asarray(past_covariates, np.float32)
h = prediction_length
T = ctx.shape[-1]
outs = self._members(ctx, h, past)
if len(outs) == 1:
q = outs[0]
elif self.backtest and T - h >= 2 * 32:
E, ok = 0.0, True
nw = 0
for wdw in range(1, self.bt_windows + 1):
if T - wdw * h < 2 * 32:
break
cut = T - wdw * h
bt = self._members(ctx[..., :cut], h, None if past is None else past[..., :cut])
if len(bt) != len(outs):
ok = False
break
y = torch.from_numpy(ctx[..., cut:cut + h])
hist = torch.from_numpy(ctx[..., :cut])
dd = hist[..., 1:] - hist[..., :-1]
fin = torch.isfinite(dd)
scale = (torch.where(fin, dd.abs(), torch.zeros_like(dd)).sum(-1) / fin.sum(-1).clamp_min(1)).clamp_min(1e-8)
errs = []
for b in bt:
e = (b[..., self.mi] - y).abs()
fe = torch.isfinite(e)
e = torch.where(fe, e, torch.zeros_like(e)).sum(-1) / fe.sum(-1).clamp_min(1) / scale
errs.append(e.mean(-1))
E = E + torch.stack(errs, 0)
nw += 1
if not ok or nw == 0:
q = torch.stack(outs).mean(0)
else:
E = E / nw
logit = -E / (self.tau * E.min(0, keepdim=True).values.clamp_min(1e-3))
logit[0] = logit[0] + self.prior
w = torch.softmax(logit, 0)
q = (torch.stack(outs) * w[:, :, None, None, None]).sum(0)
else:
q = torch.stack(outs).mean(0)
# temporal-hierarchy level correction from downsampled members
for k in self.ks:
hk = math.ceil(h / k)
pooled = _pool(ctx, k)
if pooled is None or pooled.shape[-1] < self.min_coarse_ctx:
continue
# only worth it when the coarse view sees materially more history than the fine one
if T < self.min_extra * self.base_ctx and hk < 2:
continue
pp = _pool(past, k) if past is not None else None
cq = self._one(pooled[..., -self.base_ctx:], hk, None if pp is None else pp[..., -self.base_ctx:], self.base_ctx)
cmed = cq[..., self.mi] # (B,V,hk)
fmed = q[..., self.mi]
pad = hk * k - h
fm = torch.cat([fmed, fmed[..., -1:].expand(*fmed.shape[:-1], pad)], -1) if pad else fmed
fpool = fm.view(*fm.shape[:-1], hk, k).mean(-1)
delta = (cmed - fpool).repeat_interleave(k, -1)[..., :h]
q = q + self.alpha / len(self.ks) * delta[..., None]
return _F(q)
|