Views
No views yet
Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers
Sukjun Hwang*, Aakash Lahoti*, Tri Dao, Albert Gu
Paper: https://arxiv.org/abs/2407.09941
Blogpost: https://goombalab.github.io/blog/2024/hydra-part1-matrix-mixer/
pip install mamba-ssm./hydra/bert), install additional required packages viapip install -r requirements.txt1import torch
2from .hydra import Hydra
3
4batch, length, dim = 2, 64, 16
5x = torch.randn(batch, length, dim).to("cuda")
6model = Hydra(
7 d_model=dim, # Model dimension d_model
8 d_state=64, # SSM state expansion factor
9 d_conv=7, # Local non-causal convolution width
10 expand=2, # Block expansion factor
11).to("cuda")
12y = model(x)
13assert y.shape == x.shape1from .hydra import MatrixMixer
2
3model = MatrixMixer(
4 """
5 matrix_mixer_type: options for matrix_mixer_type
6 {'dense', 'toeplitz', 'vandermonde', 'cauchy', 'low_rank', 'attention', 'quasiseparable'}
7 is_data_dependent: boolean flag to parameterize the mixer matrix to SAM
8 """
9 matrix_mixer_type=matrix_mixer_type,
10 is_data_dependent=is_data_dependent,
11 d_model=dim, # Model dimension d_model
12 qk_dim=qk_dim, # dimension for QK
13).to("cuda")
14y = model(x)
15assert y.shape == x.shapepython main.py yamls/pretrain/hydra.yamlcomposer -n 8 main.py yamls/pretrain/hydra.yamlpython glue.py yamls/finetune/hydra.yaml@article{hydra,
title={Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers},
author={Hwang, Sukjun and Lahoti, Aakash and Dao, Tri and Gu, Albert},
journal={arXiv preprint arXiv:2407.09941},
year={2024}
}