Contextual windows are growing, but there are two chairs: either classic attention, which feeds on memory and computes like crazy, or RNN-like Mamba, DeltaNet, which work quickly but start drifting and losing details in long contexts.
🟢 What NVDIA Proposed?
NVIDIA proposes a solution that tries to sit on both chairs at once - Test-Time Training with End-to-End formulation (TTT-E2E):
Usually, the model's weights are frozen after training. When you feed it data, it just holds it in the KV cache. In TTT, everything is different: the context is the training dataset itself. While the model reads your prompt (the context), it updates its weights (more precisely, performs gradient descent on the fly), thereby embedding the context's information into the model itself. This allows you to compress gigantic volumes into a fixed state size without bloating the KV cache to the sky.
➜ result - beauty and magic:
☞ Latency of inference becomes constant. It doesn't matter if there are 100 tokens in the context or a million - the time to generate the next token is the same.
☞ On a context of 128k tokens - a 2.7x speedup compared to Attention (on H100). On 2M tokens - a 35x speedup.
☞ Unlike Mamba and other RNNs, the quality doesn't degrade over long distances. TTT maintains the same level as full attention.
➜ Of course, there are a bunch of points with an asterisk
☞ Training is complex. To allow the model to learn on the fly so skillfully, it needs to be pretrained specially. This process is currently 3.4x slower than regular training.
☞ method requires calculating gradients from gradients during training. FlashAttention currently doesn't support this out of the box, requiring custom kernels or workarounds.
☞ process of consuming context during inference itself requires computations during the prefill phase.
In the end, NVIDIA compares RAG to a notebook, and its TTT to the real updating of neural connections in the brain. If you want to delve into the methodology and grasp the idea - the code and the paper are publicly available.
GitHub | Paper • #AI #ML #LLM #TTTE2E #NVIDIA
••••••••••••••••••••••••••••••••••••••
🤖 Data Science, ML & Big Data with @DataXplore