PHOTON: Hierarchical Autoregressive Modeling for Lightspeed and Memory-Efficient Language Generation

TL;DR

PHOTON introduces a hierarchical autoregressive model that significantly improves language generation speed and memory efficiency by replacing horizontal token scanning with vertical multi-resolution context scanning, especially benefiting long-context tasks.

Contribution

The paper presents PHOTON, a novel hierarchical autoregressive model that reduces inference latency and memory usage by using vertical context scanning and recursive generation, outperforming traditional transformers.

Findings

Up to 1000x higher throughput per unit memory.

Superior long-context and multi-query task performance.

Reduces decode-time KV-cache traffic.

Abstract

Transformers operate as horizontal token-by-token scanners; at each generation step, attending to an ever-growing sequence of token-level states. This access pattern increases prefill latency and makes long-context decoding more memory-bound, as KV-cache reads and writes dominate inference time over arithmetic operations. We propose Parallel Hierarchical Operation for TOp-down Networks (PHOTON), a hierarchical autoregressive model that replaces horizontal scanning with vertical, multi-resolution context scanning. PHOTON maintains a hierarchy of latent streams: a bottom-up encoder compresses tokens into low-rate contextual states, while lightweight top-down decoders reconstruct fine-grained token representations in parallel. We further introduce recursive generation that updates only the coarsest latent stream and eliminates bottom-up re-encoding. Experimental results show that PHOTON is superior to competitive Transformer-based language models regarding the throughput-quality trade-off, providing advantages in long-context and multi-query tasks. In particular, this reduces decode-time KV-cache traffic, yielding up to $10^{3}\times$ higher throughput per unit memory.