Skip to content

Add SMT (Supervised Memory Training) estimator - #3302

Open
kashif wants to merge 11 commits into
awslabs:devfrom
kashif:smt-model
Open

Add SMT (Supervised Memory Training) estimator#3302
kashif wants to merge 11 commits into
awslabs:devfrom
kashif:smt-model

Conversation

@kashif

@kashif kashif commented Jun 8, 2026

Copy link
Copy Markdown
Contributor

Adds a torch estimator implementing Supervised Memory Training (SMT, Kumar & Isola 2026, https://arxiv.org/abs/2606.06479) as an alternative to DeepAR.

What it does

  • A bidirectional Transformer encoder compresses the past (target lags + covariates) into a memory of learned register tokens.
  • A causal decoder predicts each future target from the previous future target and future covariates — a predictive-state objective with a DistributionOutput NLL head.
  • A recurrent memory cell is trained to reproduce the teacher's one-step memory transitions (m_t, x_{t+1}) -> m_{t+1}, so the RNN is trained without backpropagation through time.
  • At inference the recurrent cell is unrolled with O(1) state and sampled autoregressively.

It reuses the DeepAR feature pipeline (lags, time/age features, static embeddings, mean scaling), so SMTEstimator is a drop-in with the same covariate handling and DistributionOutput heads.

Included

  • gluonts.torch.model.smt — module, lightning module, estimator
  • unit tests in test/torch/model/test_smt.py
  • an electricity vs DeepAR comparison notebook in examples/

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant