Orthogonalizing Recurrent Memory Beats the Noise

Transformers dominate modern sequence modeling because of their associative recall capabilities. By using attention, every token has direct, unhindered access to every preceding token. But this capability comes with a well-known tax: quadratic $O(N^2)$ computational complexity. For domains like long-horizon reinforcement learning, where agents must process massive histories over extended environments, the quadratic overhead of attention is often too expensive.

Developers looking for linear-scaling alternatives have turned back to recurrent architectures. The most promising of these is the mLSTM, a variant of the classic Long Short-Term Memory network that features a matrix memory. While mLSTMs excel at clean associative recall benchmarks, real-world environments are rarely clean. They are noisy, filled with irrelevant transitions and distractor tokens. When evaluated on Noisy Associative Recall (NAR) tasks, standard recurrent models struggle. Distractor tokens crowd out the actual signal, causing the network to forget crucial key-value associations.

Recent research shows that a simple mathematical intervention, matrix orthogonalization applied strictly during the read phase, can rescue recurrent models from this memory decay.

The Problem with Noisy Recall

To understand why recurrent models fail in noisy environments, consider a typical task from the Memorization and Association Defender (MAD) benchmark. The model is fed a sequence of keys, values, and distractors:

0 9 3 10 12 13 15 14 0 9 5 8 2 9

Here, key 0 maps to value 9, and key 3 maps to 10. The tokens 12 through 15 are distractors. To succeed, the model must predict 9 when it sees the key 0 again at the end of the sequence, completely ignoring the interleaved noise.

In a standard mLSTM, the memory matrix acts as a storage buffer. As new tokens are processed, they write to this matrix. However, when noisy distractor tokens are written to the same memory space, they can easily overwrite or crowd out the weaker, older memories.

This is where orthogonalization comes in. Inspired by Muon, an optimizer that orthogonalizes its momenta to act as an equalizer of represented directions, researchers experimented with orthogonalizing the mLSTM memory matrix. By preventing a few dominant directions from taking over the update space, orthogonalization lifts up the weaker, older memories, ensuring they are not lost in the noise.

The Read-Only Orthogonalization Trick

Applying orthogonalization to a recurrent network is not a new idea, but the implementation details matter. Historically, researchers tried to constrain the recurrent weight matrices to be orthogonal throughout training to prevent vanishing and exploding gradients.

However, the breakthrough for NAR tasks relies on a specific twist: the memory matrix is orthogonalized only during reads, and this modified matrix is never written back to the state. Writing the orthogonalized matrix back to the recurrent state actually degrades performance.

To implement this, the readout process follows a specific sequence:

1. Normalize the memory matrix using the Frobenius norm.
2. Apply five Newton-Schulz iterations to orthogonalize the normalized matrix.
3. Allow gradients to flow through this entire orthogonalization process during backpropagation.
4. Use the resulting orthogonalized matrix for the read operation, while keeping the original, un-orthogonalized matrix in the recurrent state for the next step.

This approach yields massive performance gains. On difficult MAD tasks with high noise levels (where 80% of the sequence consists of distractors), orthogonalization brings mLSTMs back from the brink of failure.

xychart-beta
    title "Validation Accuracy: Orthogonalized vs Baseline mLSTM"
    x-axis ["V80 L1024 Ortho", "V80 L1024 Base", "V96 L768 Ortho", "V96 L768 Base", "V96 L1024 Ortho", "V96 L1024 Base"]
    y-axis "Accuracy (%)" 0 --> 100
    bar [98.5, 83.3, 62.4, 22.0, 68.5, 23.1]

As the task difficulty scales (larger vocabulary sizes and longer sequence lengths), the performance gap widens. In the most challenging regime tested (vocabulary size 96, sequence length 1024), the baseline mLSTM achieved a meager 23.1% accuracy, solving only 4 out of 24 random seeds. The orthogonalized variant achieved 68.5% accuracy, successfully solving 16 out of 24 seeds.

The Developer Angle: Implementation and Trade-offs

For machine learning engineers, implementing this technique in PyTorch is straightforward. The core of the method is the Newton-Schulz iteration, which iteratively refines a matrix toward orthogonality without requiring expensive singular value decomposition (SVD) operations.

Here is a clean PyTorch implementation of the read-only orthogonalization step:

import torch
import torch.nn as nn

def newton_schulz_orthogonalize(M, num_iters=5, eps=1e-6):
    """
    Orthogonalizes the input matrix M using Newton-Schulz iterations.
    M: Tensor of shape (..., rows, cols)
    """
    # Normalize by the Frobenius norm
    fro_norm = torch.linalg.norm(M, ord='fro', dim=(-2, -1), keepdim=True)
    X = M / (fro_norm + eps)
    
    # Identity matrix matching the last dimension
    I = torch.eye(X.size(-1), device=X.device, dtype=X.dtype)
    
    # Newton-Schulz iterations: Y_{k+1} = 0.5 * Y_k * (3 * I - Y_k^T * Y_k)
    for _ in range(num_iters):
        X_T_X = torch.matmul(X.transpose(-2, -1), X)
        X = 0.5 * torch.matmul(X, 3.0 * I - X_T_X)
        
    return X

class OrthogonalizedReadout(nn.Module):
    def __init__(self):
        super().__init__()
        
    def forward(self, memory_matrix, query):
        # Orthogonalize the memory matrix for the read operation only
        ortho_memory = newton_schulz_orthogonalize(memory_matrix)
        
        # Perform the read operation (e.g., query-key-value lookup)
        output = torch.matmul(query, ortho_memory)
        return output

There is no free lunch. The primary trade-off here is computational. Applying five Newton-Schulz iterations at every step introduces additional FLOPs and increases wall-clock training time. However, it does so without increasing the parameter count of the model. If you are parameter-constrained but have FLOP budget to spare, this is an incredibly efficient trade.

The Theoretical Tension: Orthogonal vs. Non-Normal Dynamics

This read-only approach also highlights an interesting theoretical debate in recurrent network design. Historically, maintaining strict orthogonality in RNNs was seen as the gold standard for preserving signal strength. But as detailed in research published on arXiv, pure orthogonality has its downsides.

In non-linear, noisy networks, orthogonal transformations are not always optimal. Because orthogonal matrices are isometries, they preserve Euclidean norms, but they do not necessarily maximize the signal-to-noise ratio (SNR). In fact, networks that maximize SNR often exhibit strongly "non-normal" dynamics, where the recurrent connectivity matrix does not commute with its transpose. These non-normal networks can outperform purely orthogonal ones on a variety of sequential tasks because they allow for transient amplification of signals, which helps pull the target data out of the noise.

This explains why writing the orthogonalized matrix back to the mLSTM memory state degrades performance. By forcing the state itself to remain orthogonal, you strip the network of its ability to utilize non-normal dynamics for signal amplification.

By keeping the orthogonalization strictly on the read path, you get the best of both worlds. The write path remains unconstrained, allowing the network to build up highly expressive, non-normal dynamical representations. Meanwhile, the read path is forced to be orthogonal, acts as an equalizer, and ensures that no single direction dominates the readout, protecting weaker memories from being drowned out by noise.

While these results are currently limited to synthetic benchmarks in small model regimes, the underlying mechanics are robust. For developers building recurrent models for noisy, long-horizon tasks, adding a read-only orthogonalization step is a low-risk, high-reward experiment.