BitNetのリポジトリ日本語化

BitNetは、大規模言語モデルのための1ビットトランスフォーマーを拡張する新しい取り組みです。この記事では、BitNetの基本から最新の進展まで、初心者でも理解しやすいように解説します。論文「BitNet: Scaling 1-bit Transformers for Large Language Models」に基づいたPyTorch実装の紹介から始め、BitNetの構造、利点、そしてどのようにして性能向上に貢献するかを掘り下げていきます。

BitNetとは？
1. 最新の進展
BitNetの利用方法
BitNetの実践的な応用
1. 推論
2. Huggingface Usage
BitNetの未来
リポジトリ
1. 元リポジトリ
2. 日本語化リポジトリ
GPTs
1. 関連

BitNetとは？

BitNetは、トランスフォーマーモデル内の線形射影（PyTorchでのnn.Linear）をBitLinearモジュールに置き換えることで実現される、シンプルながら効果的なアーキテクチャです。この変更により、モデルの量子化が可能となり、大幅な効率化を実現します。特に、BitLinearはテンソルを正規化し、二値化（-1または1の値を取る）後、絶対最大量子化を行い、その後逆量子化します。

BitNetの利用方法

BitNetの利用は、基本的に以下のステップで行います：

BitLinearレイヤーの導入 :

BitLinearレイヤーは、BitNetの核となる革新です。このレイヤーを通じて、モデルの量子化が行われ、処理効率が大幅に向上します。


# 必要なライブラリのインポート
import torch
from bitnet import BitLinear

# 入力データの作成
# 10個の512次元のランダムなテンソル（ベクトル）を生成します。
# これは、モデルに入力されるサンプルのバッチです。
x = torch.randn(10, 512)

# BitLinearレイヤーの作成
# このレイヤーは、512次元の入力を400次元の出力に変換します。
layer = BitLinear(512, 400)

# レイヤーを通じて入力xを処理して出力yを得る
# ここで、入力テンソルxがBitLinearレイヤーによって処理され、
# 出力テンソルyが生成されます。
y = layer(x)

# 出力の表示
# 処理された出力テンソルyの内容をコンソールに表示します。
print(y)

BitNetTransformerの全面的な実装 :

BitNetTransformerは、テキストだけでなく、画像やビデオ、オーディオ処理にも適用可能なトランスフォーマーモデルの完全な実装です。


# 必要なライブラリをインポートする
import torch
from bitnet import BitNetTransformer

# 整数のランダムなテンソルを作成する
# ここでは、0から20000までの値を持つ、1x1024のテンソルを生成しています。
# これは、モデルへの入力として使用されるサンプルです。
x = torch.randint(0, 20000, (1, 1024))

# BitNetTransformerモデルを初期化する
bitnet = BitNetTransformer(
    num_tokens=20000,  # 入力に含まれるユニークなトークンの数
    dim=1024,  # 入力および出力の埋め込みの次元数
    depth=6,  # トランスフォーマーレイヤーの数
    heads=8,  # 注意ヘッドの数
    ff_mult=4,  # フィードフォワードネットワークの隠れ層の次元数の倍数
)

# テンソルをトランスフォーマーモデルに通す
# 入力テンソルxがモデルを通過し、出力logitsを生成します。
logits = bitnet(x)

# 出力の形状を表示する
# 出力テンソルlogitsの形状（次元）を表示します。
print(logits)

BitAttentionによる注意機構の最適化 :

通常の多頭部注意機構に代わり、BitLinearを適用したBitAttentionを用いることで、モデルのデコード速度と処理能力が向上します。


# 必要なライブラリをインポートする
import torch
from bitnet import BitMGQA

# 形状が(1, 10, 512)のランダムなテンソルを作成する
# これは、1つのバッチに含まれる10個のデータポイントがあり、
# 各データポイントが512次元の特徴ベクトルを持っていることを意味します。
x = torch.randn(1, 10, 512)

# BitMGQAモデルのインスタンスを作成する
# 入力サイズは512、注意ヘッドの数は8、レイヤーの数は4です。
gqa = BitMGQA(512, 8, 4)

# 入力テンソルをBitMGQAモデルを通して処理し、出力と注意重みを取得する
# ここでは、自己注意メカニズムを使用しています（入力としてxを3回渡しています）。
# need_weights=Trueは、注意重みも返すように指定しています。
out, _ = gqa(x, x, x, need_weights=True)

# 出力テンソルと注意テンソルの形状を表示する
print(out)

BitFeedForwardによる前方向フィード :

BitLinearとGELUを組み合わせた構造を用いることで、より効果的な前方向フィードが実現されます。


# 必要なライブラリをインポートする
import torch
from bitnet import BitFeedForward

# 形状が(10, 512)のランダムな入力テンソルを作成する
# これは、10個のデータポイントがあり、各データポイントが512次元の特徴ベクトルを持っていることを意味します。
x = torch.randn(10, 512)

# BitFeedForwardクラスのインスタンスを作成する
# ここでは、以下のパラメータを指定しています：
# - input_dim: 入力次元数は512
# - hidden_dim: 隠れ層の次元数も512
# - num_layers: レイヤー数は4
# - swish: Swish活性化関数を使用する
# - post_act_ln: 各活性化後にレイヤーノーマライゼーションを適用する
# - dropout: ドロップアウト率を0.1（10%）とする
ff = BitFeedForward(512, 512, 4, swish=True, post_act_ln=True, dropout=0.1)

# 入力テンソルxをBitFeedForwardネットワークに適用する
y = ff(x)

# 出力テンソルyの形状を表示する
print(y)  # torch.Size([10, 512])

BitNetの実践的な応用

訓練 : train.pyファイルを使用してenwiki8という小さな1GBのWikipediaデータセットでトレーニングを行います。

推論

BitNetInferenceを使用してモデルからのテキスト生成を行うことができます。


# BitNetの推論用モジュールをインポートする
from bitnet import BitNetInference

# BitNetInferenceクラスのインスタンスを作成する
bitnet = BitNetInference()

# 事前にトレーニングされたモデルをロードする
# "../model_checkpoint.pth"は、モデルの重みが保存されているファイルのパスです。
bitnet.load_model("../model_checkpoint.pth")

# モデルを使用してテキストを生成する
# "The dog jumped over the "は、生成を開始するためのプロンプトテキストです。
# 512は、生成するテキストの最大長を指定しています。
output_str = bitnet.generate("The dog jumped over the ", 512)

# 生成されたテキストを出力する
print(output_str)

このコードの実行により、事前にトレーニングされたモデルを使って、与えられたプロンプトに基づいて新しいテキストを生成することができます。bitnet.load_model("../model_checkpoint.pth")は、指定されたパスからモデルの重みをロードします。この重みには、モデルが学習したパターンやルールが含まれており、テキスト生成の際に使用されます。

bitnet.generate("The dog jumped over the ", 512)は、指定されたプロンプトを基にして、最大512文字のテキストを生成します。この関数は、モデルが学習した知識を利用して、プロンプトに自然に続くようなテキストを作り出します。

最後に、print(output_str)によって生成されたテキストがコンソールに出力されます。これにより、モデルがどのようなテキストを生成したかを確認することができます。

Huggingface Usage


# 必要なライブラリをインポート
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from bitnet import replace_linears_in_hf

# Hugging FaceのTransformersからモデルをロード
model_name = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)  # テキストトークナイザーをロード
model = AutoModelForSequenceClassification.from_pretrained(model_name)  # テキスト分類モデルをロード

# 線形レイヤーをBitLinearに置き換え
replace_linears_in_hf(model)

# 分類するためのサンプルテキスト
text = "Replace this with your text"
# トークナイザーを使用してテキストをモデルの入力形式に変換
inputs = tokenizer(
    text, return_tensors="pt", padding=True, truncation=True, max_length=512
)

# 推論を実行
model.eval()  # モデルを評価モードに設定
with torch.no_grad():  # 勾配計算を無効化
    outputs = model(**inputs)  # モデルに入力を渡して出力を得る
    predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)  # 出力から予測確率を計算
    print(predictions)

# 予測結果の処理
predicted_class_id = predictions.argmax().item()  # 最も確率が高いクラスのIDを取得
print(f"Predicted class ID: {predicted_class_id}")  # 予測されたクラスIDを表示

# オプション：予測されたクラスIDに対応するラベルをマッピング
# labels = ["Label 1", "Label 2", ...]  # モデルのクラスに対応するラベルを定義
# print(f"Predicted label: {labels[predicted_class_id]}")  # 予測されたラベルを表示

このコードでは、まずHugging FaceのTransformersライブラリから必要なモジュールをインポートしています。AutoModelForSequenceClassificationは、シーケンス（テキスト）分類のための事前訓練済みモデルをロードするために使用され、AutoTokenizerは、テキストをモデルが理解できる形式に変換するために使用されます。

replace_linears_in_hf(model)関数は、モデル内の全ての線形レイヤーをBitNetのBitLinearレイヤーに置き換えることで、モデルの効率を向上させます。

次に、tokenizerを使用して入力テキストをトークン化し、適切な形式に変換します。この処理により、モデルが入力テキストを理解しやすくなります。

model.eval()はモデルを評価モードに設定し、torch.no_grad()は推論時に勾配計算を無効にします。これにより、推論が高速化され、メモリ使用量が削減されます。

最後に、モデルの出力から最も確率が高いクラスのIDを特定し、それを表示します。必要に応じて、このクラスIDを実際のラベルにマッピングすることも可能です。