萬字長文剖析ChatGPT

簡單來說，ChatGPT 是自然語言處理（NLP）和強化學習（RL）的一次成功結合，考慮到讀者可能只熟悉其中一個方向或者兩個方向都不太熟悉，本文會將 ChatGPT 涉及到的所有知識點儘可能通俗易懂的方式展現出來，有基礎的同學可以選擇性跳過一些內容。

一、GPT 的進化史

本節的主要目的是介紹自然語言處理中語言模型的一些基礎知識，理解語言模型到底在做什麼。

1.1 GPT

所謂的 GPT(Generative Pre-trained Transformer)，其實是 Generative Pre Training of a language model（語言模型）。那什麼是語言模型呢？可以簡單地把語言模型理解為“給定一些字或者詞，預測下一個字或者詞的模型”，這裡的字或者詞在 NLP 領域通常也被稱為 token，即給定已有 token，預測下一個 token 的模型，這裡舉個例子，我們在搜尋引擎裡進行搜尋時，自動會往後聯想就是種語言模型的體現。

那麼訓練語言模型有什麼優勢呢？答案就是它不需要人工標註資料！ 比如以“today is a good day”為例，它可以被拆解為：

接下來，讓我們來數學化地描述一下，給定一個句子，比如 ，語言模型其實就是想最大化：

【其中k 是考慮的視窗長度，條件機率p透過一個引數為O的神經網路來描述。

GPT 的神經網路採用了一個多層的 Transformer decoder，輸入經過 embedding 層（token embedding 疊加 position embedding），然後過多層解碼器，最後透過一個 position-wise 的前向網路得到輸出的分佈：

有了模型結構，有了目標函式，已經可以預訓練一個大容量的語言模型了，這也就是 GPT 的第一階段，在 GPT 的訓練流程裡還有第二個階段，利用一些有標籤資料進行微調。假設輸入為 ，標籤為【 ，】可以將輸入喂入模型，模型的輸出再疊加一個線性層作為最終的輸出：

目標函式也就是：

然而作者在微調時還發現，同時考慮語言模型的自迴歸目標函式效果更好，也就是：

在微調階段，可以最佳化的引數只有頂部的線性層已經用作分隔符的 token embedding。下圖展示的就是 GPT 做微調時對文字的一些常見做法，其實就是拼接和加分割符之類的操作，如下圖。

1.2 GPT2

GPT1 需要對特定任務再進行精調（依賴有標籤資料進行監督學習），而 GPT2 則是考慮在預訓練時考慮各種不同的任務，也就更加通用化。因此，GPT2 的模型從原本 GPT1 的：

改為 task conditioning 的形式。

也就是把任務也作為模型的輸入，具體的做法是引入一些表示任務的 token，舉幾個例子。

• 自迴歸任務input：Today is a output：good
• 翻譯任務input：Today is a [翻譯為中文] output：今天是一個
• 問答任務input：我是小明 [問題] 我是誰 [答案] output：小明上面例子中 [翻譯為中文]、[問題] 、[答案] 這些就是用於告訴模型執行什麼任務的 token。

透過這樣的方式，各種任務都能塞進預訓練裡進行了，想學的越多，模型的容量自然也需要更大，GPT2 的引數量達到了 1.5 Billions（GPT1 僅 117 Millions）。

1.3 GPT3

GPT3 可以理解為 GPT2 的升級版，使用了 45TB 的訓練資料，擁有 175B 的引數量，真正詮釋了什麼叫暴力出奇跡。

GPT3 主要提出了兩個概念：

• 情景（in-context）學習：就是對模型進行引導，教會它應當輸出什麼內容，比如翻譯任務可以採用輸入：請把以下英文翻譯為中文：Today is a good day。這樣模型就能夠基於這一場景做出回答了，其實跟 GPT2 中不同任務的 token 有異曲同工之妙，只是表達更加完善、更加豐富了。
• Zero-shot, one-shot and few-shot：GPT3 打出的口號就是“告別微調的 GPT3”，它可以透過不使用一條樣例的 Zero-shot、僅使用一條樣例的 One-shot 和使用少量樣例的 Few-shot 來完成推理任務。下面是對比微調模型和 GPT3 三種不同的樣本推理形式圖。

二、 ChatGPT

ChatGPT 使用了類似 InstructGPT 的方式來訓練模型，該方法也叫做 Learning from Human Feedback。主要分為三個步驟：

• 用有監督資料精調 GPT-3.5；
• 對於模型輸出的候選結果（因為取樣會導致同一輸入有不同輸出）進行打分，從而訓練得到一個獎勵模型；
• 使用這個獎勵模型，用 PPO 演算法來進一步對模型進行訓練。

2.1 GPT是如何訓練的

接下來我們來動手實踐一下如何訓練一個 GPT 模型出來，以從頭訓練一個程式碼補全的 GPT 模型為例。比如我們給模型一個提示，然後就能輸出提示。

from transformers import AutoTokenizer, AutoModelForSequenceClassification
# build a BERT classifier


tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')
model = AutoModelForSequenceClassification.from_pretrained('bert-base-uncased')

為了訓練這樣一個模型，首先我們需要準備用於訓練的資料，常見的程式碼補全的資料為 codeparrot，下載連結：https://huggingface.co/codepa...。這裡隨便列印一條資料，如下圖。

然而，模型是不能直接接收這樣的“文字”資訊的，所以訓練 NLP 模型前通常需要對其進行“分詞”，轉化為由一串數字表示，可以建立一個分詞器：

tokenizer = AutoTokenizer.from_pretrained("./code-search-net-tokenizer")

對上面的程式碼進行分詞轉化，就可以得到如下的一串 id：

[3, 41082, 17023, 26, 11334, 13, 24, 41082, 173, 2745, 756, 173, 2745, 4397, 173, 2745, 1893, 173, 2745, 3857, 442, 2604, 173, 973, 7880, 978, 3399, 173, 973, 10888, 978, 4582, 173, 173, 973, 309, 65, 552, 978, 6336, 4391, 173, 295, 6472, 8, ...

上面的例子展示了對單條樣本進行分詞的結果；通常我們會把分詞函式定義好，然後直接對整個資料集進行 map 就可以對整個資料集進行分詞了。

def tokenize(element):
    outputs = tokenizer(
        element["content"],
        truncation=True,
        max_length=context_length,
        return_overflowing_tokens=True,
        return_length=True,
    )
    input_batch = []
    for length, input_ids in zip(outputs["length"], outputs["input_ids"]):
        if length == context_length:
            input_batch.append(input_ids)
    return {"input_ids": input_batch}




tokenized_datasets = raw_datasets.map(
    tokenize, batched=True, remove_columns=raw_datasets["train"].column_names
)

搞定資料後，接下來就需要建立（初始化）一個模型了，GPT 的結構其實就是由 transformer 組成的，網上的輪子已經很多了，這裡就不重新造輪子了，最常見的直接用的 transformers 庫，透過配置的方式就能夠快速定義一個模型出來了。

from transformers import AutoTokenizer, GPT2LMHeadModel, AutoConfig


config = AutoConfig.from_pretrained(
    "gpt2",
    vocab_size=len(tokenizer),
    n_ctx=context_length,
    bos_token_id=tokenizer.bos_token_id,
    eos_token_id=tokenizer.eos_token_id,
)


model = GPT2LMHeadModel(config)
model_size = sum(t.numel() for t in model.parameters())
print(f"GPT-2 size: {model_size/1000**2:.1f}M parameters")

這是一個 124.2M 模型引數的 GPT2，模型的程式碼結構打給大家看看（詳細的程式碼實現可以閱讀 transformers 庫的原始碼），其實主要就是前面有個 embedding 層，中間 12 個 transformer block，最後有個線性層。

GPT2LMHeadModel(
  (transformer): GPT2Model(
    (wte): Embedding(50000, 768)
    (wpe): Embedding(1024, 768)
    (drop): Dropout(p=0.1, inplace=False)
    (h): ModuleList(
      (0): GPT2Block(
        (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
        (attn): GPT2Attention(
          (c_attn): Conv1D()
          (c_proj): Conv1D()
          (attn_dropout): Dropout(p=0.1, inplace=False)
          (resid_dropout): Dropout(p=0.1, inplace=False)
        )
        (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
        (mlp): GPT2MLP(
          (c_fc): Conv1D()
          (c_proj): Conv1D()
          (act): NewGELUActivation()
          (dropout): Dropout(p=0.1, inplace=False)
        )
      )


   #################
   中間省略重複的10層Block
   #################


      (11): GPT2Block(
        (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
        (attn): GPT2Attention(
          (c_attn): Conv1D()
          (c_proj): Conv1D()
          (attn_dropout): Dropout(p=0.1, inplace=False)
          (resid_dropout): Dropout(p=0.1, inplace=False)
        )
        (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
        (mlp): GPT2MLP(
          (c_fc): Conv1D()
          (c_proj): Conv1D()
          (act): NewGELUActivation()
          (dropout): Dropout(p=0.1, inplace=False)
        )
      )
    )
    (ln_f): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
  )
  (lm_head): Linear(in_features=768, out_features=50000, bias=False)
)

這裡也給出 GPT2 和 GPT 的模型結構圖，感興趣的同學可以仔細看看，可以發現，GPT2 的模型結構(右)較 GPT 的模型結構(左)有所改動。在 GPT2 中的一個 Transformer Block 層中，第一個 LayerNormalization 模組被移到了

Msaked-Multi-Self-Attention 模組之前, 第二個 LayerNormalization 模組也被移到了 Feed-Forward 模組之前；同時 Residual-connection 的位置也調整到了 Msaked-Multi-Self-Attention 模組與 Feed-Forward 模組之後。

資料和模型結構都確定下來後，接下來我們需要有一個訓練的流程或者框架，最簡便的那就是直接呼叫 transformers 提供的訓練器，給定一些配置，模型、分詞器、資料集。

from transformers import Trainer, TrainingArguments


args = TrainingArguments(
    output_dir="codeparrot-ds",
    per_device_train_batch_size=32,
    per_device_eval_batch_size=32,
    evaluation_strategy="steps",
    eval_steps=5_000,
    logging_steps=5_000,
    gradient_accumulation_steps=8,
    num_train_epochs=1,
    weight_decay=0.1,
    warmup_steps=1_000,
    lr_scheduler_type="cosine",
    learning_rate=5e-4,
    save_steps=5_000,
    fp16=True,
    push_to_hub=True,
)


trainer = Trainer(
    model=model,
    tokenizer=tokenizer,
    args=args,
    data_collator=data_collator,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["valid"],
)

接下來，就可以開始訓練了：

trainer.train()

自由度高一點的訓練方式也可以自行打造，依次拿到每個 batch 資料、送入模型、計算 loss、反向傳播；對於大模型來說，常見的用 accelerate 庫來進行加速，比如混合精度、梯度累積等操作。

上述的這些程式碼（使用訓練器或者 accelerate 庫進行訓練）在 transformers 的官方教程裡都有，程式碼連結：https://huggingface.co/course...。

訓練完模型後我們可以來看一下它的程式碼生成能力，那就先來跟大家 hello world 一下。比如給定 prompt：

def print_hello_world():
    """Print 'Hello World!'."""

得到的結果如下：

def print_hello_world():
    """Print 'Hello World!'."""
    print('Hello World!')

接下來，我們給定 prompt：

import numpy as np
from sklearn.ensemble import RandomForestClassifier


# create training data
X = np.random.randn(100, 100)
y = np.random.randint(0, 1, 100)


# setup train test split

它能夠幫我們劃分訓練和測試資料集：

import numpy as np
from sklearn.ensemble import RandomForestClassifier


# create training data
X = np.random.randn(100, 100)
y = np.random.randint(0, 1, 100)


# setup train test split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

這一節中我們走了一遍訓練 GPT 的一個流程，這裡和訓練 ChatGPT 的第一步的差別在於：ChatGPT 第一步採用人工寫答案的方式得到的語料對預訓練好的 GPT 進行了精調，而本節只是對一個小語料進行了一波預訓練。

2.2 訓練一個獎勵模型

接下來，我們看一下訓練 ChatGPT 的第二步：獎勵模型。

事實上，獎勵模型的學習本質就是一個句子分類或者回歸的任務。這裡找一個資料集來展示一下整個流程，比如 IMDB 影評資料集，常用於情感分析，也就是給文字的情感打分，輸入就是影評，輸出就是得分（積極或消極對應 1 或 0），那就能訓練一個獎勵模型了（比如可以用來指導 GPT 生成得分越高，也就是更加積極的文字）。

首先，匯入一下資料集，下載連結：https://ai.stanford.edu/~amaa...。

def read_imdb_split(split_dir):
    split_dir = Path(split_dir)
    texts = []
    labels = []
    for label_dir in ["pos", "neg"]:
        for text_file in (split_dir/label_dir).iterdir():
            texts.append(text_file.read_text(encoding='utf-8'))
            labels.append(0 if label_dir is "neg" else 1)
    return texts, labels
train_texts, train_labels = read_imdb_split("./DataSmall/aclImdb/train")
test_texts, test_labels = read_imdb_split("./DataSmall/aclImdb/test")

得到如下的文字，一段文字對應一個標籤：

![]()

同樣的我們需要對其進行分詞，轉為一串 id，定義分詞器，並構造資料集：

print("Tokenizing train, validate, test text ")
tokenizer = DistilBertTokenizerFast.from_pretrained('distilbert-base-uncased')
train_encodings = tokenizer(train_texts, truncation=True, padding=True)
test_encodings = tokenizer(test_texts, truncation=True, padding=True)


class IMDbDataset(torch.utils.data.Dataset):
    def __init__(self, encodings, labels):
        self.encodings = encodings
        self.labels = labels


    def __getitem__(self, idx):
        item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
        item['labels'] = torch.tensor(self.labels[idx])
        return item


    def __len__(self):
        return len(self.labels)


print("Loading tokenized text into Pytorch Datasets ")
train_dataset = IMDbDataset(train_encodings, train_labels)
test_dataset = IMDbDataset(test_encodings, test_labels)
train_loader = DataLoader(train_dataset, batch_size=10, shuffle=True)

接下來，定義一個用於預測得分的模型以及最佳化器。

print("Loading pre-trained DistilBERT model ")
model = DistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased')
model.to(device)
model.train()
optim = AdamW(model.parameters(), lr=5e-5)

然後，開始模型的訓練：

for epoch in range(3):
 for (b_ix, batch) in enumerate(train_loader):
  optim.zero_grad()
  input_ids = batch['input_ids'].to(device)
  attention_mask = batch['attention_mask'].to(device)
  labels = batch['labels'].to(device)
  outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
  loss = outputs[0]
  loss.backward()
  optim.step()

訓練完模型後，我們就可以對任意文字來進行情感的打分了，比如：

I like you,I love you

顯然這是一句很 positive 的話，模型的打分為：

詳細的細節，可以參考James D. McCaffrey 的部落格：https://jamesmccaffrey.wordpr...。

模型的選擇問題，這裡作為演示取樣了一個 distilbert 作為獎勵模型，但實際上 ChatGPT 用的獎勵模型的規模應該是和生成模型差不多大小的模型。

到目前為止，我們已經能夠訓練一個能夠生成文字的 GPT，一個能夠對文字打分的獎勵模型，接下來就需要考慮如果用獎勵模型來教會 GPT 生成高分文字了。考慮到很多同學可能沒有強化學習的基礎，這裡將插入一章強化學習的介紹。

三、強化學習 

強化學習的內容其實很多，背後涉及的數學也挺多，身邊很多人在入門過程就被勸退，為了讓大家更易於接受，這裡主要介紹一些必備的基礎知識，以及從策略梯度方法到 PPO 演算法的演進。

3.1 強化術語

在強化學習中，智慧體和環境進行互動並基於一定的獎勵機制來提高自身的水平；智慧體的決策，一般稱為策略（policy）決定了在當前環境狀態下智慧體應該去實施什麼動作；實施完動作後，智慧體會隨著環境進入下一個狀態，並獲得相應的獎勵或懲罰。這裡附上兩頁 slides 幫助大家理解。

強化學習的最終目的就是要學會一個使得智慧體能夠最大化期望回報的 policy，其中的回報就是對獎勵進行衰減求和：

和期望回報緊密相關的還有兩個概念，一個是動作狀態價值函式（觀測到狀態 ，做完決策，選中動作![]()）：

另一個是狀態價值函式（可以理解為比如下圍棋時評價當前狀態 ![]()的勝率）

瞭解了上面的定義後我們可以停下來思考一下，強化學習學什麼？

主要想學的肯定就是策略 policy 函式，也就是從狀態到動作的一個對映，如果直接學習它，那就能夠拿來使用了，這類方法也叫做基於 policy 的方法；如果採用間接點的方式，也可以學習值函式，然後根據值的大小來選擇動作，這類方法也叫做基於 value 的方法。當然，通常基於 policy 的方法也會涉及到值函式的近似。

3.2 從策略梯度到 PPO

我們要學一個策略函式，但是並不知道策略函式長什樣，怎麼去定義它才是合適的。好在有了深度學習這一工具，我們可以無腦用一個神經網路來近似策略函式，然後透過最佳化神經網路引數的方式來學習得到一個策略函式。

最佳化神經網路的引數需要有個目標函式，如果一個策略很好，那麼狀態價值的均值應當很大，因此我們定義目標函式:

這個目標函式排除掉了狀態S的因素，只依賴於策略網路π的引數Θ ；策略越好，則 ![]()越大。所以策略學習可以描述為這樣一個最佳化問題：

我們希望透過對策略網路引數Θ的更新，使得目標函式![]()越來越大，也就意味著策略 網路越來越強。想要求解最大化問題，顯然可以用梯度上升更新模型的引數。值得慶幸的是，策略函式的梯度還能被推匯出來：

策略梯度定理的詳細推導這裡就不展開了，我們需要記住的是能計算出目標函式關於引數的梯度，那就能用來更新引數，也就能學習出策略函式了。

當然這裡面還涉及動作價值函式Q的估計，如果用實際觀測的回報來近似，那就是 REINFORE 演算法，如果再用一個神經網路來近似這個價值函式，那就是演員-評論家演算法。PS：在實際使用中，策略梯度中的 Q 有多種不同的替代形式，常見效果比較好的形式是採用優勢函式A（狀態動作值函式Q減去狀態值函式V ）來替代。

傳統的策略梯度演算法的侷限性在於它是 sample-inefficient 的，也就是說每次獲取的訓練資料只被用來更新一次模型的引數後就丟掉了，因此 PPO 演算法的主要改進在於構造了新的目標函式（避免較大的引數變化），使得每次獲取的訓練資料能夠被用於多次的引數更新。

其中，比值函式為當前策略和歷史策略在狀態St下實施動作At的機率的比值：

透過這一比值也就能夠評估新舊策略的差異性，從而能夠保證策略函式在更新引數時不會跟舊策略的差異太大。有時間的同學也可以對比值在不同區間時目標函式的情況進行考慮，也就是如下表的情況。

此外，我們知道在強化學習中通常還需要去讓智慧體能夠具有 exploratory 的表現，這樣才能挖掘出更多具有高價值的行為，所以 PPO 演算法在訓練策略網路時還在目標函式中加上了和熵相關的一項獎勵H ：

這裡以離散動作為例，可以看到如果實施動作的機率分散到不同的動作上將具有更大的熵。

前面介紹策略梯度時知道策略梯度中還涉及價值函式/優勢函式的估計，在 PPO 演算法中也是採用神經網路來估計狀態價值函式，訓練價值網路的目標函式通常僅需要最小化價值網路的預測和目標的平方誤差就可以了：

綜上所述，PPO 演算法完整的最佳化目標函式由三部分組成：

3.3 PPO 演算法程式碼解讀

瞭解完原理後，我們來看一下強化的程式碼一般怎麼寫的。訓練智慧體前一般需要定義一個環境：

envs = xxxx

環境需要具有兩個主要的功能函式，一個是 step，它的輸入是動作，輸出是下一步的觀測、獎勵、以及表示環境是否結束等額外資訊：

next_obs, reward, done, info = envs.step(action)

另一個是 reset，主要用來重置環境。

然後需要定義一個智慧體，智慧體包含策略網路和價值網路兩部分（也就是演員和評論家），get_value 函式使用價值網路評估狀態的價值，get_action_and_value 函式使用策略網路給出了某個狀態下動作的機率分佈（以及對數機率）、機率分佈取樣得到的動作，機率分佈的熵、以及狀態的價值。

class Agent(nn.Module):
    def __init__(self, envs):
        super().__init__()
        self.critic = nn.Sequential(
            layer_init(nn.Linear(np.array(envs.single_observation_space.shape).prod(), 64)),
            nn.Tanh(),
            layer_init(nn.Linear(64, 64)),
            nn.Tanh(),
            layer_init(nn.Linear(64, 1), std=1.0),
        )
        self.actor = nn.Sequential(
            layer_init(nn.Linear(np.array(envs.single_observation_space.shape).prod(), 64)),
            nn.Tanh(),
            layer_init(nn.Linear(64, 64)),
            nn.Tanh(),
            layer_init(nn.Linear(64, envs.single_action_space.n), std=0.01),
        )
    def get_value(self, x):
        return self.critic(x)
    def get_action_and_value(self, x, action=None):
        logits = self.actor(x)
        probs = Categorical(logits=logits)
        if action is None:
            action = probs.sample()
        return action, probs.log_prob(action), probs.entropy(), self.critic(x)

接下來就只需要考慮如何收集資料和訓練網路了。收集資料階段主要包括兩個部分，一部分是用智慧體去和環境做互動，並儲存相應的狀態、動作等資訊，另一部分主要是根據每一步的獎勵來計算每一步的回報，從而計算用於評估動作好壞的優勢函式值。

for step in range(0, args.num_steps):
    obs[step] = next_obs
    dones[step] = next_done
    with torch.no_grad():
        action, logprob, _, value = agent.get_action_and_value(next_obs)
        values[step] = value.flatten()
    actions[step] = action
    logprobs[step] = logprob
    next_obs, reward, done, info = envs.step(action.cpu().numpy())
    rewards[step] = torch.tensor(reward).to(device).view(-1)
with torch.no_grad():
    next_value = agent.get_value(next_obs).reshape(1, -1)
    returns = torch.zeros_like(rewards).to(device)
    for t in reversed(range(args.num_steps)):
        if t == args.num_steps - 1:
            nextnonterminal = 1.0 - next_done
            next_return = next_value
        else:
            nextnonterminal = 1.0 - dones[t + 1]
            next_return = returns[t + 1]
        returns[t] = rewards[t] + args.gamma * nextnonterminal * next_return
    advantages = returns - values

訓練網路部分主要就是根據之前提到的三部分目標函式，依次計算新舊策略的差異從而計算策略網路的損失、價值網路的損失 以及 關於動作多樣性的熵獎勵；累加後再進行反向傳播就可以更新網路的引數了。

for epoch in range(args.update_epochs):
    np.random.shuffle(b_inds)
    for start in range(0, args.batch_size, args.minibatch_size):
        end = start + args.minibatch_size
        mb_inds = b_inds[start:end]
        # 計算新舊策略的差異
        _, newlogprob, entropy, newvalue = agent.get_action_and_value(b_obs[mb_inds], b_actions.long()[mb_inds])
        logratio = newlogprob - b_logprobs[mb_inds]
        ratio = logratio.exp()
        mb_advantages = b_advantages[mb_inds]
        # 策略網路損失
        pg_loss1 = -mb_advantages * ratio
        pg_loss2 = -mb_advantages * torch.clamp(ratio, 1 - args.clip_coef, 1 + args.clip_coef)
        pg_loss = torch.max(pg_loss1, pg_loss2).mean()
        # 價值網路損失
        newvalue = newvalue.view(-1)
        v_loss = 0.5 * ((newvalue - b_returns[mb_inds]) ** 2).mean()
        entropy_loss = entropy.mean()
        loss = pg_loss - args.ent_coef * entropy_loss + v_loss * args.vf_coef
        optimizer.zero_grad()
        loss.backward()
        nn.utils.clip_grad_norm_(agent.parameters(), args.max_grad_norm)
        optimizer.step()

PPO 的完整程式碼可以參考https://github.com/huggingfac...

也建議大家有空可以閱讀一下這篇文章，有助於加強對強化的理解：https://fse.studenttheses.ub....

四、利用獎勵模型強化 GPT

看到了這裡，我們已經學會啦如何訓練一個 GPT，訓練一個獎勵模型，以及強化學習演算法 PPO 的訓練過程，將它們組合起來，就能夠做出一個 ChatGPT 了！

這裡我們先定義兩個 GPT 模型，一個用於強化，一個用於參考（因為我們通常也不希望強化的模型完全朝著更高的獎勵去最佳化，所以可以約束一下強化的模型和原始模型仍然具有一定的“相似性”），參考模型不用於最佳化，所以設為 eval 模式即可。

model = GPT()
ref_model = GPT()
ref_model.eval()

需要注意的是為了適應強化這個框架，我們還需要對原來的 GPT 模型進行一定的封裝，其實主要就是加一層 value head（線性層），讓它預測每一個 token 的價值（理解為值函式，將 token 的隱藏狀態轉化為一個標量值）。

為了讓大家更加清楚地瞭解訓練的過程，我們以一條樣本為例來展示說明資料的傳遞過程。假設我們有一個 query（就是送入 gpt 的輸入），前面說過模型的輸入一般經過分詞器轉化為一串 id，比如：

![]()

送入 GPT 後，生成模型會接著輸入進行“續寫”，得到一條回答 response：

![]()

然後我們可以把這 query 和 response 拼接起來並解碼為文字後送入獎勵模型，得到一個 rewards（可以理解 chatgpt 的回答好，就意味這把提問和模型的回答拼在一起看是合理的），接下來就可以考慮使用 PPO 演算法來對 GPT 進行強化了。

首先，第一步我們就是把“提問”和“回答”拼接起來，送給獎勵模型，讓它給 GPT 的回答打個分：

texts = [q + r for q, r in zip(query, response)]
score = reward_model(texts)[1]["score"]

有了這個得分後，就可以讓 GPT 模型朝著儘可能高分的方向去最佳化了。

把拼接後的文字送入模型，並取出對應 token 的（對數）機率，這裡並不需要梯度的傳播，主要是獲得“舊”的動作（模型的回答），以及用於計算每一個 token 的獎勵：

with torch.no_grad():
    logits, _, v = self.model(**input_kwargs)
    ref_logits, _, _ = self.ref_model(**input_kwargs)
old_logprobs = logprobs_from_logits(logits[:, :-1, :], input_ids[:, 1:])
ref_logprobs = logprobs_from_logits(ref_logits[:, :-1, :], input_ids[:, 1:])

計算獎勵這裡考慮兩個部分，一部分是用於強化的 GPT 模型的回答和參考模型回答的 KL 散度（如前面所說，不能讓模型一味朝著獎勵高的方向最佳化），另一部分就是獎勵模型給出的評分：

kl = old_logprobs - ref_logprob
reward = -kl
reward[-1] += score

接下來進行模型的前向傳播（這裡是需要梯度的），並索引出 response 部分的動作機率和值函式：

logits, _, vpred = self.model(**input_kwargs)
logprob = logprobs_from_logits(logits[:, :-1, :], model_input[:, 1:])
logprob, vpred = logprob[:, -gen_len:], vpred[:, -gen_len:]

計算 loss 之前我們需要估計一下優勢函式和回報，優勢函式用於更新策略網路，回報用於更新價值網路，這裡採用經典的 GAE 方法（可以有效降低策略梯度的估計方差）來估計優勢函式：

lastgaelam = 0
advantages_reversed = []
for t in reversed(range(gen_len)):
    nextvalues = values[:, t + 1] if t < gen_len - 1 else 0.0
    delta = rewards[:, t] + self.config.gamma * nextvalues - values[:, t]
    lastgaelam = delta + self.config.gamma * self.config.lam * lastgaelam
    advantages_reversed.append(lastgaelam)
advantages = torch.stack(advantages_reversed[::-1]).transpose(0, 1)
returns = advantages + values

計算策略網路的損失：

ratio = torch.exp(logprob - old_logprobs)
pg_losses = -advantages * torch.clamp(ratio, 1.0 - self.config.cliprange, 1.0 + self.config.cliprange)
pg_loss = torch.mean(pg_losses)

價值函式的損失：

vf_losses = (vpred - returns) ** 2
vf_loss = 0.5 * torch.mean(vf_losses)

整體誤差反向傳播：

loss = pg_loss + self.config.vf_coef * vf_loss
optimizer.zero_grad()
accelerator.backward(loss)
optimizer.step()

這樣就完成了使用 PPO 來強化 GPT 的一個 step 了，也就是 ChatGPT 實現的核心思想。

可以看上面的透過強化 GPT 再用於生成的一個效果：

以及將結果再用獎勵模型打分的前後對比：

mean:
rewards (before)    0.156629
rewards (after)     1.686487
median:
rewards (before)   -0.547091
rewards (after)     2.479868

可以看到 GPT 生成結果的整體得分都有比較大的提升，這也說明了如果獎勵模型訓練得好，那用來做 ChatGPT 的效果自然也就能夠大幅度提高。

本章的完整程式碼可以參考開源專案：https://github.com/lvwerra/trl

參考：

• https://d4mucfpksywv.cloudfro...
• https://s3-us-west-2.amazonaw...
• https://arxiv.org/pdf/2005.14...
• https://openai.com/blog/chatgpt/
• https://huggingface.co/docs/t...
• https://huggingface.co/blog/rlhf
• https://huggingface.co/blog/d...
• https://fse.studenttheses.ub....
• https://github.com/lvwerra/trl
• https://github.com/wangshusen...