""" Online (TBPTT) training: Shared-salience multi-channel HiPPO on vector-token selective copying. Training uses one sampled episode per iteration, splits each episode into fixed TBPTT chunks, and applies one optimizer update per chunk. The carried state is stop_gradient'ed at chunk boundaries. """ __date__ = "January 2026" import numpy as np import matplotlib.pyplot as plt import jax import jax.numpy as jnp import jax.random as jr from jax import jit, value_and_grad from jax.lax import scan from tqdm import tqdm import optax from mpm import get_system_params from mpm.models import SalienceHiPPO from mpm.polynomials import legsval from mpm.tasks.selective_recall import make_selective_copying_task_tokens # ---------------------------- # TBPTT helpers # ---------------------------- def pad_chunk(X, Y, t0: int, L: int): """ Extract a chunk starting at t0 of length <= L and pad to exactly L. Returns Xc, Yc, mask where mask[t]=1 for valid timesteps. """ T = X.shape[0] t1 = min(t0 + L, T) ell = t1 - t0 Xc = jnp.zeros((L, X.shape[1]), dtype=X.dtype) Yc = jnp.zeros((L, Y.shape[1]), dtype=Y.dtype) mask = jnp.zeros((L,), dtype=X.dtype) Xc = Xc.at[:ell].set(X[t0:t1]) Yc = Yc.at[:ell].set(Y[t0:t1]) mask = mask.at[:ell].set(1.0) return Xc, Yc, mask def main(): # ---------------------------- # Config # ---------------------------- seed = 42 key = jr.PRNGKey(seed) # task episode_len = 30 num_episodes = 10000 # online training steps (episodes seen) num_steps = 3000 # TBPTT tbptt_steps = 30 # token dims / model dims d_in = 32 d_model = 64 # hippo n = 256 # 64 measure = "legs" base_timescale = episode_len # salience range g_max = 5.0 # opt lr = 1e-3 # ---------------------------- # Dataset # ---------------------------- X_np, Y_np, token_table, ids_np = make_selective_copying_task_tokens( episode_len=episode_len, num_episodes=num_episodes, d_in=d_in, seed=seed, ) X_all = jnp.array(X_np, dtype=jnp.float32) # (N,T,d_in) Y_all = jnp.array(Y_np, dtype=jnp.float32) ids_all = np.array(ids_np) # ---------------------------- # HiPPO system # ---------------------------- (A, b), _, _ = get_system_params(measure, n) A, b = A / base_timescale, b / base_timescale A = jnp.array(A, dtype=jnp.float32) b = jnp.array(b, dtype=jnp.float32) # ---------------------------- # Init model + optimizer # ---------------------------- model = SalienceHiPPO( d_in=d_in, d_model=d_model, n=n, A=A, b=b, g_max=g_max, sal_hidden=128, outvec_hidden=128, mem_dim=32, ) key, k1 = jr.split(key) S0 = jnp.zeros((d_model, n), dtype=jnp.float32) x0 = jnp.zeros((d_in,), dtype=jnp.float32) params = model.init(k1, x0, S0)["params"] optimizer = optax.adamw(learning_rate=lr) opt_state = optimizer.init(params) # ---------------------------- # Chunk rollout + TBPTT update # ---------------------------- def chunk_loss_and_state(params, S_init, Xc, Yc, mask): """ One TBPTT chunk. S_init: (d_model,n) Xc,Yc: (L,d_in) padded mask: (L,) in {0,1} Returns: (loss_scalar, S_final, aux) """ def step_fn(S, inputs): x_t, y_t, m_t = inputs S_next, yhat, g, _ = model.apply({"params": params}, x_t, S) mse_t = jnp.mean((yhat - y_t) ** 2) # masked mean contribution (mask will be normalized at end) return S_next, (m_t * mse_t, m_t, yhat, g) S_final, (mse_masked, m_sum, yhat_ts, g_ts) = scan( step_fn, S_init, (Xc, Yc, mask), ) denom = jnp.maximum(jnp.sum(m_sum), 1.0) loss = jnp.sum(mse_masked) / denom aux = (yhat_ts, g_ts) return loss, S_final, aux @jit def tbptt_chunk_train_step(params, opt_state, S_init, Xc, Yc, mask): """ One optimizer step for one chunk (gradients flow only within the chunk). """ def loss_fn(p): loss, S_final, aux = chunk_loss_and_state(p, S_init, Xc, Yc, mask) # pack S_final into aux so has_aux=True is satisfied return loss, (S_final, aux) (loss, (S_final, aux)), grads = value_and_grad(loss_fn, has_aux=True)(params) updates, opt_state = optimizer.update(grads, opt_state, params) params = optax.apply_updates(params, updates) return params, opt_state, S_final, loss, aux # ---------------------------- # Online training loop (episode-by-episode) with TBPTT # ---------------------------- losses = [] key = jr.PRNGKey(seed + 1) N = X_all.shape[0] pbar = tqdm(range(num_steps)) smooth_loss = None # state resets per episode (typical for this task) S = jnp.zeros((d_model, n), dtype=jnp.float32) for step in pbar: key, sub = jr.split(key) idx = jr.randint(sub, (), 0, N) X_ep = X_all[idx] # (T,d_in) Y_ep = Y_all[idx] # (T,d_in) # TBPTT over chunks T = X_ep.shape[0] n_chunks = (T + tbptt_steps - 1) // tbptt_steps loss_ep = 0.0 for c in range(n_chunks): t0 = c * tbptt_steps Xc, Yc, mask = pad_chunk(X_ep, Y_ep, t0=t0, L=tbptt_steps) params, opt_state, S_end, loss_chunk, _ = tbptt_chunk_train_step( params, opt_state, S, Xc, Yc, mask ) loss_ep = loss_ep + loss_chunk # CRITICAL: stop gradient at TBPTT boundary S = jax.lax.stop_gradient(S_end) loss_ep = loss_ep / n_chunks losses.append(float(loss_ep)) if smooth_loss is None: smooth_loss = losses[-1] else: smooth_loss = 0.98 * smooth_loss + 0.02 * losses[-1] pbar.set_description(f"loss: {smooth_loss:.7f}") # ---------------------------- # Evaluate + plot: last episode # ---------------------------- X_eval = X_all[-1] # (T,d_in) Y_eval = Y_all[-1] ids_eval = ids_all[-1] def full_rollout_collect(params, X, Y): S = jnp.zeros((d_model, n), dtype=jnp.float32) def step_fn(S, inputs): x_t, y_t = inputs S_next, yhat, g, out_vec = model.apply({"params": params}, x_t, S) loss_t = jnp.mean((yhat - y_t) ** 2) return S_next, (loss_t, yhat, g, S_next, out_vec) _, (loss_ts, yhat_ts, g_ts, S_ts, outvec_ts) = scan(step_fn, S, (X, Y)) return jnp.mean(loss_ts), yhat_ts, g_ts, S_ts, outvec_ts # outvec_ts: (T, n) loss_eval, preds_ts, gs_ts, S_ts, outvec_ts = full_rollout_collect(params, X_eval, Y_eval) preds = np.array(preds_ts) # (T,d_in) gs = np.array(gs_ts) # (T,) Xp = np.array(X_eval) Yp = np.array(Y_eval) def cos_sim(a, b, eps=1e-8): na = np.linalg.norm(a, axis=-1) nb = np.linalg.norm(b, axis=-1) return np.sum(a * b, axis=-1) / (na * nb + eps) cos = cos_sim(preds, Yp) t = np.arange(episode_len) fig, axarr = plt.subplots(7, 1, figsize=(10, 14)) axarr[0].set_title(f"Online TBPTT selective copying (eval loss={float(loss_eval):.4g})") axarr[0].plot(np.arange(len(losses)), losses) axarr[0].set_ylabel("Train loss") dims_to_plot = [0, 1, 2] for k in dims_to_plot: axarr[1].plot(t, Yp[:, k], label=f"Y dim{k}", linewidth=2) axarr[1].plot(t, preds[:, k], linestyle="--", label=f"pred dim{k}", alpha=0.8) axarr[1].set_ylabel("Selected dims") axarr[1].legend(loc="upper left", ncol=2) axarr[2].plot(t, cos) axarr[2].set_ylabel("cos(pred, target)") axarr[2].set_ylim(-0.05, 1.05) axarr[3].plot(t, gs) axarr[3].set_ylabel("Salience g") axarr[3].set_xlabel("Timestep") axarr[3].set_ylim(0, None) # ---------------------------- # Interpretability plot (sub-timestep): input + decoded memories, and induced measures # ---------------------------- # Convert to numpy S_np = np.array(S_ts) # (T, d_model, n) g_np = np.array(gs_ts) # (T,) X_np_ep = np.array(X_eval) # (T, d_in) T = episode_len # Pull learned W_in and compute projected inputs x_proj[t, j] W_in = np.array(params["W_in"]["kernel"]) # (d_in, d_model) xproj_np = X_np_ep @ W_in # (T, d_model) # Choose a channel to visualize: largest energy in this episode ch = int(np.argmax(np.sum(xproj_np**2, axis=0))) # --- Sub-timestep warp phi(t) with piecewise-constant g(t) on [k,k+1) # phi(t) = \int_0^t g(s)/base ds g_scaled = g_np / float(base_timescale) # (T,) # prefix[k] = sum_{i np.ndarray: """ Piecewise-linear integral of g_scaled on [0,T], assuming g constant on [k,k+1). t_cont can be float array in [0,T]. Returns same shape. """ t = np.clip(t_cont, 0.0, float(T)) k = np.floor(t).astype(np.int32) # handle t==T: clamp k to T-1 and set frac=1, then phi=prefix[T] k_clamped = np.clip(k, 0, T - 1) frac = (t - k).astype(np.float32) frac = np.where(t >= float(T), 1.0, frac) # ensures exact endpoint return prefix[k_clamped] + frac * g_scaled[k_clamped].astype(np.float32) # Sub-timestep time grid for plotting (absolute time u) # (dense enough to look continuous) u_grid = np.linspace(0.0, float(T - 1), 600, dtype=np.float32) # g(u) piecewise constant ku = np.floor(u_grid).astype(np.int32) ku = np.clip(ku, 0, T - 1) g_u = g_np[ku].astype(np.float32) g_u_scaled = g_u / float(base_timescale) # True projected input on that dense grid (piecewise constant) true_u = xproj_np[ku, ch].astype(np.float32) # Choose three snapshot times (integers; you can change these) first_len = (2 * T) // 3 t_snaps = [first_len - 1, first_len + 2, T - 1] # 3 timepoints def decoded_memory_on_u(t_idx: int): """ For snapshot t_idx (integer), return decoded memory as a function of absolute time u. Memory curve is defined only for u <= t_idx, else NaN. """ # coefficients at the snapshot (channel ch) c = S_np[t_idx, ch, :].astype(np.float32)[None, :] # (1, n) # tau1(u;t) = phi(t) - phi(u) phi_t = float(phi_cont(np.array([float(t_idx) + 1.0], dtype=np.float32))[0]) tau1 = (phi_t - phi_cont(u_grid)).astype(np.float32) # (len(u_grid),) # valid only for u <= t_idx valid = (u_grid <= float(t_idx)) tau1_eval = np.where(valid, tau1, 0.0).astype(np.float32) # decode \hat f(tau1) hat = legsval(tau1_eval, c)[0].astype(np.float32) # (len(u_grid),) hat = np.where(valid, hat, np.nan).astype(np.float32) return hat, valid def induced_measure_on_u(t_idx: int): """ \omega0(u|t) propto exp(-(phi(t)-phi(u))) * g(u)/base_timescale Defined for u <= t_idx, else NaN. Normalized to integrate to 1 over u<=t. """ phi_t = float(phi_cont(np.array([float(t_idx) + 1.0], dtype=np.float32))[0]) tau1 = (phi_t - phi_cont(u_grid)).astype(np.float32) valid = (u_grid <= float(t_idx)) w = np.exp(-tau1) * g_u_scaled # (len(u_grid),) w = np.where(valid, w, 0.0).astype(np.float32) # normalize as a density over u (continuous approx) Z = np.trapz(w, u_grid) + 1e-8 w = w / Z w = np.where(valid, w, np.nan).astype(np.float32) return w, valid # Build curves mem_curves = [] meas_curves = [] for t_idx in t_snaps: hat, _ = decoded_memory_on_u(t_idx) w, _ = induced_measure_on_u(t_idx) mem_curves.append((t_idx, hat)) meas_curves.append((t_idx, w)) # (1) Input (projected) + decoded memories (flipped horizontally because x-axis is absolute time u) axarr[4].plot(u_grid, true_u, linewidth=2, label=f"true input proj (channel {ch})") for (t_idx, hat) in mem_curves: axarr[4].plot(u_grid, hat, linestyle="--", linewidth=2, label=f"decoded memory @ t={t_idx}") axarr[4].set_ylabel("value (proj)") axarr[4].set_title("Input (projected) and decoded memories (absolute time axis)") axarr[4].legend(loc="upper left", ncol=2) # (2) Measures for three timepoints on the same axis for (t_idx, w) in meas_curves: axarr[5].fill_between(u_grid, 0.0, w, alpha=0.25, label=f"\omega_0(u|t={t_idx})") axarr[5].plot(u_grid, w, linewidth=2) # optional outline for readability axarr[5].set_ylim(0, None) axarr[5].set_ylabel("density") axarr[5].set_xlabel("absolute time u (unwarped)") axarr[5].set_title("Induced history measure \omega_0(u|t) (normalized densities)") axarr[5].legend(loc="upper left", ncol=2) # (3) Linear functionals (kernels) for each prediction timestep in the last third outvec_np = np.array(outvec_ts) # (T, n) write_start = first_len t_preds = list(range(write_start, T)) # each prediction in last third for t_idx in t_preds: # post-step time corresponds to t_idx + 1 phi_t = float(phi_cont(np.array([float(t_idx) + 1.0], dtype=np.float32))[0]) tau1 = (phi_t - phi_cont(u_grid)).astype(np.float32) valid = (u_grid <= float(t_idx)) tau1_eval = np.where(valid, tau1, 0.0).astype(np.float32) # Decode warped kernel k1(tau1) c = outvec_np[t_idx].astype(np.float32)[None, :] # (1, n) k1_u = legsval(tau1_eval, c)[0].astype(np.float32) # Measure factor omega0(u|t) propto exp(-tau1) * g(u)/base_timescale omega_u = np.exp(-tau1_eval) * g_u_scaled # g_u_scaled already = g(u)/base_timescale # Effective absolute-time functional k_eff = k1_u * omega_u # Mask future k_eff = np.where(valid, k_eff, np.nan).astype(np.float32) axarr[-1].plot(u_grid, k_eff, alpha=0.5, linewidth=1.5) axarr[-1].set_ylabel("kernel value") axarr[-1].set_xlabel("absolute time u (unwarped)") axarr[-1].set_title("Time-varying linear functionals (decoded kernels) for predictions in last third") plt.tight_layout() plt.savefig("temp.png") plt.show() if __name__ == "__main__": main()