TAM: Task-Aware Memory Distillation for Efficient Spatiotemporal Prediction

• Prior work supplies varied predictors and distillation techniques, but teacher accuracy alone does not guarantee effective transfer across different model capacities.These findings motivate choosing supervision that a student can exploit.Source: §1 · p. 1
• Historical references have task-specific meanings: video depends on motion or occupancy, weather on location, and traffic on observed flow and recent trend.The introduction therefore frames both the stored predictive quantity and the historical entries used for supervision as design choices.Source: §1 · p. 2
• 1.01%: TAM reduces TaxiBJ test MSE on average versus no-memory KD in the reported four-run comparison.The comparison fixes the TAU teacher and U-Net-Base student; gains vary across runs, and the comparison does not isolate conditioned retrieval from horizon grouping.Source: Figure 1 · p. 2, §1 · p. 2
• TAM stores frozen-teacher latent features, forecast changes, or flow residuals and selects references according to task structure.Relational variants match teacher and student similarity distributions over shared references; traffic can instead regress retrieved residual prototypes.Source: §1 · p. 2
• TAM supplements ground-truth prediction, output distillation, and optional feature alignment, while leaving the student’s inference architecture and cost unchanged.The evaluation compares its additional effect against matched distillation baselines across video, weather, and traffic prediction.Source: §1 · p. 2
• The stated contributions are a task-aware interface for cross-sample supervision, two memory objectives, and repeated-run evaluations across three prediction domains.The objectives are shared-reference relation matching and observation-conditioned residual prototypes.Source: §1 · p. 2

• Spatiotemporal prediction research spans recurrent, physics-guided, attention-based, and recurrent-free architectures, whereas TAM augments supervision without changing student inference architecture.Examples include ConvLSTM, PredRNN, PhyDNet, SimVP, TAU, Earthformer, and Earthfarsser.Source: §2 · p. 3
• Distillation methods transfer hints, attention, activation boundaries, feature stages, channel-wise maps, frequency representations, or semantic priors to students.TAM complements these objectives by organizing teacher features or forecast changes across training sequences and selecting task-specific references.Source: §2 · p. 3
• Relational distillation and memory methods provide precedents, including CIRKD’s shared-reference alignment, while TAM adapts teacher history and retrieval to forecasting tasks.Video uses motion or occupancy, weather uses geographic location, and traffic conditions residual retrieval on observed flow; traffic also uses residual-prototype regression.Source: §2 · p. 3

3 METHOD

• TAM supplements per-sample distillation with a bounded history, using task-specific quantities and references to supervise relations or retrieved residual targets.The framework retains the original student-only inference path.Source: §3 · p. 3, Figure 2 · p. 4

3.1 PROBLEM FORMULATION

• With a fixed-variance isotropic Gaussian observation model, maximizing conditional likelihood is equivalent to minimizing mean squared prediction error.The deterministic predictor supplies the forecast used at inference.Source: §3.1 · p. 4

3.2 TASK-AWARE TEACHER MEMORY

• Task-aware memory represents latent spatial states, forecast changes, or flow residuals, then selects eligible historical references using task metadata and similarity.Examples include video anchors selected by motion or occupancy, weather references restricted by geographic location, and observation-key retrieval for related input conditions.Source: §3.2 · pp. 4–5

3.3 DISTILLATION OVER SHARED HISTORICAL REFERENCES

• Shared-reference distillation matches teacher and student similarity distributions over the same historical teacher vectors, transferring relative relationships to past samples.Video anchors use uniform weights; geographic-grid variants may use latitude-dependent weights.Source: §3.3 · p. 5

3.4 OBSERVATION-CONDITIONED RESIDUAL PROTOTYPES

• For traffic, memory stores teacher residuals relative to the latest observation and retrieves near- and far-horizon prototypes using similar observed-input keys.The student regresses both prototypes; separating horizon groups preserves horizon-specific targets, and the objective transfers residual values directly.Source: §3.4 · p. 5

3.5 JOINT OPTIMIZATION AND MEMORY UPDATES

• Joint training combines prediction, output distillation, optional feature alignment, and the task’s memory loss; retrieval precedes insertion of detached teacher entries.The bounded memory retains recent history subject to task constraints, and the teacher, memory, and adapters are absent at inference.Source: §3.5 · pp. 5–6

4.1 EXPERIMENTAL SETUP

• Seven video benchmarks plus WeatherBench and TaxiBJ compare supervised training, KD, and TAM-augmented matched teacher–student configurations.BAIR, WeatherBench, and TaxiBJ use episode-aware retrieval, location-specific history, and observation-conditioned residual prototypes, respectively. Teachers remain frozen and inference uses only the student; seeds 42–45 support matched comparisons, while four video implementations share validation and test data for checkpoint selection.Source: §4.1 · p. 6

4.2 COMPARISON WITH DISTILLATION BASELINES

• With teacher, student, and KD objective fixed, TAM raises mean SSIM on all six video datasets and lowers mean MSE on five.Gains vary by metric: HMDB51 and BAIR improve SSIM without corresponding LPIPS improvement.Source: Table 1 · p. 7, §4.2 · p. 7

4.3 WEATHER AND TRAFFIC FORECASTING

• 1.01%: TaxiBJ paired MSE reduction over Feature-MSE KD; WeatherBench mean MSE, MAE, and RMSE fall with both teachers.Feature-AB KD remains stronger on TaxiBJ; WeatherBench benefits depend on the teacher and do not consistently reduce variance.Source: Table 2 · p. 8, §4.3 · p. 8

4.4 PREDICTION HORIZON ANALYSIS

• 4.33% to 0.20%: BAIR’s MSE reduction from the first to last prediction step; TaxiBJ’s four-step reductions include a −0.17% reversal.On BAIR, LPIPS improves only at the first two steps, and gains across horizons are not monotonic.Source: Figure 3 · p. 8, §4.4 · p. 8

4.5 SPATIAL FREQUENCY ANALYSIS

• 15.44%, 11.14%, and 6.73%: HMDB51 spectral-error reductions in the low-, mid-, and high-frequency bands.KTH reductions are smaller in the first two bands, with high-frequency error essentially unchanged.Source: Figure 4 · p. 9, §4.5 · p. 9

4.6 MOTION-CONDITIONED PREDICTION ANALYSIS

• 9.68%, 14.31%, and 12.12%: HMDB51 MSE reductions across low-, middle-, and high-motion groups; KittiCaltech reductions are 4.78%, 4.85%, and 4.45%.Gains extend across motion conditions but do not increase monotonically with motion strength.Source: Figure 5 · p. 9, §4.6 · p. 9

5 CONCLUSION

• TAM turns a frozen teacher’s historical predictions and representations into cross-sample supervision through task-specific representations and reference selection.The student aligns relations over shared history or regresses observation-conditioned residual prototypes.Source: §5 · pp. 9–10
• Experiments across video, weather, and traffic prediction report improvements across multiple teacher–student configurations without additional student inference cost.Prediction-horizon, frequency, and motion-conditioned analyses find that gains vary, and pixel and structural improvements do not necessarily yield consistent perceptual gains.Source: §5 · p. 10
• The framework relies on task-specific representations and retrieval rules, and available comparisons do not fully isolate individual design contributions.Future work will investigate more unified reference selection and controlled ablations.Source: §5 · p. 10

AI USE STATEMENT

• ChatGPT and Codex assisted with literature search and synthesis, framework development, mathematical presentation, drafting, and Chinese-English translation.Source: AI USE STATEMENT · p. 11
• The tools also assisted with implementation and debugging, experimental and post-hoc analysis design, result interpretation, and figure creation and typesetting.Source: AI USE STATEMENT · p. 11
• The authors say they reviewed and verified AI-assisted work and take responsibility for the final text, code, results, and scientific conclusions.Source: AI USE STATEMENT · p. 11

REPRODUCIBILITY STATEMENT

• Section 3 and Appendix A describe the method and task-specific memory designs.Source: REPRODUCIBILITY STATEMENT · p. 11
• Section 4.1 and Appendices B–C document preprocessing, data splits, checkpoint selection, model configurations, training hyperparameters, and evaluation rules.Source: REPRODUCIBILITY STATEMENT · p. 11
• Appendices D–E specify sampling and computation for frequency, motion-conditioned, and temporal analyses; the paper reports repeated runs and paired comparisons and identifies settings sharing validation and test data.Code and experimental configurations will be released.Source: REPRODUCIBILITY STATEMENT · p. 11

A TASK-SPECIFIC MEMORY INSTANTIATIONS

• The memory interface specializes training objectives by task, while the prediction network does not access memory at inference.Source: A TASK-SPECIFIC MEMORY INSTANTIATIONS · p. 15
• Table A1 separates the listed representations and anchor selection from reference organization, and defines region r as non-overlapping r×r average pooling.Source: Table A1 · p. 15
• Video anchors are selected by observed motion or occupancy; BAIR uses one motion-pooled token per sequence, at most one entry per episode, and excludes the current episode from retrieval.KTH adds spatial coordinates to regional forecast differences, while weather memory compares features at corresponding locations and weights loss by normalized cosine latitude.Source: A TASK-SPECIFIC MEMORY INSTANTIATIONS · p. 15
• For TaxiBJ, spatial average pooling and vectorization preserve horizon and channel order, and the observation notation is chronological.The specified groups are Hnear = {1, 2} and Hfar = {3, 4}; pooling and concatenation retain individual horizons within each group.Source: A TASK-SPECIFIC MEMORY INSTANTIATIONS · p. 15
• The same observation-key weights aggregate near- and far-horizon residual banks, and both targets are detached during student optimization.Source: A TASK-SPECIFIC MEMORY INSTANTIATIONS · p. 15
• The global-residual control changes both the retrieval key and horizon organization relative to observation-conditioned near/far prototypes, so their individual effects are not isolated.The control retrieves using the student’s normalized predicted residual and combines all future horizons into one vector.Source: A TASK-SPECIFIC MEMORY INSTANTIATIONS · p. 15

B DATA AND EVALUATION PROTOCOLS

• A time step denotes a sampled video frame, hourly temperature field, or traffic observation; image values are scaled to [0, 1] and temperature fields use training-set standardization.Table A2 summarizes the prediction tasks and protocols.Source: B DATA AND EVALUATION PROTOCOLS · p. 16, Table A2 · p. 16
• Video evaluation uses dataset-specific splits, including held-out subjects for KTH and Human3.6M and official split 1 with a class-stratified validation holdout for HMDB51.Moving MNIST and Moving FMNIST use online-generated training sequences and fixed test sets.Source: B DATA AND EVALUATION PROTOCOLS · p. 16
• BAIR uses episode-balanced training and all-window validation and testing, while KittiCaltech tests on Caltech Pedestrian recordings after training on disjoint KITTI recordings.KittiCaltech preprocessing samples every third decoded frame and center-crops to 128 × 160 pixels.Source: B DATA AND EVALUATION PROTOCOLS · p. 16
• WeatherBench separates 2010–2015, 2016, and 2017–2018 for training, validation, and testing; TaxiBJ uses the official test set after a purged temporal split of training windows.WeatherBench normalization uses statistics from the training period.Source: B DATA AND EVALUATION PROTOCOLS · p. 16
• Checkpoint selection is MSE-based, but four video datasets reuse evaluation data for validation; metric conventions distinguish elementwise errors from spatial/channel-sum errors.Relative improvements average per-seed relative changes, and multi-run deviations use the population convention.Source: B DATA AND EVALUATION PROTOCOLS · p. 16

C TRAINING CONFIGURATIONS AND HYPERPARAMETERS

• Table A3 documents configured training budgets and loss weights, while the listed configurations identify teacher–student pairs and keep teachers frozen in evaluation mode during student optimization.Output distillation has weight λ_out = 1; nonzero feature weights use standardized feature MSE.Source: C TRAINING CONFIGURATIONS AND… · p. 17, Table A3 · p. 17
• Video and traffic use Adam with OneCycleLR peaking at 10^-3, while WeatherBench uses Adam with an initial rate of 5 × 10^-3 and cosine decay.Feature adapters map student channels to teacher dimensions before spatial standardization and MSE.Source: C TRAINING CONFIGURATIONS AND… · p. 17
• Memory references use a 4,096-entry relational FIFO with 1,024 sampled references, while task-specific rules select motion, occupancy, location, episode, or flow-conditioned entries.WeatherBench retains 64 entries per location in Base configurations; TaxiBJ stores 2,048 entries and retrieves 16 neighbors.Source: C TRAINING CONFIGURATIONS AND… · p. 17

D SPATIAL FREQUENCY ANALYSIS PROTOCOL

• Spatial-frequency analysis samples 256 fixed test windows per dataset, evaluates every predicted frame, and uses the same windows across models and four paired runs.The analysis is descriptive, and overlapping windows are not treated as independent evidence of statistical significance.Source: D SPATIAL FREQUENCY ANALYSIS PROTOCOL · p. 18
• The spectral procedure removes a Hann-windowed weighted spatial mean before the orthonormal 2D Fourier transform; DC coefficients are excluded, so brightness discrepancies are omitted.Power is averaged over native channels and prediction times.Source: D SPATIAL FREQUENCY ANALYSIS PROTOCOL · p. 18
• Radial power is averaged across 48 frequency annuli, and band errors retain phase discrepancies while normalizing each band by its own ground-truth energy.The low, middle, and high bands span (0, 1/8], (1/8, 1/4], and (1/4, 1/2] cycles per pixel; greater power alone need not mean more accurate detail.Source: D SPATIAL FREQUENCY ANALYSIS PROTOCOL · p. 18

E MOTION AND TEMPORAL ANALYSIS PROTOCOL

• The analysis uses 1,024 evenly spaced windows per test set, validation-selected checkpoints, and four paired runs; overlapping windows are not independent statistical replicates.HMDB51 uses U-Net-Base, while KittiCaltech uses ResNet-FCN.Source: E MOTION AND TEMPORAL ANALYSIS PROTOCOL · p. 19
• Motion groups use dataset-specific tertiles of mean input change, while the high-change region is the top 20% of input-change pixels shared across models and future frames.The input proxy can also capture camera and illumination changes rather than semantic foreground or tracked objects.Source: E MOTION AND TEMPORAL ANALYSIS PROTOCOL · p. 19
• Temporal-change error compares predicted and target frame differences at each transition, measuring agreement with observed changes rather than visual smoothness alone.The reported HMDB51 reductions use the same 1,024 windows.Source: E MOTION AND TEMPORAL ANALYSIS PROTOCOL · p. 19
• Figure 6 reports HMDB51 temporal-change error reductions for TAM versus S0 of 4.74%, 1.68%, and 0.96% across successive transitions.The gains diminish at later transitions; KittiCaltech predicts one future frame and does not support this within-forecast analysis.Source: Figure 6 · p. 19, E MOTION AND TEMPORAL ANALYSIS PROTOCOL · p. 19