Merge pull request #56 from srush/dss

[WIP] Diagonal State Spaces

Makefile

@@ -25,15 +25,19 @@ autoformat:
 	black s4/s4.py s4/data.py s4/train.py s4/sample.py
 	flake8 --show-source s4/s4.py s4/data.py s4/train.py s4/sample.py
 
-notebook: s4/s4.py
+notebook: s4/s4.py s4/dss.py
 	jupytext --to notebook s4/s4.py -o s4.ipynb
+	jupytext --to notebook s4/dss.py -o dss.ipynb
 
-html: s4/s4.py
+html: s4/s4.py s4/dss.py
 	jupytext --to notebook s4/s4.py -o s4.ipynb
 	jupyter nbconvert --to html s4.ipynb
+	jupytext --to notebook s4/dss.py -o dss.ipynb
+	jupyter nbconvert --to html dss.ipynb
 
-s4/s4.md: s4/s4.py
+s4/s4.md: s4/s4.py s4/dss.py
 	jupytext --to markdown s4/s4.py
+	jupytext --to markdown s4/dss.py
 
 blog: s4/s4.md
 	pandoc docs/header-includes.yaml s4/s4.md  --katex=/usr/local/lib/node_modules/katex/dist/ --output=docs/index.html --to=html5 --css=docs/github.min.css --css=docs/tufte.css --no-highlight --self-contained --metadata pagetitle="The Annotated S4"

README.md

@@ -37,9 +37,13 @@ Gets "best" 97.76% accuracy in 10 epochs @ 40s/epoch on a TitanRTX.
 ```
 # Following @frederick0329's/@albertgu's results: https://github.com/srush/annotated-s4/pull/43#issuecomment-1065444261
 python -m s4.train --dataset cifar-classification --model s4 --epoch 100 --bsz 64 --n_layers 6 --p_dropout 0.25 --lr 5e-3 --d_model 512
+
+# DSS Model
+python -m s4.train --dataset cifar-classification --model dss --epoch 100 --bsz 64 --n_layers 6 --p_dropout 0.25 --lr 5e-3 --d_model 512
 ```
 
-Gets "best" 85.81% accuracy after 100 epochs @ 3m8s/epoch on a TitanRTX
+S4 gets "best" 87.05% accuracy after 100 epochs @ 3m8s/epoch on a TitanRTX
+DSS gets "best" 88.90% accuracy after 100 epochs @ 3m11s/epoch on a TitanRTX
 
 ---