Create TART.py

TART.py

@@ -0,0 +1,453 @@
+import math
+from typing import Optional
+
+import numpy as np
+import jax
+import jax.numpy as jnp
+import flax.linen as nn
+import optax
+from flax.training import train_state
+
+def zscore(x: np.ndarray, eps: float = 1e-6):
+    x = np.asarray(x, dtype=np.float32)
+    mu = x.mean(axis=0, keepdims=True)
+    sd = x.std(axis=0, keepdims=True) + eps
+    return (x - mu) / sd, (mu, sd)
+
+
+def make_nextstep_windows(X: np.ndarray, L: int):
+    """
+    Simple 1-step-ahead dataset:
+      input  X_in[i]  = X[i : i+L]      (L, D)
+      target Y_out[i] = X[i+L]          (D,)
+    """
+    X = np.asarray(X, dtype=np.float32)
+    T, D = X.shape
+    if T <= L:
+        raise ValueError(f"T={T} must be > L={L}")
+    N = T - L
+    X_in = np.stack([X[i:i+L] for i in range(N)], axis=0)   # (N, L, D)
+    Y_out = np.stack([X[i+L] for i in range(N)], axis=0)    # (N, D)
+    return X_in, Y_out
+
+
+def choose_heads(d_model: int, target_head_dim: int = 32, max_heads: int = 8) -> int:
+    """Pick a reasonable number of heads given d_model."""
+    n_heads = max(1, min(max_heads, d_model // target_head_dim))
+    # ensure divisibility
+    while d_model % n_heads != 0 and n_heads > 1:
+        n_heads -= 1
+    return n_heads
+
+def build_rotary_inv_freq(head_dim: int, base: float = 10000.0):
+    assert head_dim % 2 == 0, "head_dim must be even for RoPE."
+    return 1.0 / (base ** (jnp.arange(0, head_dim, 2, dtype=jnp.float32) / head_dim))
+
+
+def apply_rope(x, inv_freq):
+    """
+    Apply RoPE to q/k.
+
+    x:        (batch, heads, seq_len, head_dim)
+    inv_freq: (head_dim/2,)
+    """
+    b, h, t, d = x.shape
+    half = d // 2
+    # positions
+    positions = jnp.arange(t, dtype=jnp.float32)     # (t,)
+    freqs = jnp.einsum("i,j->ij", positions, inv_freq)  # (t, half)
+    cos = jnp.cos(freqs)[None, None, :, :]  # (1,1,t,half)
+    sin = jnp.sin(freqs)[None, None, :, :]
+
+    x1 = x[..., :half]   # (b,h,t,half)
+    x2 = x[..., half:]   # (b,h,t,half)
+
+    x1_rot = x1 * cos - x2 * sin
+    x2_rot = x1 * sin + x2 * cos
+    x_rot = jnp.concatenate([x1_rot, x2_rot], axis=-1)  # (b,h,t,d)
+    return x_rot
+
+class MultiHeadSelfAttention(nn.Module):
+    d_model: int
+    n_heads: int
+    dropout: float = 0.0
+    use_rope: bool = True
+
+    @nn.compact
+    def __call__(self, x, deterministic: bool):
+        """
+        x: (batch, seq_len, d_model)
+        returns: (batch, seq_len, d_model)
+        """
+        b, t, d_model = x.shape
+        assert d_model == self.d_model
+        assert self.d_model % self.n_heads == 0, "d_model must be divisible by n_heads"
+        head_dim = self.d_model // self.n_heads
+        assert head_dim % 2 == 0, "head_dim must be even for RoPE"
+
+        # project to qkv
+        qkv = nn.Dense(3 * self.d_model, use_bias=False, name="qkv")(x)  # (b, t, 3*d_model)
+        qkv = qkv.reshape(b, t, 3, self.n_heads, head_dim)
+        qkv = qkv.transpose(2, 0, 3, 1, 4)  # (3, b, h, t, head_dim)
+        q, k, v = qkv[0], qkv[1], qkv[2]    # each: (b, h, t, head_dim)
+
+        if self.use_rope:
+            inv_freq = build_rotary_inv_freq(head_dim)
+            q = apply_rope(q, inv_freq)
+            k = apply_rope(k, inv_freq)
+
+        # scaled dot-product attention with causal mask
+        scale = 1.0 / math.sqrt(head_dim)
+        attn_logits = jnp.einsum("bhqd, bhkd -> bhqk", q, k) * scale  # (b,h,t,t)
+
+        # causal mask
+        mask = jnp.tril(jnp.ones((t, t), dtype=jnp.bool_))
+        attn_logits = jnp.where(mask, attn_logits, -1e9)
+
+        attn_weights = nn.softmax(attn_logits, axis=-1)
+        attn_weights = nn.Dropout(rate=self.dropout)(attn_weights, deterministic=deterministic)
+
+        attn_output = jnp.einsum("bhqk, bhkd -> bhqd", attn_weights, v)  # (b,h,t,head_dim)
+        attn_output = attn_output.transpose(0, 2, 1, 3).reshape(b, t, self.d_model)  # (b,t,d_model)
+        out = nn.Dense(self.d_model, name="out_proj")(attn_output)
+        out = nn.Dropout(rate=self.dropout)(out, deterministic=deterministic)
+        return out
+
+class FeedForward(nn.Module):
+    d_model: int
+    mlp_ratio: float = 4.0
+    dropout: float = 0.0
+
+    @nn.compact
+    def __call__(self, x, deterministic: bool):
+        hidden_dim = int(self.d_model * self.mlp_ratio)
+        x = nn.Dense(hidden_dim)(x)
+        x = nn.gelu(x)
+        x = nn.Dropout(rate=self.dropout)(x, deterministic=deterministic)
+        x = nn.Dense(self.d_model)(x)
+        x = nn.Dropout(rate=self.dropout)(x, deterministic=deterministic)
+        return x
+
+
+class TransformerBlock(nn.Module):
+    d_model: int
+    n_heads: int
+    dropout: float = 0.0
+    mlp_ratio: float = 4.0
+    use_rope: bool = True
+
+    @nn.compact
+    def __call__(self, x, deterministic: bool):
+        # Self-attention
+        h = nn.LayerNorm()(x)
+        h = MultiHeadSelfAttention(
+            d_model=self.d_model,
+            n_heads=self.n_heads,
+            dropout=self.dropout,
+            use_rope=self.use_rope,
+        )(h, deterministic=deterministic)
+        x = x + h
+
+        # FFN
+        h2 = nn.LayerNorm()(x)
+        h2 = FeedForward(
+            d_model=self.d_model,
+            mlp_ratio=self.mlp_ratio,
+            dropout=self.dropout,
+        )(h2, deterministic=deterministic)
+        x = x + h2
+        return x
+
+class TARTModel(nn.Module):
+    """
+    Core autoregressive Transformer that maps
+      (batch, L, D_in) -> (batch, L, D_in)
+    and we train it to predict the *next* token at each step.
+    For forecasting, we use the last position as the next-step prediction.
+    """
+    input_dim: int
+    d_model: int
+    n_heads: int
+    depth: int
+    max_len: int
+    dropout: float = 0.0
+    use_learned_pos: bool = True
+    use_rope: bool = True
+
+    @nn.compact
+    def __call__(self, x, deterministic: bool = True):
+        """
+        x: (batch, seq_len, input_dim)
+        returns: (batch, seq_len, input_dim)
+        """
+        b, t, d_in = x.shape
+        assert d_in == self.input_dim
+        if t > self.max_len:
+            raise ValueError(f"seq_len={t} exceeds max_len={self.max_len}")
+
+        # Project input to d_model
+        h = nn.Dense(self.d_model, name="token_embed")(x)
+
+        # Learned positional embedding
+        if self.use_learned_pos:
+            pos_emb = self.param(
+                "pos_emb",
+                nn.initializers.normal(stddev=0.02),
+                (self.max_len, self.d_model),
+            )  # (max_len, d_model)
+            positions = jnp.arange(t)[None, :]          # (1, t)
+            h = h + pos_emb[positions, :]               # broadcast to (b,t,d_model)
+
+        # Stacked transformer blocks
+        for i in range(self.depth):
+            h = TransformerBlock(
+                d_model=self.d_model,
+                n_heads=self.n_heads,
+                dropout=self.dropout,
+                mlp_ratio=4.0,
+                use_rope=self.use_rope,
+                name=f"block_{i}",
+            )(h, deterministic=deterministic)
+
+        h = nn.LayerNorm(name="final_ln")(h)
+        out = nn.Dense(self.input_dim, name="out_proj")(h)  # (b,t,D_in)
+        return out
+
+class TART:
+    """
+    Time-series AutoReg Transformer (TART)
+
+    Usage:
+      forecaster = TART(R_tX, L=128, d_model=64, depth=4, ...)
+      preds = forecaster(F_cX, steps=200)
+
+    - R_tX: training time-series, shape (T, D)
+    - L:    context length (sequence length seen by the Transformer)
+    """
+
+    def __init__(self,
+                 R_tX: np.ndarray,
+                 L: int,
+                 d_model: int = 64,
+                 depth: int = 4,
+                 n_heads: Optional[int] = None,
+                 use_learned_pos: bool = True,
+                 use_rope: bool = True,
+                 dropout: float = 0.0,
+                 seed: int = 0,
+                 val_split: float = 0.1,
+                 batch_size: int = 64,
+                 max_epochs: int = 50,
+                 init_lr: float = 3e-4,
+                 min_lr: float = 1e-5,
+                 lr_decay: float = 0.5,
+                 tol_rel_improve: float = 1e-3,
+                 patience: int = 5):
+
+        R_tX = np.asarray(R_tX, dtype=np.float32)
+        if R_tX.ndim == 1:
+            R_tX = R_tX[:, None]
+        T, D = R_tX.shape
+        if T <= L:
+            raise ValueError(f"T={T} must be > L={L}")
+
+        self.L = int(L)
+        self.D = int(D)
+
+        # choose d_model / n_heads
+        d_model = int(max(16, d_model))
+        if n_heads is None:
+            n_heads = choose_heads(d_model)
+        self.d_model = d_model
+        self.n_heads = int(n_heads)
+        self.depth = int(depth)
+
+        # normalize training data
+        F_norm, (mu, sd) = zscore(R_tX)
+        self._mu = mu
+        self._sd = sd
+
+        # build 1-step-ahead dataset
+        X_in, Y_out = make_nextstep_windows(F_norm, L=self.L)  # (N,L,D), (N,D)
+        N = X_in.shape[0]
+        n_val = max(1, int(val_split * N))
+        n_tr  = N - n_val
+        Xtr, Ytr = X_in[:n_tr], Y_out[:n_tr]
+        Xva, Yva = X_in[n_tr:], Y_out[n_tr:]
+
+        print(f"[TART] Train N={N} (train={n_tr}, val={n_val}) | "
+              f"L={self.L} D={self.D} | d_model={d_model} heads={n_heads} depth={depth}")
+
+        # build model
+        max_len = self.L          # always feed sequences of length L
+        model = TARTModel(
+            input_dim=self.D,
+            d_model=self.d_model,
+            n_heads=self.n_heads,
+            depth=self.depth,
+            max_len=max_len,
+            dropout=dropout,
+            use_learned_pos=use_learned_pos,
+            use_rope=use_rope,
+        )
+        self.model = model
+        self.max_len = max_len
+
+        rng = jax.random.PRNGKey(seed)
+        dummy_x = jnp.zeros((1, self.L, self.D), dtype=jnp.float32)
+        params = model.init(rng, dummy_x, deterministic=True)["params"]
+
+        def create_state(lr):
+            tx = optax.adamw(learning_rate=lr, weight_decay=0.0)
+            return train_state.TrainState.create(apply_fn=model.apply,
+                                                 params=params,
+                                                 tx=tx)
+
+        state = create_state(init_lr)
+        self._rng = rng
+
+        @jax.jit
+        def train_step(state, x_batch, y_batch, rng):
+            """One training step (MSE on next-step prediction of last token)."""
+            dropout_rng, new_rng = jax.random.split(rng)
+
+            def loss_fn(p):
+                preds = state.apply_fn({"params": p},
+                                       x_batch,
+                                       deterministic=False,
+                                       rngs={"dropout": dropout_rng})
+                pred_last = preds[:, -1, :]            # (B,D)
+                loss = jnp.mean((pred_last - y_batch) ** 2)
+                return loss
+
+            loss, grads = jax.value_and_grad(loss_fn)(state.params)
+            new_state = state.apply_gradients(grads=grads)
+            return new_state, loss, new_rng
+
+        @jax.jit
+        def eval_step(state, x_batch, y_batch):
+            preds = state.apply_fn({"params": state.params},
+                                   x_batch,
+                                   deterministic=True)
+            pred_last = preds[:, -1, :]
+            loss = jnp.mean((pred_last - y_batch) ** 2)
+            return loss
+
+        # move training data to JAX arrays once
+        Xtr_j = jnp.asarray(Xtr)
+        Ytr_j = jnp.asarray(Ytr)
+        Xva_j = jnp.asarray(Xva)
+        Yva_j = jnp.asarray(Yva)
+
+        best_params = state.params
+        best_val = float("inf")
+        curr_lr = init_lr
+        epochs_no_gain = 0
+
+        num_batches = lambda N: int(np.ceil(N / batch_size))
+
+        for epoch in range(1, max_epochs + 1):
+            # shuffle training indices
+            idx = np.arange(n_tr)
+            np.random.default_rng(epoch + seed).shuffle(idx)
+
+            # training loop
+            train_losses = []
+            rng = self._rng
+            for bi in range(num_batches(n_tr)):
+                s = bi * batch_size
+                e = min((bi + 1) * batch_size, n_tr)
+                batch_idx = idx[s:e]
+                xb = Xtr_j[batch_idx]
+                yb = Ytr_j[batch_idx]
+                state, loss, rng = train_step(state, xb, yb, rng)
+                train_losses.append(float(loss))
+
+            self._rng = rng
+            tr_loss = float(np.mean(train_losses))
+
+            # validation
+            val_losses = []
+            for bi in range(num_batches(n_val)):
+                s = bi * batch_size
+                e = min((bi + 1) * batch_size, n_val)
+                xb = Xva_j[s:e]
+                yb = Yva_j[s:e]
+                val_losses.append(float(eval_step(state, xb, yb)))
+            va_loss = float(np.mean(val_losses))
+
+            print(f"[TART] epoch {epoch:03d} | train {tr_loss:.6e} | val {va_loss:.6e} | lr {curr_lr:.2e}")
+
+            # track best
+            if va_loss + 1e-8 < best_val:
+                rel_gain = (best_val - va_loss) / max(best_val, 1e-8) if best_val < float("inf") else 1.0
+                best_val = va_loss
+                best_params = state.params
+                epochs_no_gain = 0
+            else:
+                rel_gain = 0.0
+                epochs_no_gain += 1
+
+            # LR schedule on plateau
+            if epochs_no_gain >= patience:
+                if curr_lr > min_lr * (1.0 + 1e-9):
+                    curr_lr = max(min_lr, curr_lr * lr_decay)
+                    print(f"[TART] plateau → lowering LR to {curr_lr:.2e}")
+                    # rebuild optimizer with new LR, keep params
+                    params = state.params
+                    state = train_state.TrainState.create(
+                        apply_fn=state.apply_fn,
+                        params=params,
+                        tx=optax.adamw(learning_rate=curr_lr, weight_decay=0.0),
+                    )
+                    epochs_no_gain = 0
+                else:
+                    print(f"[TART] early stop: lr at min and no improvement (best val {best_val:.6e})")
+                    break
+
+        # store best params
+        self.state = state.replace(params=best_params)
+        self.params = self.state.params
+
+        # parameter count (for paper)
+        self.param_count = sum(p.size for p in jax.tree_util.tree_leaves(self.params))
+        print(f"[TART] params: {self.param_count:,}")
+
+    def __call__(self, F_cX: np.ndarray, steps: int) -> np.ndarray:
+        """
+        Forecast `steps` points given any context slice F_cX that ENDS at
+        the forecast start. Uses last L points (with left-padding if needed).
+        """
+        X = np.asarray(F_cX, dtype=np.float32)
+        if X.ndim == 1:
+            X = X[:, None]
+        if X.shape[1] != self.D:
+            raise ValueError(f"Expected D={self.D}, got {X.shape[1]}")
+
+        # take last L rows (pad on the left if needed)
+        if X.shape[0] < self.L:
+            pad_len = self.L - X.shape[0]
+            pad = np.repeat(X[:1], repeats=pad_len, axis=0)
+            ctx = np.concatenate([pad, X], axis=0)
+        else:
+            ctx = X[-self.L:]
+
+        # normalize with training stats
+        ctx_n = (ctx - self._mu) / self._sd
+        ctx_n = jnp.asarray(ctx_n[None, :, :])  # (1, L, D)
+
+        preds_n = []
+        params = self.params
+
+        for _ in range(steps):
+            # deterministic forward (no dropout)
+            out = self.model.apply({"params": params}, ctx_n, deterministic=True)  # (1,L,D)
+            next_n = out[:, -1, :]          # (1,D)
+            preds_n.append(np.array(next_n[0]))
+
+            # update context: drop oldest, append new
+            ctx_n = jnp.concatenate([ctx_n[:, 1:, :], next_n[:, None, :]], axis=1)
+
+        preds_n = np.stack(preds_n, axis=0)      # (steps, D) normalized
+        preds = preds_n * self._sd + self._mu    # unnormalize
+        return preds