82 lines
2.5 KiB
Markdown
82 lines
2.5 KiB
Markdown
|
---
|
||
|
|
language:
|
||
|
|
- en
|
||
|
|
license: apache-2.0
|
||
|
|
library_name: transformers
|
||
|
|
pipeline_tag: text-generation
|
||
|
|
tags:
|
||
|
|
- mamba
|
||
|
|
- state-space-model
|
||
|
|
- ssm
|
||
|
|
- causal-lm
|
||
|
|
- pytorch
|
||
|
|
- pretrained
|
||
|
|
datasets:
|
||
|
|
- wikimedia/wikipedia
|
||
|
|
---
|
||
|
|
|
||
|
|
# Mamba-50M
|
||
|
|
|
||
|
|
A ~50M-parameter [Mamba](https://arxiv.org/abs/2312.00752) (selective state-space) causal
|
||
|
|
language model, **pretrained from scratch on English Wikipedia**. Mamba replaces the attention
|
||
|
|
mechanism of a Transformer with a selective state-space layer, giving linear-time sequence
|
||
|
|
processing instead of the quadratic cost of self-attention.
|
||
|
|
|
||
|
|
> **This is a base model: pretrained only.** It has **not** been fine-tuned, instruction-tuned,
|
||
|
|
> RLHF'd, or aligned in any way. It is a raw next-token predictor intended for research.
|
||
|
|
|
||
|
|
## Model details
|
||
|
|
|
||
|
|
| | |
|
||
|
|
|---|---|
|
||
|
|
| Architecture | Mamba (selective SSM) |
|
||
|
|
| Parameters | ~50M |
|
||
|
|
| Context length | 512 tokens |
|
||
|
|
| Tokenizer | [GPT-NeoX-20B](https://huggingface.co/EleutherAI/gpt-neox-20b) (BPE, ~50k vocab) — the same tokenizer used by the original `state-spaces/mamba-*` models |
|
||
|
|
| Language | English |
|
||
|
|
| License | Apache-2.0 |
|
||
|
|
|
||
|
|
## Limitations
|
||
|
|
|
||
|
|
- **Pretrained only.** No fine-tuning, instruction tuning, or alignment. It does not follow instructions and has no safety filtering; it simply continues text.
|
||
|
|
- **Small.** At ~50M parameters it has limited fluency and reasoning; expect frequent hallucination and repetition.
|
||
|
|
- **English only.** Trained solely on English Wikipedia; other languages are out of distribution.
|
||
|
|
- **Domain-narrow.** Only Wikipedia was used as training data.
|
||
|
|
|
||
|
|
## Training data
|
||
|
|
|
||
|
|
Pretrained on the English subset of Wikipedia: over 3 million articles.
|
||
|
|
|
||
|
|
### Training procedure
|
||
|
|
|
||
|
|
| Hyperparameter | Value |
|
||
|
|
|---|---|
|
||
|
|
| Learning rate | 5e-4 |
|
||
|
|
| Sequence length | 512 |
|
||
|
|
| Batch size | 64 |
|
||
|
|
| Tokenizer | GPT-NeoX-20B (`EleutherAI/gpt-neox-20b`) |
|
||
|
|
| Optimizer | `AdamW` |
|
||
|
|
| LR schedule / warmup | `constant` / `10000` |
|
||
|
|
| Total tokens seen | `~ 2.5-2.9B` |
|
||
|
|
| Hardware | `2x Nvidia Quadro RTX 6000 24GB` |
|
||
|
|
|
||
|
|
## Evaluation
|
||
|
|
|
||
|
|
Evaluated on a held-out set of **10,000 Wikipedia articles** that were not seen during training.
|
||
|
|
The training and evaluation loss curves are shown below.
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
## Citation
|
||
|
|
|
||
|
|
If you use this model, please cite the Mamba paper:
|
||
|
|
|
||
|
|
```bibtex
|
||
|
|
@article{gu2023mamba,
|
||
|
|
title={Mamba: Linear-Time Sequence Modeling with Selective State Spaces},
|
||
|
|
author={Gu, Albert and Dao, Tri},
|
||
|
|
journal={arXiv preprint arXiv:2312.00752},
|
||
|
|
year={2023}
|
||
|
|
}
|
||
|
|
```
|