SpatialLM-Mini sped up training with a 16K context by 2.47× on an RTX 3050

Long context is expensive even before a model starts delivering value. With full attention, each new token can attend to all preceding tokens, and computation grows quadratically. The SpatialLM-Mini experiment restricted that access to nearby positions and a few long-range connections. On an RTX 3050 with 8 GB of memory, the prototype trained 2.47× faster than the baseline Transformer with a window of 16,384 tokens.

For developers, there is a concrete takeaway: reducing the number of connections can speed up experiments with long sequences, even on an inexpensive graphics card. But the result applies to a small model, a single corpus, and a particular compute kernel. The benefit of the specifically “spatial” method of calculating attention remains an open question.

How tokens choose whom to listen to

Imagine a long document. With standard causal attention, a token at the end can directly take any preceding position into account. In SpatialLM-Mini, the list of accessible positions is defined in advance: the token itself, several nearby predecessors, and positions at distances of 1, 2, 4, 8, and so on. Looking ahead is prohibited.

The model then decides how much attention to give each permitted neighbor. To do this, it transforms the hidden state into learnable three-dimensional query and key coordinates. The shorter the distance between them, the higher the connection weight can be; additional learnable parameters adjust this calculation.

The coordinates serve as a way to calculate weights. The model does not search the entire document for the nearest tokens or use the graphics card’s RT cores. The neighbor list is fixed, while the distribution of attention within that list depends on the text.

This graph reduces computation but limits direct access to distant facts. The author therefore also tested a hybrid: two spatial blocks alternate with two full-attention blocks. This allows some layers to attend to the entire preceding context.

What the comparison on 8 GB showed

Three models with approximately 9.3 million parameters each were trained from scratch on Jane Austen’s Pride and Prejudice. Each had four layers, the same context length, and three runs with different random seeds. Only the tokenizer was taken from GPT-2.

With a budget of 180 seconds of measured training, the baseline Transformer processed 13,126 tokens per second, the spatial model 32,364, and the hybrid 18,725. The hybrid achieved the lowest average loss on the test split. However, the faster models processed more tokens and completed more updates in that time.

When compared after the same 100 steps, the difference in loss narrowed considerably. The spatial model needed 50.64 seconds, versus 125.10 seconds for the Transformer. The main demonstrated benefit of this experiment is faster updates on a long sequence.

Memory told a different story: peak memory allocated by PyTorch rose from approximately 499 to 671 MiB. The sparse implementation creates intermediate tensors when gathering neighbors, while full attention uses an optimized kernel. These figures do not represent the computer’s total occupied GPU memory.

What to check before adopting it in your project

The baseline Transformer used PyTorch Efficient Attention. FlashAttention was unavailable in the Windows build used. The speedup therefore cannot automatically be carried over to another system or to comparisons with other kernels.

There is also a more substantial gap: both the set of connections and the method of calculating their weights were changed at the same time. To identify the contribution of the coordinates, a comparison is needed using the same sparse graph with a standard dot product or fixed weights. For now, the advantage may primarily come from reducing the number of interactions.

Accepting 16 thousand tokens also does not prove that the model can retrieve the relevant fact from the beginning of a document. That requires a separate test. Generation speed was not measured, the prototype has no KV cache, and brief training on a single novel cannot establish its suitability for production tasks.

The practical next step is to repeat the comparison on your own corpus and hardware, then add a baseline using the same graph without coordinates. According to the author’s description, the code, configurations, and original measurements are available in the SpatialLM-Mini repository. For a team exploring low-cost training with long contexts, it is a useful starting point for an experiment of its own.