Adaptive Tokenization and Memory in Foundation Models for Efficient and Long-Horizon AI (AToM ⚛︎)

The recent revolution in generative AI has been driven by the rapid growth in training and inference compute for Foundation Models (FMs). This scaling, however, brings unintended consequences, including high energy demand and associated carbon emissions. My research sets out to reverse this trend by targeting a fundamental inefficiency in dominant FM architectures such as Transformers. Currently, these map data into sequences of internal representations, whose length bottlenecks both prompt processing (compute-bound by typically quadratic attention) and output generation (memory-bandwidth bound by key–value cache reads). Yet this length is largely determined upfront by input segmentation (tokenisation) and typically remains fixed across layers; during decoding, it also accumulates unchanged in the key–value cache.

I have prototyped a new class of FM architectures that learn, end-to-end, to compress sequences of internal representations, effectively redefining the model’s “atomic units” for processing and memorising information. To accelerate adoption, I have repurposed existing state-of-the-art open-weight FMs (such as Qwen 3 in collaboration with NVIDIA and OLMo in collaboration with AI2) into adaptive variants and released them publicly.

This leads not only to substantial gains in efficiency (with 8× speedups without accuracy degradation) but also to the emergence of new capabilities: adaptive FMs can operate over broader effective horizons, as they can perceive longer inputs and generate longer outputs under a fixed budget. This enables (1) lifelong learning via a permanent, sub-linearly growing memory, (2) inference-time hyper-scaling for reasoning-intensive tasks (maths, science, coding), and (3) enhanced world modelling for multimodal planning and simulation. Adaptive FMs thus open a path towards greener, more capable generative AI.

Models

Hugging Face checkpoints

allenai/Bolmo-7B

State-of-the-art, fully open-source large language model with latent tokenization. Available in 1B and 7B sizes.

nvidia/Qwen3-8B-DMS-8x

8x KV cache compression without quality degradation. Ideal for inference-time scaling.

Publications

Related papers

Current hierarchical attention methods, such as NSA and InfLLMv2, select the top-k relevant key-value (KV) blocks based on coarse …

Sparse attention has been proposed as a way to alleviate the quadratic cost of transformers, a central bottleneck in long-context …

Internal modelling of the world — predicting transitions between previous states X and next states Y under actions Z — is essential to …

Recent advances in generative AI have been largely driven by large language models (LLMs), deep neural networks that operate over …

Multi-token prediction (MTP) is a prominent strategy to significantly speed up generation in large language models (LLMs), including …

Join AToM ⚛︎

2 PhD studentships and 1 postdoctoral position

AToM is funded by an ERC Starting Grant and hosted at the Department of Computing of Imperial College London. I am recruiting two PhD students and one postdoctoral research associate. All positions are fully funded.

Research directions

Each position will focus on one or more of the project’s 4 main foci:

  1. Latent tokenization: end-to-end learned segmentation of raw, modality-agnostic data into hierarchical units (prototyped in Dynamic Token Pooling and brought to scale in Bolmo).
  2. Permanent adaptive-size memory and sparse attention: learned compression of the memory for lifelong learning and long-context inference (see Dynamic Memory Compression, Dynamic Memory Sparsification, The Sparse Frontier, and DashAttention).
  3. Retrofitting of SoTA models into adaptive architectures and public release of data, recipe, and models (see the released checkpoints, such as Qwen3-8B-DMS-8x with NVIDIA and Bolmo-7B with AI2).
  4. Inference-time hyper-scaling and long-horizon world modelling: using the freed compute budget for reasoning-intensive tasks (see Inference-Time Hyper-Scaling) and multimodal planning and simulation (see Self-Improving World Modelling and Action-Grounded Visual Dynamics).

What we offer

  • Extensive compute: the project’s own B200 GPUs and cloud credits, plus H200s from the Department of Computing and the Faculty of Engineering.
  • Close collaboration with industry and academic partners (NVIDIA, AI2) and a track record of releasing retrofitted open-weight models.
  • Funding for two international conferences per year.
  • For the postdoc: freedom to define an independent agenda within the project, to co-supervise PhD and MSc students, and to deliver tutorials at conferences & summer schools.

Postdoctoral position

Research Associate in Adaptive and Efficient LLM Architectures (job number ENG04037), Department of Computing, South Kensington campus.

  • Contract: full-time, fixed term for 2 years, starting around winter 2026/27 (flexible).
  • Salary: £50,733 to £59,484 per annum (Research Assistant scale, £45,399 to £48,876, if the PhD is not yet awarded).
  • Closing date: 30 September 2026.
  • Requirements: a PhD in computer science, a strong publication record in ML/NLP venues, and hands-on experience with LLM training and evaluation in PyTorch or JAX. Familiarity with CUDA or Triton kernels and inference engines is desirable.

Apply through the official Imperial job posting, attaching a full CV with publication list and a research statement (max 2 pages) on the research questions in adaptive and efficient architectures that you find most exciting and why your expertise is relevant. You are also welcome to fill in the expression-of-interest form so we can get in touch early.

PhD studentships

PhD applications go through the Department of Computing’s central admissions. Please follow the PhD application guidelines, select ‘Computing Research’ as a degree, and name me as your proposed supervisor. In parallel, fill in the expression-of-interest form (attaching your CV, a link to your best piece of work, and a paragraph on which of the 4 foci above interest you the most). I review submissions on a rolling basis and reach out to promising candidates before the formal deadline.

Process and timeline

  • Postdoc: applications close on 30 September 2026. Shortlisted candidates will be invited to an interview in October, with a research talk and a technical discussion, and the outcome communicated shortly after. The start date is flexible around winter 2026/27.
  • PhD: I read expression-of-interest submissions on a rolling basis and reach out to promising candidates to invite formal applications. Formal offers follow Imperial’s admissions and scholarship timelines and eligibility requirements.
  • I try to acknowledge every form submission that is a plausible fit. If you have not heard back within three weeks of submitting, assume I could not take the application further.

If you have any specific questions, email me with [phd-atom] or [postdoc-atom] in the subject line. However, please note that due to the volume of request, despite my best effort, I am not always capable of replying to email enquiries.