safetensors llama_enc custom_code ja arxiv:2302.13971 arxiv:2002.05202 arxiv:2104.09864 license:apache-2.0 region:us
---
language:
- ja
license: apache-2.0
---

(English part follows Japanese one.)

# TohokuNLP BERT-alpha 500M

長系列 (4,096, 8,192 トークン) の入力を可能にした日本語 [BERT](https://aclanthology.org/N19-1423/) モデルです。

## 利用方法

```python
from transformers import AutoModelForMaskedLM, AutoTokenizer

model = AutoModelForMaskedLM.from_pretrained(
  "tohoku-nlp/tohokunlp-bert-500m-sq4096-alpha", trust_remote_code=True
)
tokenizer = AutoTokenizer.from_pretrained("tohoku-nlp/tohokunlp-bert-500m-sq4096-alpha")
```

[transformers version 4.46.2](https://github.com/huggingface/transformers/releases/tag/v4.46.2) において、動作確認をしています。


## モデルアーキテクチャ

[Llama](https://arxiv.org/abs/2302.13971) アーキテクチャをベースとし、Causal Attention Mask を取り除くことで、Encoder 型言語モデルとして利用しています。  
具体的には、以下のモジュールを採用しています。

- [SwiGLU](https://arxiv.org/abs/2002.05202)
- [Rotary Positional Embeddings (RoPE)](https://arxiv.org/abs/2104.09864)
- [Grouped Query Attention (GQA)](https://aclanthology.org/2023.emnlp-main.298/)


### モデルパラメータの詳細

`tohoku-nlp/tohokunlp-bert-500m-sq4096-alpha`, `tohoku-nlp/tohokunlp-bert-500m-sq8192-alpha` のモデルパラメータの詳細は以下の表の通りです。

<table>
  <tr>
    <td>Num Layers</td>
    <td>24</td>
  </tr>
  <tr>
    <td>Hidden Size</td>
    <td>1,024</td>
  </tr>
  <tr>
    <td>FFN Hidden Size</td>
    <td>4,096</td>
  </tr>
  <tr>
    <td>Num Attention Heads</td>
    <td>16</td>
  </tr>
  <tr>
    <td>Num Key-Value Heads</td>
    <td>8</td>
  </tr>
</table>

## 学習データ

[llm-jp-corpus-v3](https://gitlab.llm-jp.nii.ac.jp/datasets/llm-jp-corpus-v3)  の日本語コーパスのサブセット (ja\_cc, ja\_warp\_html, ja\_warp\_pdf, ja\_wiki, kaken) を使用しました。  
Tokenizer には、[llm-jp-v3 tokenizer](https://github.com/llm-jp/llm-jp-tokenizer) を採用しています。  
また、学習時には Whole Word Masking を実施しています。  
Whole Word Masking 単語分割器には、[vibrato](https://github.com/daac-tools/vibrato) を利用しました。  
辞書は [bccwj-suw+unidic-cwj-3_1_1](https://github.com/daac-tools/vibrato/releases#:~:text=Compact%2Ddual-,bccwj%2Dsuw%2Bunidic%2Dcwj%2D3_1_1,-618%20MB) を用いています。


## 学習時の設定
 
モデルの重みを初期化した Llama アーキテクチャベースの Encoder モデルを from scratch で学習させています。  
まず、Sequence Length: 4096 で約 114 B トークンを学習させた (`tohoku-nlp/tohokunlp-bert-500m-sq4096-alpha`) 後、継続学習の形で、Sequence Length: 8192 に拡張し、約 34 B トークンを学習させました (`tohoku-nlp/tohokunlp-bert-500m-sq8192-alpha`)。  
各モデルの学習設定は以下の通りです。 

| | Params. | Tokens | Steps | checkpoint averaging |
| --- | --- | --- | --- | --- |
| tohoku-nlp/tohokunlp-bert-500m-sq4096-alpha | 581 M | 114 B | 100,000 | 95,000 and 100,000 steps |
| tohoku-nlp/tohokunlp-bert-500m-sq8192-alpha | 581 M | + 34 B | +15,000 | N/A |

学習には、Masked Language Modeling (MLM) のみ実施し、Next Sentence Prediction (NSP) は実施していません。

### 学習設定の詳細

|      | tohoku-nlp/tohokunlp-bert-500m-sq4096-alpha | tohoku-nlp/tohokunlp-bert-500m-sq8192-alpha |
| ---- | ---- | ---- |
| Batch Size (tokens) | 1,146,880 | 2,293,760 |
| Max Learning Rate | 1.0E-4 | 1.0E-4 |
| Min Learning Rate | 1.0E-6 | N/A |
| Learning Rate Warmup Steps | 10,000 | N/A |
| Scheduler | cosine | constant |
| Optimizer | AdamW | AdamW |
| Optimizer Config | beta_1 = 0.9, beta_2 = 0.999, eps = 1.0E-8 | beta_1 = 0.9, beta_2 = 0.999, eps = 1.0E-8 |
| Weight Decay | 0.01 | 0.01 |
| Gradient Clipping | 1 | 1 |
| Sequence Length | 4,096 | 8,192 |
| MLM Probability | 0.15 | 0.15 |
| Replace Masked-token Probability | 0.8 | 0.8 |
| Replace Random-token Probability | 0.1 | 0.1 |


## 評価

評価指標として、[JMTEB](https://www.sbintuitions.co.jp/blog/entry/2024/05/16/130848) を利用しました。  
評価結果は以下の通りです。

|  | Classification | Clustering | PairClassification | Reranking | Retrieval | STS |
|-----------|---------------|------------|-------------------|-----------|-----------|-----|
| tohoku-nlp/tohokunlp-bert-500m-sq4096-alpha | .577 | **.494** | .623 | .894 | .229 | .583 |
| tohoku-nlp/tohokunlp-bert-500m-sq8192-alpha | .460 | .420 | .623 | .885 | .295 | .433 |
| [tohoku-nlp/bert-base-japanese-v3](https://huggingface.co/tohoku-nlp/bert-base-japanese-v3) | **.644** | .486 | **.624** | **.903** | **.328** | **.693** |
| [retrieva-jp/bert-1.3b](https://huggingface.co/retrieva-jp/bert-1.3b) | .637 | .450 | **.624** | .897 | .260 | .517 |

|  | Param Size | Sequence Length |
| ---- | ---- | ---- |
| tohoku-nlp/tohokunlp-bert-500m-sq4096-alpha | 581 M | 4,096 |
| tohoku-nlp/tohokunlp-bert-500m-sq8192-alpha | 581 M | 8,192 |
| [tohoku-nlp/bert-base-japanese-v3](https://huggingface.co/tohoku-nlp/bert-base-japanese-v3) | 136 M | 512 |
| [retrieva-jp/bert-1.3b](https://huggingface.co/retrieva-jp/bert-1.3b) | 1.45 B | 2,048 |

2024 年 11 月現在、長系列の入力を扱うことができる日本語 Encoder モデルである [RetrievaBERT](https://huggingface.co/retrieva-jp/bert-1.3b) と比較して、同等の性能でパラメータサイズは約半分となっています。  
一方で、我々が学習したモデルを含め長系列の入力を扱うことのできるモデルは、パラメータサイズが大きいにも関わらず、入力可能系列長が比較的短く小パラメータサイズの [bert-base-japanese-v3](https://huggingface.co/tohoku-nlp/bert-base-japanese-v3) よりも性能が劣っています。  
この原因と改善方法については現在調査中です。  
今後もより高性能な日本語 Encoder 型言語モデルの開発を見据えた研究を継続していくことを計画しています。

※ ここで示した評価結果は、あくまで基盤モデルどうしの性能比較であり、 **(fine-tuning を実施した) 文書埋め込みモデルにおける評価結果ではない** ことに留意してください。


## ライセンス

このモデルは Apache License 2.0 の下で配布しています。

# 免責事項

本モデルの作者は本モデルを作成するにあたって、その内容、機能等について細心の注意を払っておりますが、モデルの出力が正確であるかどうか、安全なものであるか等について保証をするものではなく、何らの責任を負うものではありません。  
本モデルの利用により、万一、利用者に何らかの不都合や損害が発生したとしても、モデルやデータセットの作者や作者の所属組織は何らの責任を負うものではありません。

## 謝辞

このモデルの学習にあたり様々な面でご協力いただきました [Tohoku NLP Group](https://www.nlp.ecei.tohoku.ac.jp/) の皆様に感謝いたします。

## 作成者
- [Keito Kudo](https://x.com/k8kudo)
- [Daiki Shiono](https://x.com/onely7_deep)
- [Jun Suzuki](https://x.com/drJunSuzuki)


<br>
<br>
<br>
<br>

# TohokuNLP BERT-alpha 500M

A Japanese [BERT](https://aclanthology.org/N19-1423/) model capable of processing long sequences (4,096, 8,192 tokens).

## Usage

```python
from transformers import AutoModelForMaskedLM, AutoTokenizer

model = AutoModelForMaskedLM.from_pretrained(
  "tohoku-nlp/tohokunlp-bert-500m-sq4096-alpha", trust_remote_code=True
)
tokenizer = AutoTokenizer.from_pretrained("tohoku-nlp/tohokunlp-bert-500m-sq4096-alpha")
```

Operation has been confirmed with [transformers version 4.46.2](https://github.com/huggingface/transformers/releases/tag/v4.46.2).

## Model Architecture

Based on the [Llama](https://arxiv.org/abs/2302.13971) architecture, we removed the Causal Attention Mask to use it as an Encoder-type language model.  
Specifically, we adopted the following modules:

- [SwiGLU](https://arxiv.org/abs/2002.05202)
- [Rotary Positional Embeddings (RoPE)](https://arxiv.org/abs/2104.09864)
- [Grouped Query Attention (GQA)](https://aclanthology.org/2023.emnlp-main.298/)

### Model Parameter Details

The model parameters for `tohoku-nlp/tohokunlp-bert-500m-sq4096-alpha` and `tohoku-nlp/tohokunlp-bert-500m-sq8192-alpha` are as follows:

<table>
  <tr>
    <td>Num Layers</td>
    <td>24</td>
  </tr>
  <tr>
    <td>Hidden Size</td>
    <td>1,024</td>
  </tr>
  <tr>
    <td>FFN Hidden Size</td>
    <td>4,096</td>
  </tr>
  <tr>
    <td>Num Attention Heads</td>
    <td>16</td>
  </tr>
  <tr>
    <td>Num Key-Value Heads</td>
    <td>8</td>
  </tr>
</table>

## Training Data

We used a subset of the Japanese corpus from [llm-jp-corpus-v3](https://gitlab.llm-jp.nii.ac.jp/datasets/llm-jp-corpus-v3) (ja\_cc, ja\_warp\_html, ja\_warp\_pdf, ja\_wiki, kaken).  
We adopted the [llm-jp-v3 tokenizer](https://github.com/llm-jp/llm-jp-tokenizer) as our tokenizer.  
During training, we implemented Whole Word Masking.  
For Whole Word Masking word segmentation, we used [vibrato](https://github.com/daac-tools/vibrato).  
We used [bccwj-suw+unidic-cwj-3_1_1](https://github.com/daac-tools/vibrato/releases#:~:text=Compact%2Ddual-,bccwj%2Dsuw%2Bunidic%2Dcwj%2D3_1_1,-618%20MB) as our dictionary.

## Training Configuration

We trained a Llama architecture-based Encoder model from scratch after initializing the model weights.  
First, we trained with Sequence Length: 4096 for about 114B tokens (`tohoku-nlp/tohokunlp-bert-500m-sq4096-alpha`), then continued training with an extended Sequence Length: 8192 for about 34B tokens (`tohoku-nlp/tohokunlp-bert-500m-sq8192-alpha`).  
The training settings for each model are as follows:

| | Params. | Tokens | Steps | checkpoint averaging |
| --- | --- | --- | --- | --- |
| tohoku-nlp/tohokunlp-bert-500m-sq4096-alpha | 581 M | 114 B | 100,000 | 95,000 and 100,000 steps |
| tohoku-nlp/tohokunlp-bert-500m-sq8192-alpha | 581 M | + 34 B | +15,000 | N/A |

We only implemented Masked Language Modeling (MLM) during training, without Next Sentence Prediction (NSP).

### Detailed Training Settings

|      | tohoku-nlp/tohokunlp-bert-500m-sq4096-alpha | tohoku-nlp/tohokunlp-bert-500m-sq8192-alpha |
| ---- | ---- | ---- |
| Batch Size (tokens) | 1,146,880 | 2,293,760 |
| Max Learning Rate | 1.0E-4 | 1.0E-4 |
| Min Learning Rate | 1.0E-6 | N/A |
| Learning Rate Warmup Steps | 10,000 | N/A |
| Scheduler | cosine | constant |
| Optimizer | AdamW | AdamW |
| Optimizer Config | beta_1 = 0.9, beta_2 = 0.999, eps = 1.0E-8 | beta_1 = 0.9, beta_2 = 0.999, eps = 1.0E-8 |
| Weight Decay | 0.01 | 0.01 |
| Gradient Clipping | 1 | 1 |
| Sequence Length | 4,096 | 8,192 |
| MLM Probability | 0.15 | 0.15 |
| Replace Masked-token Probability | 0.8 | 0.8 |
| Replace Random-token Probability | 0.1 | 0.1 |

## Evaluation

We used [JMTEB](https://www.sbintuitions.co.jp/blog/entry/2024/05/16/130848) as our evaluation metric.  
The evaluation results are as follows:

|  | Classification | Clustering | PairClassification | Reranking | Retrieval | STS |
|-----------|---------------|------------|-------------------|-----------|-----------|-----|
| tohoku-nlp/tohokunlp-bert-500m-sq4096-alpha | .577 | **.494** | .623 | .894 | .229 | .583 |
| tohoku-nlp/tohokunlp-bert-500m-sq8192-alpha | .460 | .420 | .623 | .885 | .295 | .433 |
| [tohoku-nlp/bert-base-japanese-v3](https://huggingface.co/tohoku-nlp/bert-base-japanese-v3) | **.644** | .486 | **.624** | **.903** | **.328** | **.693** |
| [retrieva-jp/bert-1.3b](https://huggingface.co/retrieva-jp/bert-1.3b) | .637 | .450 | **.624** | .897 | .260 | .517 |

|  | Param Size | Sequence Length |
| ---- | ---- | ---- |
| tohoku-nlp/tohokunlp-bert-500m-sq4096-alpha | 581 M | 4,096 |
| tohoku-nlp/tohokunlp-bert-500m-sq8192-alpha | 581 M | 8,192 |
| [tohoku-nlp/bert-base-japanese-v3](https://huggingface.co/tohoku-nlp/bert-base-japanese-v3) | 136 M | 512 |
| [retrieva-jp/bert-1.3b](https://huggingface.co/retrieva-jp/bert-1.3b) | 1.45 B | 2,048 |

As of November 2024, compared to [RetrievaBERT](https://huggingface.co/retrieva-jp/bert-1.3b), a Japanese Encoder model capable of handling long sequences, our model achieves comparable performance with about half the parameter size.  
However, models capable of handling long sequences, including ours, despite their larger parameter sizes, perform worse than [bert-base-japanese-v3](https://huggingface.co/tohoku-nlp/bert-base-japanese-v3), which has a relatively shorter input sequence length and smaller parameter size.  
We are currently investigating the cause and potential improvements.  
We plan to continue research aimed at developing higher-performance Japanese Encoder language models in the future.

Note: The evaluation results shown here are comparisons between base models and are **not evaluation results for document embedding models (with fine-tuning)**.

## License

This model is distributed under the Apache License 2.0.

# Disclaimer

While the developers of this model have taken utmost care in creating its content and functionality, we do not guarantee the accuracy or safety of the model's outputs and assume no responsibility for them.  
The model's developers and their affiliated organizations bear no responsibility for any inconvenience or damages that may occur to users through the use of this model.

## Acknowledgments

We would like to thank the members of [Tohoku NLP Group](https://www.nlp.ecei.tohoku.ac.jp/) for their various forms of cooperation in training this model.

## Developers
- [Keito Kudo](https://x.com/k8kudo)
- [Daiki Shiono](https://x.com/onely7_deep)
- [Jun Suzuki](https://x.com/drJunSuzuki)

