Fast Spatial Memory with Elastic Test-Time Training
| Authors | Ziqiao Ma et al. |
| Year | 2026 |
| HF Upvotes | 6 |
| arXiv | 2604.07350 |
| Download | |
| HF Page | View on Hugging Face |
Abstract
Large Chunk Test-Time Training (LaCT) has shown strong performance on long-context 3D reconstruction, but its fully plastic inference-time updates remain vulnerable to catastrophic forgetting and overfitting. As a result, LaCT is typically instantiated with a single large chunk spanning the full input sequence, falling short of the broader goal of handling arbitrarily long sequences in a single pass. We propose Elastic Test-Time Training inspired by elastic weight consolidation, that stabilizes LaCT fast-weight updates with a Fisher-weighted elastic prior around a maintained anchor state. The anchor evolves as an exponential moving average of past fast weights to balance stability and plasticity. Based on this updated architecture, we introduce Fast Spatial Memory (FSM), an efficient and scalable model for 4D reconstruction that learns spatiotemporal representations from long observation sequences and renders novel view-time combinations. We pre-trained FSM on large-scale curated 3D/4D data to capture the dynamics and semantics of complex spatial environments. Extensive experiments show that FSM supports fast adaptation over long sequences and delivers high-quality 3D/4D reconstruction with smaller chunks and mitigating the camera-interpolation shortcut. Overall, we hope to advance LaCT beyond the bounded single-chunk setting toward robust multi-chunk adaptation, a necessary step for generalization to genuinely longer sequences, while substantially alleviating the activation-memory bottleneck.
Engineering Breakdown
Plain English
This paper addresses a critical limitation in Large Chunk Test-Time Training (LaCT) for long-context 3D reconstruction: the model suffers from catastrophic forgetting and overfitting when updating weights during inference on arbitrarily long sequences. The authors propose Elastic Test-Time Training (ETAT), which stabilizes these inference-time updates by maintaining an anchor state that evolves as an exponential moving average, combined with Fisher-weighted elastic priors to balance learning plasticity against stability. They introduce Fast Spatial Memory (FSM), an efficient instantiation of this approach that enables handling of long sequences in a single pass without the previous limitation of requiring a single large chunk spanning the entire input. The method demonstrates stronger performance on long-context 3D reconstruction tasks while remaining computationally tractable at inference time.
Core Technical Contribution
The core innovation is adapting elastic weight consolidation principles to the test-time training setting, where the model must update rapidly during inference while maintaining performance across arbitrarily long sequences. Rather than locking weights after training (as in traditional EWC), ETAT introduces a dynamic anchor state that tracks an exponential moving average of past fast weights, providing a regularization target that evolves to prevent catastrophic forgetting. The Fisher information matrix weighted elastic prior enables the method to identify which weight dimensions are most critical to preserve, allowing selective plasticity—some dimensions can update aggressively while others remain constrained. This is fundamentally different from prior work like LaCT, which used fully plastic updates that inevitably degrade on long sequences without explicit stability mechanisms.
How It Works
The system maintains two weight states during test-time inference: the fast weights (being updated) and an anchor state (reference for regularization). At each test-time update step, the model computes gradients from the current loss and applies them to fast weights, but adds a Fisher-weighted regularization term that penalizes deviation from the anchor state. The anchor state itself is not static—it updates as an exponential moving average of the fast weights (θ_t = α·θ_fast + (1-α)·θ_anchor), where α is a momentum coefficient controlling how quickly the anchor adapts. The Fisher information matrix F is precomputed or estimated during training and weights the elastic penalty—dimensions with high Fisher values (important for the training task) are regularized more strongly, while dimensions with low Fisher values can drift more freely. Fast Spatial Memory (FSM) instantiates this by decomposing spatial memory into chunks that are processed sequentially, with the anchor state enabling consistent performance across chunk boundaries. The output at each step is the model's prediction conditioned on the current fast weights, which have been regularized to stay close to the evolving anchor.
Production Impact
For production 3D reconstruction systems handling long sequences, this approach eliminates the architectural constraint of requiring a single giant chunk—you can now process arbitrarily long inputs in a single forward pass with stable inference-time adaptation. The compute overhead is modest: precomputing Fisher information adds training cost, but test-time forward passes only incur the cost of one additional regularization term computation per weight update, typically 10-20% overhead compared to standard inference. The stability guarantee (formalized through the elastic prior) reduces the risk of deployment failures where models catastrophically degrade on out-of-distribution long sequences, a critical concern for production systems. However, integrating this requires maintaining two weight copies during inference (fast and anchor), increasing memory usage by roughly 2x the model size, which may be prohibitive for resource-constrained edge deployments. The exponential moving average anchor introduces a hyperparameter (momentum coefficient α) that must be tuned per application—different sequence lengths and domains may require different values, adding validation complexity.
Limitations and When Not to Use This
The paper does not address how to select the momentum coefficient α and Fisher regularization strength in a principled way without extensive tuning, making it unclear how to deploy FSM across diverse applications without per-domain validation. The method assumes Fisher information computed during training remains relevant for test-time updates—this assumption breaks down when test data distributions differ substantially from training, limiting applicability to truly out-of-distribution scenarios. The approach is specific to 3D reconstruction tasks and the paper provides limited evidence on whether ETAT generalizes to other long-context domains (language modeling, video understanding, etc.), leaving this as an open question. Additionally, maintaining an anchor state that evolves during inference could become stale or drift into invalid regions if α is poorly chosen, potentially causing subtle degradation that's hard to debug in production—the paper lacks failure mode analysis or guidance on detecting when the anchor has become misaligned.
Research Context
This work builds directly on Large Chunk Test-Time Training (LaCT) by addressing its core limitation of catastrophic forgetting during long-sequence inference. It draws inspiration from Elastic Weight Consolidation (Kirkpatrick et al., 2017), a continual learning technique that uses Fisher information to prevent forgetting when learning new tasks sequentially, creatively adapting that principle to the single-task-per-inference setting. The paper contributes to the broader research direction of test-time adaptation (TTA) and online learning under distribution shift, where models must improve or adapt during deployment without access to retraining. By enabling efficient long-context processing in a single pass, FSM advances the state-of-the-art in 3D reconstruction benchmarks and opens questions about whether similar elastic stabilization can improve other long-context models (transformers, state-space models, etc.) that suffer from attention or memory degradation at scale.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::