Sasha Rush, a machine learning researcher, writes annotated, code-driven explanations of models like S4, Mamba, and transformers.
HN commenters call the work from his group on S4 layers essential reading, and link his Thinking Like Transformers and Mamba posts.
> Why does, 'mat' follow from 'the cat sat on the ...'
You're confidently incorrect by oversimplifying all LLMs to a base model performing a completion from a trivial context of 5 words.
This is tantamount to a straw man. Not only do few people use untuned base models, it completely ignores in-context learning that allows the model to build complex semantic structures from the relationships learnt from its training data.
Unlike base models, instruct and chat fine-tuning teaches models to 'reason' (or rather, perform semantic calculations in abstract latent spaces) with their "conditional probability structure", as you call it, to varying extents. The model must learn to use its 'facts', understand semantics, and perform abstractions in order to follow arbitrary instructions.
You're also confabulating the training metric of "predicting tokens" with the mechanisms required to satisfy this metric for complex instructions. It's like saying "animals are just performing survival of the fittest". While technically correct, complex behaviours evolve to satisfy this 'survival' metric.
You could argue they're "just stitching together phrases", but then you would be varying degrees of wrong:
For one, this assumes phrases are compressed into semantically addressable units, which is already a form of abstraction ripe for allowing reasoning beyond 'stochastic parroting'.
For two, it's well known that the first layers perform basic structural analysis such as grammar, and later…
Still chewing but my takeaways thus far:
SSM models strengths are with continuous data like audio and video, they struggle with discrete data like text/ DNA. This newest architecture uses selective attention to try to address the weaknesses around discrete data with some loss in performance in continuous tasks, empirically shown here with audio. The empirical exploration was limited to smaller size models, the performance as larger scales is yet to be explored in practice.
I found my deepest understanding of the selection mechanism came from struggling with the discretization in section 2, followed by the deeper explanations of the variables involved in 3.5.2. This video gives excellent background to SSMs, along with a detailed walk through of the paper itself.[a]
I am still coming to understand S4, SSMs in general but the video suggested this annotated explainer that has been helping a lot[b].
I would also point out section 3.1 and it’s discussion of the tradeoffs between compression and effectiveness as particularly interesting.
I do wonder how many different GPUs / hardware architectures will be able to execute the optimizations that are described as critical. I think the nature of the optimizations is the part of the paper I understand least well.
The paper taken at face value looks very exciting. The promise of a very large context window with 5x throughput for inference would be huge if it proves to scale well. I do wonder if it will make sense to train SSMs…
The work out of that group, starting with S4 layers, is 10000% the stuff to be paying attention to.
https://srush.github.io/annotated-s4/
HiPPO was brilliant - instead of working with the raw sequence, you work with its weighted laplace transform, and instead of actually computing the laplace transform you find the rule to update it when new data is added. Furthermore, we can 'band limit' the Laplace transform (similar to PCA) to keep only the 'most important' information while still preserving most of the information in the sequence - this is a common and quite effective compression technique.
Any 'fast' transformer is going to be working with some kind of sampling or aggregation or compression of the long sequence. Sampling is ultimately going to be too noisy, and standard aggregations are going to be too coarse. So the thing to bet on is better compression techniques, which is what the S4/RWKV group are ultimately working on.
Oh hey, I wrote this. Been a long time. I had the lucky of break of working in machine translation / parsing when the most important invention of the century happened in my niche field.
I'm pretty interested in the intersection of code / ML. If that's your thing here are some other writing you might be interested in.
* Thinking about cuda: http://github.com/srush/gpu-puzzles
* Tensors considered harmful: https://nlp.seas.harvard.edu/NamedTensor
* Differentiating SVG: https://srush.github.io/DiffRast/
* Annotated S4: https://srush.github.io/annotated-s4/
Recently moved back to industry, so haven't had a chance to write in a while.
A related line of work is "Thinking Like Transformers" [1]. They introduce a primitive programming language, RASP, which is composed of operations capable of being modeled with transformer components, and demonstrate how different programs can be written with it, e.g. histograms, sorting. Sasha Rush and Gail Weiss have an excellent blog post on it as well [2]. Follow on work actually demonstrated how RASP-like programs could actually be compiled into model weights without training [3].
[1] https://arxiv.org/abs/2106.06981
[2] https://srush.github.io/raspy/
[3] https://arxiv.org/abs/2301.05062
This is an article about "S4" that uses "Structured State Spaces".
https://srush.github.io/annotated-s4/
The matrices that make up the state space (A, B and C) are constant in S4. This allowed them to represent some of the math operations as a convolution (which can be parallelized).
The difference between S4 and Mamba is that these matrices are input-dependent in Mamba. Plus they add in some CUDA stuff ("parallel scan") to make it faster to compute on a GPU even if these matrices are not constant.
Yannic Kilcher's video on Mamba might also be a good resource: https://youtu.be/9dSkvxS2EB0