Sampling Strategies - From Neuron to GPT

In the previous ten chapters, we trained a GPT model. In Chapters 11 to 14, we study how to use the model.

The basic function of a GPT model is to generate text, but how to generate it takes some care.

The simplest approach is greedy search: at every step, pick the token with the highest probability. But the text generated this way is often rigid and easily falls into repetition.

The solution is sampling: pick tokens at random according to the probability distribution. But complete randomness may generate incoherent content.

Below are several commonly used sampling strategies, which seek a balance between generation quality and diversity.

Temperature sampling¶

Temperature is the simplest parameter for controlling sampling. Before the Softmax, divide the logits by the temperature $T$ :

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

(1)

$T = 1$ : standard sampling, choosing by the original probability distribution;
$T = 0$ : greedy search, always choosing the most probable token;
$T > 1$ : the probability distribution is flatter and more random;
$T < 1$ : the probability distribution is sharper and more deterministic.

Top-k sampling¶

Top-k sampling samples only from the $k$ tokens with the highest probability and sets the probability of all other tokens to 0.

Advantage: filters out noisy tokens in the long tail and improves generation quality;
Disadvantage: $k$ is fixed and cannot adapt to different contexts.

Top-p sampling¶

Top-p sampling samples from the smallest set of tokens whose cumulative probability exceeds $p$ .

Advantage: the size of the candidate set adapts, with a wide range when the distribution is flat and a narrow range when it is sharp;
Disadvantage: the algorithm is slightly more complex.

Combined use¶

In practice, several strategies are usually combined:

temperature + top-k: adjust the temperature first, then filter;
temperature + top-p: adjust the temperature first, then filter adaptively;
temperature + top-k + top-p: triple control.

import numpy as np

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

np.random.seed(42)

GPT sampling trainer¶

class GPTTrainerWithSampling(GPTTrainer):

    def generate(self, dataset, prompt, steps=256, temperature=1.0, top_k=None, top_p=None):
        # Inference mode
        self.model.eval()
        tokens = dataset.encode(prompt)

        # Pause training (disable gradients)
        with Tensor.no_grad():
            for _ in range(steps):
                # Inference
                feature = Tensor([tokens[-dataset.context_size:]])
                logits = self.model(feature)
                last_logits = logits.data[0, -1]

                # Greedy search
                if temperature == 0:
                    probs = np.zeros(len(last_logits))
                    probs[np.argmax(last_logits)] = 1.0
                else:
                    # Temperature sampling
                    last_logits /= temperature

                    # Top-k sampling
                    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 sampling
                    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)

Hyperparameters¶

Model parameters¶

The model parameters must stay consistent with the pre-training stage.

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

Modeling¶

Loading the model¶

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)

Greedy search¶

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

Temperature sampling¶

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 sampling¶

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 sampling¶

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

Combined sampling¶

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

In this chapter, we implemented three sampling strategies: temperature controls the randomness, top_k filters the long tail, and top_p adapts the candidate set. Used together, they can also balance generation quality and diversity flexibly.

Exercises¶

Try adjusting the combination of temperature, top_k, and top_p and observe how the generated text changes. Which combination is best for generating Shakespeare-style scripts?