Improving Reasoning with
Inference-Time Scaling

Making AI 3× Better at Math Without Retraining

Chapter 4 from "Build a Reasoning Model (From Scratch)"
Sebastian Raschka

What Are Reasoning Models?

The Challenge: Teaching AI to Think Step-by-Step

What is Inference-Time Scaling?

Improving Performance During Generation

Chain-of-Thought Results

To solve the problem, we need to find the value of x...

Step 1: Set up the equation
1/2(3x - 9) = x + 37

Step 2: Eliminate the fraction
Multiply both sides by 2:
3x - 9 = 2x + 74

Step 3: Solve for x
Subtract 2x from both sides:
x - 9 = 74
Add 9 to both sides:
x = 83

Final Answer: \boxed{83}

✓ Correct!

Problem with Deterministic Decoding

We're Stuck with One Answer

How LLMs Select the Next Token

Temperature Scaling

Rescaling Logits to Control Randomness

def scale_logits_by_temperature(logits, temperature):
    if temperature <= 0:
        raise ValueError("Temperature must be positive")
    return logits / temperature

Temperature Effects

From Logits to Probabilities

Sampling Instead of Argmax

rescaled_logits = logits / temperature

probabilities = torch.softmax(rescaled_logits, dim=-1)
# Now probabilities sum to 1.0

next_token = torch.multinomial(probabilities, num_samples=1)
# Each token has a chance proportional to its probability

Temperature Sampling Implementation

@torch.inference_mode()
def generate_text_temp_stream_cache(
    model, token_ids, max_new_tokens,
    eos_token_id=None, temperature=0.0
):
    model.eval()
    cache = KVCache(n_layers=model.cfg["n_layers"])
    out = model(token_ids, cache=cache)[:, -1]

    for _ in range(max_new_tokens):
        if temperature is None or temperature == 0.0:
            # Greedy decoding
            next_token = torch.argmax(out, dim=-1, keepdim=True)
        else:
            # Temperature sampling
            logits = scale_logits_by_temperature(out, temperature)
            probas = torch.softmax(logits, dim=-1)
            next_token = torch.multinomial(probas, num_samples=1)

        yield next_token
        out = model(next_token, cache=cache)[:, -1]

Problem: Temperature Alone Isn't Enough

Too Random vs Too Conservative

Top-p Sampling

Only Sample from the "Nucleus" of High-Probability Tokens

Top-p Example (p = 0.8)

Token probabilities (sorted):
Token 1: 45.4%  →  Cumulative: 45.4%  ✓ Keep
Token 2: 27.5%  →  Cumulative: 72.9%  ✓ Keep
Token 3:  8.3%  →  Cumulative: 81.2%  ✓ Keep (crosses threshold)
Token 4:  6.8%  →  Cumulative: 88.0%  ✗ Remove
Token 5:  3.7%  →  Cumulative: 91.7%  ✗ Remove
...

Result: Sample only from top 3 tokens, then renormalize

Top-p Implementation

def top_p_filter(probas, top_p):
    if top_p is None or top_p >= 1.0:
        return probas

    # Step 1: Sort by descending probability
    sorted_probas, sorted_idx = torch.sort(
        probas, dim=1, descending=True)

    # Step 2: Cumulative sum
    cumprobas = torch.cumsum(sorted_probas, dim=1)

    # Step 3: Keep tokens where prefix mass < top_p
    prefix = cumprobas - sorted_probas
    keep = prefix < top_p
    keep[:, 0] = True  # Always keep at least one

    # Step 4: Zero out and renormalize
    kept_sorted = torch.where(
        keep, sorted_probas, torch.zeros_like(sorted_probas))
    filtered = torch.zeros_like(probas).scatter(
        1, sorted_idx, kept_sorted)
    return filtered / torch.sum(filtered, dim=1, keepdim=True)

Self-Consistency Implementation

def self_consistency_vote(
    model, tokenizer, prompt, device,
    num_samples=10, temperature=0.8, top_p=0.9
):
    full_answers, short_answers = [], []

    # 1) Sample multiple answers with diversity
    for i in range(num_samples):
        answer = generate_text_stream_concat_flex(
            model, tokenizer, prompt, device,
            generate_func=generate_text_top_p_stream_cache,
            temperature=temperature, top_p=top_p,
        )

        # 2) Extract final answer
        short = extract_final_candidate(answer)
        full_answers.append(answer)
        short_answers.append(short)

    # 3) Vote: choose most frequent answer
    counts = Counter(short_answers)
    most_common = counts.most_common()
    final_answer = most_common[0][0] if most_common else None

    return {
        "full_answers": full_answers,
        "short_answers": short_answers,
        "counts": dict(counts),
        "final_answer": final_answer
    }

Self-Consistency Example Output

results = self_consistency_vote(
    model, tokenizer,
    prompt + "\n\nExplain step by step.",  # Use CoT!
    device=device,
    num_samples=5,
    temperature=0.8,
    top_p=0.9
)

Console Output:
[Sample 1/5] → '83'
[Sample 2/5] → '83'
[Sample 3/5] → '83'
[Sample 4/5] → '83'
[Sample 5/5] → '83'

Final answer: 83

All 5 samples converged to the correct answer!

Results: Method Comparison

All Techniques Benchmarked on MATH-500

Accuracy vs Computational Cost

Real-World Applications

Modern AI Systems Use These Techniques

These techniques are production-critical in state-of-the-art AI!

Conclusion & Key Takeaways

Thank You!

Presented by: AmirHasan Aref Asl