采样策略 - 从零开始大模型

前面十章，我们训练了一个 GPT 模型。接下来的第 11 到 14 章，我们研究一下怎么使用模型。

GPT 模型的基本功能是生成文本，但是如何生成，还有很多讲究。

最简单的办法是贪心搜索（Greedy Search）：每一步都选概率最大的词元。但是这样生成的文本往往很死板，容易陷入重复。

解决的办法是随机采样（Sampling）：按概率分布随机选择词元。但是完全随机又可能生成不连贯的内容。

下面是几种常用的采样策略，在生成质量和多样性之间寻找平衡。

温度采样¶

温度（Temperature）是最简单的采样控制参数。在 Softmax 之前，把逻辑值除以温度 $T$ ：

p_i = \frac{\exp(z_i / T)}{\sum_j \exp(z_j / T)}

(1)

$T = 1$ ：标准采样，按照原始概率分布选择；
$T = 0$ ：贪心搜索，总是选概率最大的；
$T > 1$ ：概率分布更平坦，随机性更强；
$T < 1$ ：概率分布更尖锐，确定性更强。

Top-k 采样¶

Top-k 采样只从概率最高的 $k$ 个词元中采样，其余词元的概率设为 0。

优点：过滤掉长尾的噪音词元，提高生成质量；
缺点： $k$ 是固定的，不能适应不同的上下文。

Top-p 采样¶

Top-p 采样从累积概率超过 $p$ 的最小词元集合中采样。

优点：候选词元集合大小自适应，概率分布平坦时选择范围大，尖锐时选择范围小；
缺点：算法稍复杂。

组合使用¶

实践中，通常组合使用多种策略：

温度 + Top-k：先调温度，再过滤；
温度 + Top-p：先调温度，再自适应过滤；
温度 + Top-k + Top-p：三重控制。

import numpy as np

from src.core import *
from src.gpt import *

np.random.seed(42)

GPT 采样训练器¶

GPT 训练器的 generate 函数支持按原始概率分布采样。

GPT 采样训练器 GPTTrainerWithSampling 继承自 GPT 训练器，扩展了 generate 函数：

temperature（温度）：根据温度调整概率分布。温度越高，随机性越强；温度越低，确定性越强；
top_k：找到第 $k$ 大的值作为阈值，低于阈值的设置为 $-\infty$ ；
top_p：排序后计算累积概率，找到超过 $p$ 的截断点，截断点后的设置为 $-\infty$ 。

温度、top-k、top-p 可以组合使用，顺序是：先温度，再 top-k，最后 top-p。

class GPTTrainerWithSampling(GPTTrainer):

    def generate(self, dataset, prompt, steps=256, temperature=1.0, top_k=None, top_p=None):
        # 推理模式
        self.model.eval()
        tokens = dataset.encode(prompt)

        # 暂停训练
        with Tensor.no_grad():
            for _ in range(steps):
                # 推理
                feature = Tensor([tokens[-dataset.context_size:]])
                logits = self.model(feature)
                last_logits = logits.data[0, -1]

                # 贪心搜索
                if temperature == 0:
                    probs = np.zeros(len(last_logits))
                    probs[np.argmax(last_logits)] = 1.0
                else:
                    # 温度采样
                    last_logits /= temperature

                    # Top-k 采样
                    if top_k is not None:
                        threshold = np.partition(last_logits, -top_k)[-top_k]
                        last_logits = np.where(last_logits < threshold, -np.inf, last_logits)

                    # Top-p 采样
                    if top_p is not None:
                        sorted_indices = np.argsort(last_logits)[::-1]
                        sorted_logits = last_logits[sorted_indices]
                        exp = np.exp(sorted_logits - np.max(sorted_logits))
                        cumulative_probs = np.cumsum(exp / np.sum(exp))
                        cutoff = np.searchsorted(cumulative_probs, top_p)
                        last_logits[sorted_indices[cutoff + 1:]] = -np.inf

                    exp = np.exp(last_logits - np.max(last_logits))
                    probs = exp / np.sum(exp)

                token = np.random.choice(len(probs), p=probs)
                tokens.append(token)

        return dataset.decode(tokens)

超参数¶

模型参数¶

模型参数需要保持和预训练阶段一致。

BATCH_SIZE = 16
CONTEXT_SIZE = 256
EMBEDDING_SIZE = 384
HEADS = 6
BLOCKS = 6

建模¶

加载模型¶

DATA_FILE = "../../../../data/tinyshakespeare.txt"
MODEL_FILE = "../../../../data/tinyshakespeare-gpt.npz"

dataset = CharDataset(DATA_FILE, BATCH_SIZE, CONTEXT_SIZE)
model = GPT(dataset.vocab_size, CONTEXT_SIZE, EMBEDDING_SIZE, HEADS, BLOCKS)
loss_fn = GPTLoss()
optimizer = AdamWOptimizer(model.parameters)
trainer = GPTTrainerWithSampling(model, loss_fn, optimizer)
trainer.load(MODEL_FILE)

贪心搜索¶

print(trainer.generate(dataset, prompt="ROMEO:", temperature=0.0))

ROMEO:
What says the senate of the court?

PETRUCHIO:
I will not stay the state of the country.

POMPEY:
I would not be so so so so so so so so so this son,
And the shall be the straight of the death.

POLIXENES:
I will not be so.

POMPEY:
I would not be so so s

温度采样¶

print(trainer.generate(dataset, prompt="ROMEO:", temperature=0.8))

ROMEO:
He may the body majesty. But, thy state, that I had
men my true princely from our grown.

QUEEN MARGARET:
Unferal the news have been and lords these better
What he sates on the best and will place
That set I think to be a repus before.

QUEEN ELIZABETH:
M

Top-k 采样¶

print(trainer.generate(dataset, prompt="ROMEO:", temperature=1.0, top_k=5))

ROMEO:
What's my mistress. They say, the senate is to
hold my bear the will be come to be the sometime
As a seat, are any bearing treads,
We are nother'd their severences, by me,
As it brace off this bloody, and worthy
To stay the perform'd. The darest is nothin

Top-p 采样¶

print(trainer.generate(dataset, prompt="ROMEO:", temperature=1.0, top_p=0.9))

ROMEO:
Ay, and no more of in this looks: but I not
fare not to book'd in out thy body
To fight, that was for your honours and may consul!

SLY:
He lost the company of my death, some party,
The shall be stop and stay what now means are mother;
For honour provides

组合采样¶

print(trainer.generate(dataset, prompt="ROMEO:", temperature=0.9, top_k=5, top_p=0.9))

ROMEO:
I saw the did three thy banish'd.

POMPEY:
I have did my lord, I speak a stop, and I say
The consul: that have my loved in their
Thou wish'd of mine and that she wars and with their sease,
Thou art the bearing borness and sentender to
bring the country to

本章，我们实现了三种采样策略：温度控制随机性，top_k 过滤长尾，top_p 自适应候选集。组合使用还可以在生成质量和多样性之间灵活平衡。

课后练习¶

试着调整 temperature、top_k、top_p 的组合，观察生成文本的变化。什么组合最适合莎士比亚风格剧本的生成？