End-to-End Test-Time Training for Long Context in JAX

Repositorytest-time-training/e2eResearch

This repository provides the official JAX implementation of End-to-End Test-Time Training (TTT) for long contexts. It approaches long-context language modeling as a continual learning challenge, using a standard Transformer with sliding-window attention that compresses input contexts into its weights at test time via next-token prediction.

Key points

  1. Frames long-context language modeling as continual learning rather than architectural redesign.
  2. Employs a standard Transformer architecture with sliding-window attention.
  3. Compresses context directly into model weights at test time using next-token prediction.
  4. Uses meta-learning at pre-training time to optimise parameter initialisation for test-time adaptation.
  5. Built with JAX, package-managed via uv, and configured through Hydra and Submitit.

Why it matters

Instead of relying on increasingly large attention matrices or custom architectures, E2E TTT adapts model weights on the fly during inference, offering an alternative mechanism for handling expansive contexts.

Analysis

The project tackles the computational and representational limits of long-context LLMs. By framing context retention as test-time continual learning, a sliding-window Transformer updates its parameters end-to-end on input tokens. The tech stack includes JAX, CUDA 12.8, cuDNN 9.8, NCCL 2.26, uv for dependency management, Hydra for configuration, Submitit for multi-node Slurm execution, and Weights & Biases for logging. Datasets are tokenised with Llama-3. The repository is aimed at AI researchers and machine learning practitioners investigating meta-learning and inference-time adaptation. In terms of maturity, it is a research-grade artifact built to replicate paper experiments (such as a 125M parameter run); it currently lacks an explicit open-source license.

Found in

Summary generated by AI (Gemini 3.8 Flash, 27 September 2026). It may contain errors – the original source is authoritative.

Similar in content, found by the AI search.

  • Repository:volotat/mini-AGI

    mini-AGI: Byte-Level Continual Learning on 8 GB VRAM Consumer Hardware

    mini-AGI is an experimental open-source byte-level language model designed to demonstrate continual learning from a single stream of data (batch size 1) on modest 8 GB VRAM hardware. By paging model weights from disk into GPU memory as needed, it decouples parameter capacity from VRAM limits while dynamically scaling its architecture.

    793starsPython

    AI & LLMs· Research

  • Repository:vectorize-io/hindsight

    Hindsight: Agent Memory That Learns

    Hindsight is an MIT-licensed memory system designed to help AI agents learn over time. Rather than merely recalling conversation histories, it aims to overcome the limitations of conventional RAG and knowledge graphs on long-term memory tasks.

    41.3kstarsPython

    AI & LLMs· Tool

  • Article:Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp, Günter Klambauer, Johannes Brandstetter, Sepp Hochreiter

    xLSTM: Extended Long Short-Term Memory

    A research team led by Sepp Hochreiter introduces xLSTM, an extension of the original Long Short-Term Memory architecture designed for large-scale language modelling. The approach addresses historical scaling limits of LSTMs and demonstrates competitive performance against state-of-the-art Transformers and State Space Models.

    AI & LLMs· Research

  • Repository:hhhuang/CAG

    Cache-Augmented Generation (CAG): A KV-Cache Alternative to RAG

    The CAG open-source repository introduces Cache-Augmented Generation as an alternative to Retrieval-Augmented Generation (RAG). By preloading the entire knowledge base into the context window and reusing the precomputed KV cache, it bypasses real-time document retrieval entirely.

    1.5kstarsPython

    AI & LLMs· Resource

  • Repository:fudan-generative-vision/Remory

    Remory: Learning Residual Memory for Context Compaction in LLM Agents

    Remory enhances standard text summaries in LLM agents with learned soft memory extracted from execution history. Conditioned on the summary and appended to it, this residual memory enables frozen base models to sustain continuity across context compactions.

    236starsPython

    AI & LLMs· Tool

  • Article:Opher Lieber, Barak Lenz, Hofit Bata et al.

    Jamba: A Hybrid Transformer-Mamba Language Model

    Researchers introduce Jamba, a base large language model utilizing a hybrid Transformer-Mamba mixture-of-experts (MoE) architecture. By interleaving Transformer and Mamba layers with MoE modules, Jamba fits into a single 80GB GPU while supporting context lengths of up to 256K tokens.

    AI & LLMs· Research

Let's talk about your project

Locations

  • Mattersburg
    Johann Nepomuk Bergerstraße 7/2/14
    7210 Mattersburg, Austria
  • Vienna
    Ungargasse 64-66/3/404
    1030 Wien, Austria

Parts of this content were created with the assistance of AI.