見出し画像

RetrievaBERTの公開

株式会社レトリバ

AIコンサルティング&ソリューション部、Chief Research Officerの西鳥羽( @jnishi )です。この度 RetreivaBERTモデルを構築し、Huggingface Hubで公開いたしました。

一般的なHuggingface TransformersのBERTと同じように利用できます。ただしcustom modelになりますので trust_remote_code=True を指定する必要がありますのでご注意ください。以下がサンプルコードになります。

from transformers import AutoModelForMaskedLM, AutoTokenizer, pipeline

model_id = "retrieva-jp/bert-1.3b"
model = AutoModelForMaskedLM.from_pretrained(model_id, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(model_id)
pipe = pipeline("fill-mask", model=model, tokenizer=tokenizer)

text = "こんにちは！私の名前は<MASK|LLM-jp>です！"
print(pipe(text))

RetreivaBERTの開発

RAGにおけるSentence Representationの重要性

最近ではGPTをはじめとするDecoder-onlyのモデルが日本語も含めて数多く公開されていて日進月歩の勢いで性能が良くなっています。更にRetrieval Augumented Generation (RAG)によってHullcinationの抑制や社内文書の反映させることがある程度可能になったため、RAGを利用されている方が多いかと思います。RAGのRetrievalの部分においては如何に精度良く関連する情報を検索してくるかが大事になります。文章の検索では文章の埋め込みベクトルを作成して類似度の高い文章を取得します。Sentence Representation(文章のベクトル化)については以前から研究開発しており、弊社の研究開発ブログでも SimCSEの紹介、DiffCSEの紹介、ME5の紹介などアルゴリズムの紹介や、Ruby on Rails & Postgresqlでベクトル近傍検索を手軽に実装するためのTipsの紹介など行ってきました。

系列長の大きいBERTの必要性

過去のブログで取り上げた手法のようにSentence RepresentationにおいてはBERTを利用している手法が多くあります。
BERTは数多く公開されており、日本語に対応したものには東北大BERT、早稲田大学RoBERTa、京大DeBERTaなどがあります。おかげで日本語のSentence Representationを作ることも可能になっています。一方で、RAGにより用途が広がり、長い文書を対象とするシーンも出てきました。その場合、BERTの事前学習時の系列長によりEmbeddingを作成できる文書長が制限されてしまいます。我々の知る限りでは現在公開されているBERTは系列長512までとなっているものがほとんどだと思います。そこで今回、Sentence Representationで用いるためのBERTとして系列長が2048まで対応しているBERTを構築しました。パラメータサイズについては、日本語のBERTとしては最大規模の1.3Bとしています。

RetrievaBERTのアーキテクチャ

今回BERTを構築するにあたり、アーキテクチャも見直しました。ここ数年でGPTをはじめとするDecoder-onlyのモデルの学習が活発になり、SwiGLUを始めとした新しいアーキテクチャやChinchilla則やSequence length warmupなど学習における知見などが発表されています。これらはEncoder-onlyのモデルの学習にも有用かと考え、以下のものを採用しました。

  • SwiGLU活性化関数

  • PreNorm

  • Grouped Query Attention (Multi Query Attention)

TokenizerにはLLM-jp Tokenizerのver 3.0b1を用いています。
また、学習コーパスには日本語と英語の両方を用いたほか、少量の中国語と韓国語を用いています。

学習には系列長におけるCurriculum Learningを行い、始めは系列長128で学習を行い、その後256, 512, 2048と段々長くしていきました。Sequence length warmupと似ていますが、今回の学習はそれほど厳密なスケジューリングは行っておりません。Batch sizeは1024で固定しており、約180B Tokenの学習を実施しました。具体的な系列長とステップ数は以下のようになっています。

  • 系列長128: 31,000 steps.

  • 系列長256: 219,000 steps.

  • 系列長512: 192,000 steps.

  • 系列長2048: 12,000 steps.

おわりに

RAGで用いられる長文向けSentence RepresentationのベースとなるRetrievaBERTを構築し、公開いたしました。
Apache 2.0ライセンスで公開いたします。商用利用可です。気軽に使っていただければと思いますが、何かありましたら pr@retrieva.jpにフィードバックを頂けたら幸いです。

最後に、このモデルの学習は東京工業大学のTSUBAMEミニキャンプで実施させていただきました。数多くのメンターにアドバイスを頂きつつ、TSUBAME 4.0を用いて学習することができました。この場を借りて感謝いたします。