File size: 4,461 Bytes
e68eb1c | 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 | """Causal sequence models for the Wunder stepwise next-state task.
Design constraints derived from the data analysis (see memory / report):
- Features are pre-whitened (~N(0,1), no heavy tails) so no internal
normalization is required; the model predicts the next *level* directly.
- Inference is online, one row at a time, on a single CPU core within a
60-minute budget. A unidirectional GRU carries hidden state across steps,
giving O(1) work per prediction.
- Training uses full 1000-step sequences with back-prop-through-time and a
loss masked to the scored steps (100..998). Because a GRU is a pure
recurrence, the full-sequence forward and the stepwise stateful forward
produce identical outputs — so what we train is exactly what we score.
`forward(x, h0)` returns per-step predictions of the *next* state:
preds[:, t] predicts x[:, t + 1].
"""
from __future__ import annotations
from typing import Optional, Tuple
import torch
import torch.nn as nn
class CausalGRUForecaster(nn.Module):
"""Unidirectional GRU that predicts the next state at every step.
Args:
n_features: input/output width (32).
d_model: GRU hidden size and embedding width.
n_layers: stacked GRU layers.
dropout: dropout between GRU layers and inside the head.
head_hidden: width of the MLP head hidden layer (defaults to d_model).
"""
def __init__(
self,
n_features: int = 32,
d_model: int = 256,
n_layers: int = 2,
dropout: float = 0.1,
head_hidden: Optional[int] = None,
rnn_type: str = "gru",
):
super().__init__()
self.n_features = int(n_features)
self.d_model = int(d_model)
self.n_layers = int(n_layers)
self.rnn_type = str(rnn_type).lower()
head_hidden = int(head_hidden) if head_hidden else self.d_model
# Per-step input embedding (applied identically in batch and stepwise modes).
self.input_proj = nn.Linear(n_features, d_model)
self.input_norm = nn.LayerNorm(d_model)
rnn_cls = nn.LSTM if self.rnn_type == "lstm" else nn.GRU
self.gru = rnn_cls(
input_size=d_model,
hidden_size=d_model,
num_layers=n_layers,
batch_first=True,
dropout=dropout if n_layers > 1 else 0.0,
)
self.head = nn.Sequential(
nn.LayerNorm(d_model),
nn.Linear(d_model, head_hidden),
nn.GELU(),
nn.Dropout(dropout),
nn.Linear(head_hidden, n_features),
)
def forward(
self,
x: torch.Tensor,
h0: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Args:
x: (batch, seq_len, n_features)
h0: optional initial hidden state (n_layers, batch, d_model)
Returns:
preds: (batch, seq_len, n_features) where preds[:, t] predicts x[:, t+1]
h_n: final hidden state (n_layers, batch, d_model)
"""
z = self.input_norm(self.input_proj(x))
out, h_n = self.gru(z, h0)
preds = self.head(out)
return preds, h_n
@torch.no_grad()
def step(
self,
state: torch.Tensor,
h: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Single online step.
Args:
state: (n_features,) or (1, n_features) current state.
h: previous hidden state (n_layers, 1, d_model) or None.
Returns:
pred: (n_features,) prediction of the next state.
h_new: updated hidden state.
"""
if state.dim() == 1:
state = state.view(1, 1, -1)
elif state.dim() == 2:
state = state.unsqueeze(1)
preds, h_new = self.forward(state, h)
return preds[0, -1], h_new
def build_gru(config: dict) -> CausalGRUForecaster:
"""Construct a CausalGRUForecaster from a plain dict config."""
return CausalGRUForecaster(
n_features=config.get("n_features", 32),
d_model=config.get("d_model", 256),
n_layers=config.get("n_layers", 2),
dropout=config.get("dropout", 0.1),
head_hidden=config.get("head_hidden"),
rnn_type=config.get("rnn_type", "gru"),
)
def count_parameters(model: nn.Module) -> int:
return sum(p.numel() for p in model.parameters() if p.requires_grad)
|