Sampling Algorithms for Text Generation

Implements various sampling strategies for autoregressive language model inference.

Sampling Strategies

1. Greedy Decoding (Deterministic)

next_token = argmax(logits)

Use case: Reproducible outputs, factual tasks Pros: Deterministic, fast Cons: Repetitive, lacks creativity

2. Top-k Sampling

top_k_logits, indices = topk(logits, k)
probs = softmax(top_k_logits / temperature)
next_token = sample(indices, probs)

Use case: Creative generation with controlled diversity Pros: Balances quality and diversity Cons: Fixed k may be too restrictive or permissive

3. Nucleus (Top-p) Sampling

sorted_probs = sort(softmax(logits), descending=True)
cumsum = cumulative_sum(sorted_probs)
nucleus = sorted_probs[cumsum <= p]
next_token = sample(nucleus)

Use case: Adaptive diversity based on distribution Pros: Adapts to confidence (dynamic cutoff) Cons: More compute than top-k

Temperature Scaling

Temperature τ controls randomness:

scaled_logits = logits / τ

τ → 0:   Deterministic (greedy)
τ = 1:   Normal distribution
τ > 1:   More random/creative

References

• "The Curious Case of Neural Text Degeneration" (Holtzman et al., 2019) - Nucleus sampling
• "Hierarchical Neural Story Generation" (Fan et al., 2018) - Top-k sampling
• GPT-2: Uses both top-k and top-p
• llama.cpp: common/sampling.cpp

Data Structures

structure Hesper.Inference.Sampling.SamplingConfig : Type

Sampling configuration

• temperature : Float
• topK : Nat
• topP : Float
• seed : Option Nat

instance Hesper.Inference.Sampling.instReprSamplingConfig : Repr SamplingConfig

Equations

• Hesper.Inference.Sampling.instReprSamplingConfig = { reprPrec := Hesper.Inference.Sampling.instReprSamplingConfig.repr }

def Hesper.Inference.Sampling.instReprSamplingConfig.repr : SamplingConfig → Nat → Std.Format

Equations

• One or more equations did not get rendered due to their size.

Core Utilities

def Hesper.Inference.Sampling.argmax (logits : Array Float) : Nat

Find argmax (index of maximum value)

@param logits Array of logit scores @return Index of maximum value

Equations

• Hesper.Inference.Sampling.argmax logits = if (logits.size == 0) = true then 0 else Hesper.Inference.Sampling.argmax.loop logits 1 0 logits[0]!

@[irreducible]

def Hesper.Inference.Sampling.argmax.loop (logits : Array Float) (idx maxIdx : Nat) (maxVal : Float) : Nat

Equations

• One or more equations did not get rendered due to their size.

def Hesper.Inference.Sampling.applyTemperature (logits : Array Float) (temperature : Float) : Array Float

Apply temperature scaling

@param logits Original logits @param temperature Temperature value (τ > 0) @return Scaled logits

Equations

• Hesper.Inference.Sampling.applyTemperature logits temperature = if (temperature == 1.0) = true then logits else Array.map (fun (x : Float) => x / temperature) logits

def Hesper.Inference.Sampling.softmax (logits : Array Float) : Array Float

Compute softmax probabilities

Uses numerically stable version: exp(x - max(x))

@param logits Log-probabilities @return Probabilities (sum to 1.0)

Equations

• One or more equations did not get rendered due to their size.

Greedy Sampling

def Hesper.Inference.Sampling.sampleGreedy (logits : Array Float) : Nat

Greedy sampling: select token with highest probability

Algorithm:

next_token = argmax(logits)

@param logits Logit scores [vocab_size] @return Selected token ID

Equations

• Hesper.Inference.Sampling.sampleGreedy logits = Hesper.Inference.Sampling.argmax logits

Top-k Sampling

def Hesper.Inference.Sampling.topK (logits : Array Float) (k : Nat) : Array Float × Array Nat

Find indices of top-k largest values

@param logits Array of scores @param k Number of top elements @return (top_k_values, top_k_indices)

Equations

• One or more equations did not get rendered due to their size.

@[irreducible]

def Hesper.Inference.Sampling.topK.buildIndexed (logits : Array Float) (idx : Nat) (acc : Array (Float × Nat)) : Array (Float × Nat)

Equations

• Hesper.Inference.Sampling.topK.buildIndexed logits idx acc = if idx ≥ logits.size then acc else Hesper.Inference.Sampling.topK.buildIndexed logits (idx + 1) (acc.push (logits[idx]!, idx))

def Hesper.Inference.Sampling.categoricalSample (probs : Array Float) (randomValue : Float) : Nat

Sample from categorical distribution

Simple implementation using linear search through cumulative probabilities.

@param probs Probability distribution (must sum to ~1.0) @param randomValue Random value in [0, 1) @return Sampled index

Equations

• Hesper.Inference.Sampling.categoricalSample probs randomValue = Hesper.Inference.Sampling.categoricalSample.loop probs randomValue 0 0.0

@[irreducible]

def Hesper.Inference.Sampling.categoricalSample.loop (probs : Array Float) (randomValue : Float) (idx : Nat) (cumulative : Float) : Nat

Equations

• One or more equations did not get rendered due to their size.

def Hesper.Inference.Sampling.sampleTopK (logits : Array Float) (k : Nat) (temperature randomValue : Float) : Nat

Top-k sampling: sample from k tokens with highest probability

Algorithm:

1. Find top-k logits and their indices
2. Apply temperature scaling
3. Compute softmax over top-k
4. Sample from categorical distribution

@param logits Logit scores [vocab_size] @param k Number of top candidates @param temperature Sampling temperature @param randomValue Random number in [0, 1) for sampling @return Selected token ID

Equations

• One or more equations did not get rendered due to their size.

Nucleus (Top-p) Sampling

def Hesper.Inference.Sampling.sampleNucleus (logits : Array Float) (p temperature randomValue : Float) : Nat

Nucleus (top-p) sampling: sample from smallest set with cumulative probability >= p

Algorithm:

1. Sort logits in descending order
2. Compute softmax probabilities
3. Find nucleus: tokens where cumsum(probs) <= p
4. Renormalize and sample from nucleus

@param logits Logit scores [vocab_size] @param p Cumulative probability threshold (typically 0.9) @param temperature Sampling temperature @param randomValue Random number in [0, 1) for sampling @return Selected token ID

Equations

• One or more equations did not get rendered due to their size.

@[irreducible]

def Hesper.Inference.Sampling.sampleNucleus.buildIndexed (logits : Array Float) (idx : Nat) (acc : Array (Float × Nat)) : Array (Float × Nat)

Equations

• One or more equations did not get rendered due to their size.

@[irreducible]

def Hesper.Inference.Sampling.sampleNucleus.findNucleusSize (p : Float) (probs : Array Float) (idx : Nat) (cumulative : Float) : Nat

Equations

• One or more equations did not get rendered due to their size.

Repetition Penalty

def Hesper.Inference.Sampling.applyRepetitionPenalty (logits : Array Float) (prevTokens : Array Nat) (penalty : Float) : Array Float

Apply repetition penalty to logits. Tokens that appeared in prevTokens have their logits divided (if positive) or multiplied (if negative) by the penalty factor. penalty = 1.0 means no penalty. Typical values: 1.1 - 1.3.

Equations

• One or more equations did not get rendered due to their size.

High-Level Sampling Interface

inductive Hesper.Inference.Sampling.Strategy : Type

Sampling strategy

• Greedy : Strategy
• TopK (k : Nat) (temperature : Float) : Strategy
• Nucleus (p temperature : Float) : Strategy

def Hesper.Inference.Sampling.instReprStrategy.repr : Strategy → Nat → Std.Format

Equations

• One or more equations did not get rendered due to their size.

instance Hesper.Inference.Sampling.instReprStrategy : Repr Strategy

Equations

• Hesper.Inference.Sampling.instReprStrategy = { reprPrec := Hesper.Inference.Sampling.instReprStrategy.repr }

instance Hesper.Inference.Sampling.instToStringStrategy : ToString Strategy

Equations

• One or more equations did not get rendered due to their size.

def Hesper.Inference.Sampling.sample (logits : Array Float) (strategy : Strategy) (randomValue : Float := 0.5) : Nat

Sample next token using specified strategy

@param logits Logit scores [vocab_size] @param strategy Sampling strategy @param randomValue Random value in [0, 1) (only used for stochastic strategies) @return Selected token ID

Equations

• Hesper.Inference.Sampling.sample logits Hesper.Inference.Sampling.Strategy.Greedy randomValue = Hesper.Inference.Sampling.sampleGreedy logits
• Hesper.Inference.Sampling.sample logits (Hesper.Inference.Sampling.Strategy.TopK a a_1) randomValue = Hesper.Inference.Sampling.sampleTopK logits a a_1 randomValue
• Hesper.Inference.Sampling.sample logits (Hesper.Inference.Sampling.Strategy.Nucleus a a_1) randomValue = Hesper.Inference.Sampling.sampleNucleus logits a a_1 randomValue

Random Number Generation

def Hesper.Inference.Sampling.lcgNext (seed : UInt32) : Float × UInt32

Simple linear congruential generator (LCG) for reproducible sampling

Parameters: Numerical Recipes (a = 1664525, c = 1013904223, m = 2^32)

@param seed Current seed @return (random_float in [0,1), new_seed)

Equations

• Hesper.Inference.Sampling.lcgNext seed = ((1664525 * seed + 1013904223).toFloat / UInt32.size.toFloat, 1664525 * seed + 1013904223)

structure Hesper.Inference.Sampling.RNG : Type

Random number generator state

• seed : UInt32

def Hesper.Inference.Sampling.instReprRNG.repr : RNG → Nat → Std.Format

Equations

• One or more equations did not get rendered due to their size.

instance Hesper.Inference.Sampling.instReprRNG : Repr RNG

Equations

• Hesper.Inference.Sampling.instReprRNG = { reprPrec := Hesper.Inference.Sampling.instReprRNG.repr }

def Hesper.Inference.Sampling.RNG.create (seed : Option Nat := none) : RNG

Create RNG from optional seed

Equations

• Hesper.Inference.Sampling.RNG.create (some n) = { seed := n.toUInt32 }
• Hesper.Inference.Sampling.RNG.create = { seed := 42 }

def Hesper.Inference.Sampling.RNG.next (rng : RNG) : Float × RNG

Generate next random float

Equations

• rng.next = match Hesper.Inference.Sampling.lcgNext rng.seed with | (randomFloat, newSeed) => (randomFloat, { seed := newSeed })

Sampling with RNG

def Hesper.Inference.Sampling.sampleWithRNG (logits : Array Float) (strategy : Strategy) (rng : RNG) : Nat × RNG

Sample with automatic RNG threading

@param logits Logit scores @param strategy Sampling strategy @param rng Random number generator state @return (selected_token, new_rng)

Equations

• One or more equations did not get rendered due to their size.
• Hesper.Inference.Sampling.sampleWithRNG logits Hesper.Inference.Sampling.Strategy.Greedy rng = (Hesper.Inference.Sampling.sampleGreedy logits, rng)

Utilities

def Hesper.Inference.Sampling.printStats (logits : Array Float) (selectedToken : Nat) : IO Unit

Print sampling statistics

@param logits Logit scores @param selectedToken Sampled token

Equations

• One or more equations did not get rendered due to their size.