New Attention variant - SSM (Adopted from Mamba and Hymba)

[Mars-Cat2023]

New Attention Variant - SSM

We support one more command option --attention_variant="ssm".

TODO

• Adding a hyperparameter list for selection of different attention variants (i.e. "causal", "linear" or "ssm") for different layers.

Training Results

You can replicate the following training results by running commands:

python3 train.py --compile
python3 train.py --compile --attention_variant="ssm"

Comparison between --attention_variant="causal" and --attention_variant="ssm"

--attention_variant="causal" step 0: train loss 4.2484, train_stdev 0.0031, val loss 4.2482, val_stdev 0.0031, lr 0.0010
--attention_variant="ssm" step 0: train loss 4.3725, train_stdev 0.0045, val loss 4.3620, val_stdev 0.0048, lr 0.0010

[msaligane]

Our Transformer Attention is called by python3 train.py.

Now you can try this new command to use all "SSM" blocks.
python3 train.py --attention_variant="ssm"

• TODO: How to specify the "SSM" for Layer 1 and "Attention" for Layer

Great work Qilong!

[gkielian]

Awesome! I looked through the dependency checks, and seems we'll need to add a few more pip installs for this to pass our CI.

Could you add these to the requirements_cpu.txt and README?

[Mars-Cat2023]

Updated and just waiting for solving CI and LICENSE problems.