about
GateLoop: Data-Controlled Linear Recurrence for Sequence Modeling (arxiv.org)
26 points by zxexz on Dec 7, 2023 | hide | past | pdf | 2 comments on HN

In plain words: Instead of just adding each new word to a running memory, this model lets the input also decide how that memory is scaled and turned at every step. That fuller control beat earlier step-by-step memory models at predicting text.

Abstract · GateLoop: Fully Data-Controlled Linear Recurrence for Sequence Modeling

Linear Recurrence has proven to be a powerful tool for modeling long sequences efficiently. In this work, we show that existing models fail to take full advantage of its potential. Motivated by this finding, we develop GateLoop, a foundational sequence model that generalizes linear recurrent models such as S4, S5, LRU and RetNet, by employing data-controlled state transitions. Utilizing this theoretical advance, GateLoop empirically outperforms existing models for auto-regressive language modeling. Our method comes with a low-cost $O(l)$ recurrent mode and an efficient $O(l \log_{2} l)$ parallel mode making use of highly optimized associative scan implementations. Furthermore, we derive an $O(l^2)$ surrogate attention mode, revealing remarkable implications for Transformer and recently proposed architectures. Specifically, we prove that our approach can be interpreted as providing data-controlled relative-positional information to Attention. While many existing models solely rely on data-controlled cumulative sums for context aggregation, our findings suggest that incorporating data-controlled complex cumulative products may be a crucial step towards more powerful sequence models.

Tobias Katsch
arXiv:2311.01927 · cs.LG, cs.AI, cs.CL, cs.DS · submitted Nov 3, 2023 · updated Jan 27, 2024
abstract · pdf · html · Minor updates: Clarified tensor shapes

add comment on HN

Lucidrains has a re-implementation of this: https://github.com/lucidrains/gateloop-transformer and was unable to beat the transformer baseline at equal numbers of parameters.
I saw that. The WandB report is here [0]. The losses are quite close, but in my mind we'd need to see parameter counts at least 1 and 2 params more to make conclusions (with the dataset scaling proportionally). If the training performance can be investigated, there may be some wins in this area!

[0] https://api.wandb.ai/links/lucidrains/lgz368mf