|
--- |
|
license: apache-2.0 |
|
--- |
|
|
|
The *TokenFormer* is a **fully attention-based architecture** |
|
that unifies the computations of token-token and token-parameter interactions |
|
by entirely employing the attention mechanism, **maximizes the flexibility of neural network**.[(see paper)](https://arxiv.org/pdf/2410.23168). |
|
It contains four models of sizes |
|
150M, 450M, 900M, 1.5B. For each size, it's trained based on [gpt-neox](https://github.com/EleutherAI/gpt-neox) code base and uses [Pile](https://huggingface.co/datasets/EleutherAI/pile) with 300B tokens. |
|
All 4 model sizes are trained on the exact |
|
same data, in the exact same order. |
|
|
|
# TokenFormer-900M |
|
|
|
## Model Details |
|
|
|
- Developed by: [Haiyang Wang](https://haiyang-w.github.io/) |
|
- Model type: TokenFormer-based Language Model |
|
- Language: English |
|
- Learn more: [TokenFormer's GitHub repository](https://github.com/Haiyang-W/TokenFormer) |
|
for training procedure, config files, and details on how to use. |
|
[See paper](https://arxiv.org/pdf/2410.23168) for more evals and implementation |
|
details. |
|
- Library: [GPT-NeoX](https://github.com/EleutherAI/gpt-neox) |
|
- License: Apache 2.0 |
|
- Contact: to ask questions about this model, please email Haiyang Wang. |
|
|
|
<figure> |
|
|
|
| TokenFormer model | Layers | #QKV Param Tokens | #Output Param Tokens | #FFN Param Tokens | Model Dim | Heads | Batch Size | Learning Rate | Training Iterations | |
|
| ----------------: | -----: | :---------------: | :------------------: | :---------------: | :-------: | :---: | :--------: | :-------------------: | :-------------------------: | |
|
| 150M | 12 | 768 | 768 | 3072 | 768 | 12 | 2M | 6.0 x 10<sup>-4</sup> | 143000 | |
|
| 450M | 24 | 1024 | 1024 | 4096 | 1024 | 16 | 2M | 6.0 x 10<sup>-4</sup> | 143000 | |
|
| 900M | 32 | 1280 | 1280 | 5120 | 1280 | 16 | 2M | 6.0 x 10<sup>-4</sup> | 143000 | |
|
| 1.5B | 40 | 1536 | 1536 | 6144 | 1536 | 16 | 2M | 6.0 x 10<sup>-4</sup> | 143000 | |
|
<figcaption>Engineering details for the <i>TokenFormer</i>. </figcaption> |
|
</figure> |
|
|
|
## Training |
|
|
|
### Training data |
|
|
|
[The Pile](https://pile.eleuther.ai/) is a 825GiB general-purpose dataset in |
|
English. It was created by EleutherAI specifically for training large language |
|
models. It contains texts from 22 diverse sources, roughly broken down into |
|
five categories: academic writing (e.g. arXiv), internet (e.g. CommonCrawl), |
|
prose (e.g. Project Gutenberg), dialogue (e.g. YouTube subtitles), and |
|
miscellaneous (e.g. GitHub, Enron Emails). See [the Pile |
|
paper](https://arxiv.org/abs/2101.00027) for a breakdown of all data sources, |
|
methodology, and a discussion of ethical implications. Consult [the |
|
datasheet](https://arxiv.org/abs/2201.07311) for more detailed documentation |
|
about the Pile and its component datasets. The Pile can be downloaded from |
|
the [official website](https://pile.eleuther.ai/), or from a [community |
|
mirror](https://the-eye.eu/public/AI/pile/).<br> |
|
|
|
### Training procedure |
|
We follow the default training strategy of [Pythia](https://arxiv.org/abs/2304.01373) in [gpt-neox](https://github.com/EleutherAI/gpt-neox), |
|
including the dataset processing, hyper-parameter and code base. |
|
All models were trained on the exact same data, in the exact same order. Each |
|
model saw 299,892,736,000 tokens during training. |
|
|
|
All *TokenFormer* models trained for 143000 steps at a batch size |
|
of 2M (2,097,152 tokens).<br> |
|
See [GitHub](https://github.com/Haiyang-W/TokenFormer) for more details on training |
|
procedure.<br> |
|
TokenFormer uses the same tokenizer as [GPT-NeoX- |
|
20B](https://huggingface.co/EleutherAI/gpt-neox-20b). |
|
|
|
## Evaluations |
|
|
|
All *TokenFormer* models were evaluated using the [LM Evaluation |
|
Harness](https://github.com/EleutherAI/lm-evaluation-harness). |
|
You can run the evaluation with our [instruction](https://github.com/Haiyang-W/TokenFormer?tab=readme-ov-file#evaluations).<br> |
|
Expand the sections below to see plots of evaluation results for all |
|
TokenFormer compared with Opensource Transformer-based LLMs. |
|
|
|
<figure> |
|
|
|
| Model | #Param | LAMBADA | HellaSwag | PIQA | Arc-E | Arc-C | WinoGrande | Average | |
|
| :----: | :------: | :------: | :-------: | :--: | :---: | :---: | :--------: | :------: | |
|
| Pythia | 150M | 35.4 | 30.3 | 62.3 | 43.6 | 23.6 | 51.3 | 40.1 | |
|
| **TokenFormer** | 150M | **45.0** | **35.5** | **64.9** | **47.3** | **24.9** | **50.4** | **44.7** | |
|
| Pythia | 410M | 51.4 | 40.6 | 66.9 | 52.1 | 24.6 | 53.8 | 48.2 | |
|
| **TokenFormer** | 450M | **57.3** | **47.5** | **69.5** | **56.2** | **26.7** | **54.6** | **52.0** | |
|
| Pythia | 1B | 56.1 | 47.2 | 70.7 | 57.0 | 27.1 | 53.5 | 51.9 | |
|
| **TokenFormer** | 900M | **64.0** | **55.3** | **72.4** | **59.9** | **30.6** | **56.4** | **56.4** | |
|
| GPT-Neo | 1.3B | 57.2 | 48.9 | 71.1 | 56.2 | 25.9 | 54.9 | 52.4 | |
|
| OPT | 1.3B | 58.0 | 53.7 | 72.4 | 56.7 | 29.6 | 59.5 | 55.0 | |
|
| Pythia | 1.3B | 61.7 | 52.1 | 71.0 | 60.5 | 28.5 | 57.2 | 55.2 | |
|
| GPT-Neo | 2.7B | 62.2 | 55.8 | 71.1 | 61.1 | 30.2 | 57.6 | 56.5 | |
|
| OPT | 2.7B | 63.6 | 60.6 | 74.8 | 60.8 | 31.3 | 61.0 | 58.7 | |
|
| Pythia | 2.8B | 64.7 | 59.3 | 74.0 | 64.1 | 32.9 | 59.7 | 59.1 | |
|
| **TokenFormer** | 1.5B | **64.7** | 60.0 | **74.8** | **64.8** | 32.0 | 59.7 | **59.3** | |
|
<figcaption>Zero-shot evaluation of Language Modeling. </figcaption> |
|
</figure> |
