State-Space Blocks

A language model is a computer program that has learned to continue texts. Such models consist of many similar building blocks connected one after another. State-space blocks are a particular type of these building blocks. Their defining feature: they read a text from beginning to end and remember the important parts in a small internal memory. This memory is called the state, and that is where the name comes from. With each new word, the state is updated, but the old word itself is no longer looked at.

The idea originally comes from control theory, that is, from the design of controllers for machines. There, a system is described by its current state and a rule for how it changes. Since around 2021, this principle has been used for AI models. Well-known representatives are called S4 and Mamba.

The price of long texts

Almost all of today’s language models use a different architecture, the so-called Transformer. Its central mechanism compares every word with every other word in the text. This is very powerful, but quickly becomes expensive. If you double the text length, the computational effort quadruples. For very long inputs, this becomes a problem.

State-space blocks scale differently. They process each word once and update their state in the process. Twice as much text therefore means roughly twice as much work, not four times as much. Memory requirements even stay almost constant, because the state is always the same size. This is exactly why research and industry are interested in these blocks.

This becomes practically relevant for tasks with very long inputs. An entire book, an hour of audio, or a genome segment can easily comprise hundreds of thousands of units. For Transformers this is expensive, for state-space models it is well manageable. In return, Transformers have so far been better at retrieving a single detail from the middle of a long text exactly.

State, filter, and the role of training

You can picture the state as a notepad of fixed size. When reading each new word, the model writes something down and deletes something else in return. How strongly old information fades and how strongly new information is weighted is stored in numerical values that were learned during training. Nobody manually decides what the model should remember.

This computation rule is essentially a simple formula that is applied over and over again. This sounds slow, because each word has to wait for the previous one. During training, however, there is a mathematical trick: the entire sequence can be rewritten into a form that graphics cards can process in parallel. Only this made state-space blocks practically usable in the first place.

An important advance came in 2023 with Mamba. There, the numerical values for remembering and forgetting depend on the current word. The model can therefore skip unimportant filler words and look more closely at important passages. Older variants treated every word according to the same rigid rule.

Mamba, hybrids, and what’s in the news

State-space blocks are not directly visible to users. They sit inside models, similar to how the engine type sits inside a car. But anyone reading news about new AI models regularly comes across names like Mamba, Jamba, or Falcon Mamba. Nvidia and several research labs have also released models with such blocks.

In practice, many developers rely on hybrid forms. A model then contains predominantly state-space blocks with a few Transformer layers interspersed. The former provide speed for long texts, the latter provide precise lookup of individual passages. Such hybrids are currently considered the most promising path.

A common misconception is that state-space blocks will soon replace Transformers. So far, little points to this. They are especially interesting where computing time and power consumption matter, for example on phones or in data centers with many requests. Whether they become widely adopted depends on whether they deliver similarly good answers at the same size.