From Bytes to Ideas: Language Modeling with Autoregressive U-Nets

TL;DR

This paper introduces an autoregressive U-Net that learns to embed its own tokens from raw bytes, enabling flexible, multi-scale language modeling that adapts to different tasks and languages.

Contribution

It proposes a novel U-Net based architecture that learns tokenization during training, moving beyond fixed schemes like BPE, and demonstrates its effectiveness across various language modeling tasks.

Findings

Shallow hierarchies match strong BPE baselines.

Deeper hierarchies show promising trends in performance.

Model handles character-level tasks and low-resource languages.

Abstract

Tokenization imposes a fixed granularity on the input text, freezing how a language model operates on data and how far in the future it predicts. Byte Pair Encoding (BPE) and similar schemes split text once, build a static vocabulary, and leave the model stuck with that choice. We relax this rigidity by introducing an autoregressive U-Net that learns to embed its own tokens as it trains. The network reads raw bytes, pools them into words, then pairs of words, then up to 4 words, giving it a multi-scale view of the sequence. At deeper stages, the model must predict further into the future -- anticipating the next few words rather than the next byte -- so deeper stages focus on broader semantic patterns while earlier stages handle fine details. When carefully tuning and controlling pretraining compute, shallow hierarchies tie strong BPE baselines, and deeper hierarchies have a promising trend. Because tokenization now lives inside the model, the same system can handle character-level tasks and carry knowledge across low-resource languages.