Compressing History into Memory: Distilling Transformers into Recurrent Transformers

Recurrent Transformers address computational limitations in long-horizon sequential data processing by distilling the compression strategy from full-history transformers into fixed-size memory. This work proposes a distillation approach that transfers the compression mechanism of a classical full-history transformer to a recurrent variant, enabling recurrent models to retain linear-time complexity while substantially narrowing the performance gap with full-history transformers.

The gist

We propose a distillation approach that transfers the compression strategy of a classical full-history transformer to a recurrent variant by designing a teacher model that explicitly compresses its observation history into a fixed-size bottleneck representation, and then directly supervising the student’s memory with this bottleneck representation.

Motivation and Problem Context

Transformers are powerful for sequential data but suffer from quadratic complexity, O(T2), when processing long sequences, making them impractical for applications like map-free pose estimation where storing and maintaining a history of observations is prohibitive. Recurrent Transformers (RTs) offer O(1) complexity at each step by maintaining a fixed-size memory bank, but their performance lags behind full-history Transformers (HTs). The authors argue this gap stems from differences in how they learn to compress past information: HTs implicitly compress history via attention and gradients, whereas RTs must explicitly decide at each time step what information to store in their fixed-size memory without having access to future requirements for prediction, making the learning problem significantly harder.

Proposed Method: Distillation Architecture

The core of the proposal is to bridge this gap by turning memory learning into a supervised compression task. The method involves two main components:

1. A Teacher model, the Latent Bottleneck History Transformer (LBHT), which explicitly compresses its observation history into a fixed-size bottleneck representation, similar to the Perceiver Resampler architecture. This teacher learns how to take a full sequence of observations and compress them into a memory bank in the form of contextualized tokens, interpreted as a latent representation, or latent 'map', of the scene.

2. A Student model, implemented as a Recurrent Transformer (RT), which maintains and updates a fixed-size memory state Mt using an update equation: Mt = fθ(Mt−1, ot). The student must learn a smart lossy compression mechanism to decide which information from the current observation ot needs to be retained in its limited memory.

Distillation Strategy

The distillation process aligns the two models by ensuring they share the same latent representation dimensions. Key steps include:

1. Aligning dimensions: The teacher’s latent representation B˜t and the student’s memory representation Mt are aligned using an identical number of embeddings and embedding dimensions, allowing direct distillation via an L1 loss between the flattened representations B˜t and Mt.

2. Segment-wise distillation: To best benefit from how the teacher learned to compress sequences of different lengths, a rollout of the student over a full sequence is decomposed into N segments. The final end memory bank Men for each segment is distilled with an L1 loss targeting the teacher representation B˜en obtained by a teacher having observed that specific sub-sequence.

Results and Findings

The approach, called Chimera for Compressing History Into Memory with Transformers, was validated on a long-horizon map-free pose estimation task. The experiments confirm two primary hypotheses:

1. "It is easier to learn to compress a full sub-sequence, i.e. the compression of a full sub-sequence, i.e., it is easier to learn to compress a full sub-sequence, i.e., that it is easier to learn to compress a full sub-sequence."

2. Once the LBHT teacher is trained and applied to sequences of increasing lengths T, thus obtaining sequences of bottleneck representations, these sequences can be sufficiently approximated through a learned recurrence function.

The resulting Chimera model retains the favorable O(1) inference complexity of recurrent models while achieving performance levels previously reserved for computationally expensive transformers. Chimera significantly outperforms the state-of-the-art recurrent models, including Kinaema, and performs close to the LBHT teacher across various query and memory timesteps. Furthermore, analysis shows that the distilled model develops a more skewed distribution of memory tokens, suggesting it learns to maintain a stable core of memory tokens that are updated less frequently while capturing discriminative information in a few highly dynamic tokens.

Architectural Details

The student architecture is based on the Kinaema model, featuring a recurrent memory update step and a decoding step that queries memory via cross-attention. The teacher (LBHT) uses 7 transformer layers compared to the student's 3 layers, and both architectures were optimized. Observation encoding for both models involves a fine-tuned DINO-v2 ViT-s backbone, where visual embeddings are down-projected to 64 values, concatenated with odometry embeddings.

Based on the provided paper, here are the specific improvements to AI systems that can be achieved by implementing its proposed architecture, Chimera:

1. Improve performance in long-horizon streaming vision and robotics applications (e.g., map-free pose estimation).

2. Enable agents to accurately answer complex visual queries about their environment relative to a moving coordinate frame (Where is this thing with respect to your current position?).

3. Achieve efficient, linear-time complexity inference for memory-based tasks, overcoming the quadratic complexity limitations of standard full-history Transformers in sequential processing.

4. Reduce the performance gap between recurrent models (Recurrent Transformers) and full-history Transformers by explicitly learning effective compression strategies rather than relying on implicit history access.

5. Develop agents that can maintain a fixed-size, relevant memory bank of past observations without needing to store or attend over the entire observation history, leading to more scalable and deployment-friendly models (O(1) update cost).

6. Improve the stability and generalization of learned memory representations by enforcing a compression mechanism distilled from a high-capacity teacher model (Latent Bottleneck History Transformer), resulting in memory states with better temporal stability.

7. Create models that learn to selectively retain critical information, allowing them to maintain a stable core of persistent memory tokens while capturing highly dynamic information in specialized, frequently updated tokens.

In summary, the improved AI system (Chimera) will be a more efficient and capable agent for real-world embodied AI tasks where long sequences of observations are common but full history processing is computationally prohibitive.