FreeSolo: Linear-time Sparse Attention with Zero-Shot Length Extrapolation

FreeSolo. Linear-time Sparse Attention with Zero-Shot Length Extrapolation. By Zekun Shi and one other, from the National University of Singapore, the University of Toronto, and one other institution. Published on the twelfth of October. Its code repository is worked out: about one hundred source files, with a readme and a licence.

Linear-time sequence models typically keep long-context information in a recurrent state, and a fixed-size state limits how much can be stored. Hybrid architectures add full attention to recall better, but pay a quadratic cost in sequence length. Sparse attention with continued pretraining can extend the effective context instead. Generalizing past the trained length is still hard, though, and continued pretraining can be costly.

FreeSolo is a native sparse attention mechanism with two branches. A short-range branch uses a sliding window. A long-range branch selects keypoint cells and then attends directly to their tokens. Nearby context and distant context are handled by different routes.

It comes in two variants. The linear-time one, FS-L, learns to prune the prefix before the query-dependent selection happens, and it learns this through a short round of continued pretraining. The other, FS-B, searches balanced trees, at an algorithmic cost of order N times log squared N, and that figure includes building the structure with radix sort.

The reported result is about length. Without any long-context fine-tuning, FreeSolo shows better zero-shot length extrapolation than the baselines.

The aim is to reach distant context at close to linear cost, and to keep working beyond the lengths seen in training, with only a brief stretch of continued pretraining needed for the linear-time version and none of the long-context fine-tuning that the problem usually calls for.