Layerwise learning improves nonlinear memory in recurrent networks

DeltaTTT: Layerwise Optimization for Nonlinear Recurrent Memory

Machine LearningArtificial IntelligenceComputation and Language

Summary

Some computer models remember past information to make better predictions, but when these memories get complicated, they’re hard to train effectively. The authors found that updating the whole memory at once doesn’t always work well for these complex memories. They created a new way called DeltaTTT that updates the memory layer by layer, making it easier to optimize and faster to compute. This method improved performance in tasks like predicting language and finding related information.

What this means in practice

  • For language model developers: Improve language prediction models by applying layerwise updates to nonlinear recurrent memory components.
  • For search engine engineers: Enhance information retrieval systems using nonlinear recurrent memories optimized with layerwise delta rules for better accuracy.

Authors

Yining Li, Dongchen Han, Jie Fu, Gao Huang

Abstract

Sequential test-time training adapts a memory network through successive updates, each computing an inner-loop gradient based on the network's previous state. Intuitively, this state dependence should allow each update to account for what the memory has already learned and better incorporate new information. However, we find that this expected advantage does not consistently materialize in nonlinear memories: a fixed-base parallel TTT baseline outperforms its serial counterpart. Our exploratory experiments point to a key underlying difficulty: nonlinear memories can be harder to optimize than linear ones within a single pass over the sequence. To alleviate this optimization difficulty, we introduce DeltaTTT, which replaces joint inner-loop optimization of a two-layer memory network with layerwise learning. Each layer is assigned a local prediction target and updated through a state-dependent delta rule. This formulation retains a nonlinear readout while enabling chunkwise parallel computation. Experiments on DeltaNet and LaCT backbones show improvements in language modeling and retrieval over their recurrent baselines.