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

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

Traditionally, we’ve been using tokenizers to represent input for language models. Honestly, that has always felt a bit… artificial to me. The granularity is synthetically enforced, and the model gets stuck with a predefined vocabulary.

How is the model supposed to do symbol-level reasoning? (Just think about how car and cars relate, for instance.) Or what happens with dialects, rare languages, or funky edge cases? Well… tokenization with BPE is a pretty hard constraint.

In this awesome paper: From Bytes to Ideas:Language Modeling with Autoregressive U-Nets, by Mathurin Videau, Badr Youbi Idrissi, Alessandro Leite, Marc Schoenauer, Olivier Teytaud and David Lopez-Paz, they propose a very interesting alternative: U-Nets

🔍 U-Nets are so back!

I loved revisiting U-Nets when learning about Stable Diffusion. (The diffusion world is honestly awesome.) And as it turns out, U-Nets absolutely have a place in modern-modern ML. In this paper, they propose using them in an autoregressive way

The idea is processing raw bytes and learn the hierarchical token representation (words, pairs of them, up to four), creating a multo-scañle, learned view of the sequence. In order to get taht functioning, there is a technical decision that stands out: the “pooling points” are identified by a splitting function (regex/spaces, etc). This is not the most “canonical” way, but it is definitely something that works

🏗️ How does it work?

• Instead of being locked into a single “token” (like a fixed word or sub-word piece) as traditional models are, AU-Net processes raw bytes
• It then simultaneously learns to “pool” these bytes into a hierarchy of increasingly larger, more abstract “chunks”
• Deeper Stages = Broad Semantics/Strategic View: The “deeper” layers of the AU-Net model work on these larger, more compressed chunks. Their job is to predict further ahead into the future
• Shallower Stages = Fine Details/Operational Execution: Conversely, the “shallower” layers operate on the finer details, like individual bytes or words. They are responsible for handling the precise spelling and character-level accuracy, guided by the broader semantic patterns established by the deeper stages
• There are two paths the model takes: there is a “contracting path” on which it takes the raw bytes and compresses them to the higher level representations, and the “expanding path”, for reconstructiong the original sequence
• Keep in mind there are skip-connections from one path to the other, so that no crucial info is lost on compression
• AU then learns representations in a similar way LMs do: by predicting the next information unit. In this case, from the bytes

Going AU

When learning, the hierarchically deeper AU-Net layers are tasked with predicting “further into the future”, as in guessing next words, similarly to the masked learning we all know. Those predictions at deep layers then “guide shallower stages at the semantic level”, allowing the earlier stages to handle detalis like spelling. So, the representations for words and chunks are learned through this process.

At training time, the model would process the whole input sequence in parallel across all stages, but at inference time, it would work as a true autoregressive model. This “cascading conditional activation” would increase coherence and efficiency at text generation level

Results

AU_Nets showed compelling performance against Byte Pair Encoding (BPE) baselines and notable advantages in multilingual contexts. When carefully tuning and controlling pre-training compute, AU-Net models consistently match or outperform BPE-based counterparts, with multi-stage AU-Net models.

AU-Net initially underperformed on knowledge-intensive tasks like TQA, this discrepancy tends to vanish with increased model size and training data. On scaling laws, AU-Net (2 and 3 stages) can match the performance of the BPE baseline for tasks like Hellaswag, ARC Easy, and NQ, and the 3-stage model catches up to BPE on TQA at higher compute levels

Furthermore, on the FLORES-200 benchmark, AU-Net showed consistent gains in translation tasks into English for regional and low-resource languages, demonstrating its ability to generalize based on subword morphology and shared linguistic roots without requiring these languages in the tokenizer or training corpus.

However, a limitation observed is its lower Chinese MMLU scores compared to BPE baselines, indicating that its current rule-based splitting function, designed for space-based languages, is not well-suited for non-space-based languages

🤔 Final thoughts

I think this is a pretty exciting step. Moving beyond rigid tokenization feels like an overdue direction, and combining it with the multiscale power of U-Nets? Very clever. Definitely curious to see how this scales in practice — and whether it becomes the backbone of future multilingual or low-resource LLMs.