Predicting the next token

Count context and select a model before the final test

natural language generation

next-token prediction

language models

Learn how a language model assigns probabilities to the next token. Evaluate a count model whose values readers can examine.

Nia builds an autocomplete demonstration for an archive. After a reader types we, the tool must decide which token could come next.

A language model assigns probabilities to token sequences. In next-token prediction, the input is the context already seen and the output is a probability for each possible next token. This is one language-model training objective. It is not a definition of every language model or every text-generation system.

What you will learn

This lesson shows how to:

keep training, validation, and test speeches separate
build next-token rows without crossing paragraph boundaries
fit unigram, bigram, and trigram count models
compare greedy choice with probability-weighted sampling
select context length and interpolation strength with validation perplexity
explain what held-out perplexity does not measure.

Load the split and study results

The inaugural corpus contains 60 public-domain speeches. The committed study uses 37 speeches for training, 10 for model selection, and 13 for the final test. A speech belongs to one split only. Assignment by era places both eras in every part. The split separates speeches, not speakers. Presidents with more than one inaugural address can appear in multiple splits.

library(digest)
library(dplyr)
library(tidyr)
library(tibble)
library(stringr)
library(purrr)
library(readr)

source("R/inaugural-corpus.R")

split_path <- "data/inaugural/next-token-splits.csv"
validation_path <- "data/inaugural/next-token-validation.csv"
test_path <- "data/inaugural/next-token-test.csv"
metadata_path <- "data/inaugural/next-token-study-metadata.csv"

split_assignments <- read_csv(
  split_path,
  na = character(),
  col_types = cols(
    speech_id = col_character(),
    era = col_character(),
    split = col_character()
  )
)
validation_results <- read_csv(
  validation_path,
  na = c("", "NA"),
  col_types = cols(
    split = col_character(),
    model = col_character(),
    context_tokens = col_integer(),
    interpolation_strength = col_double(),
    next_token_rows = col_integer(),
    oov_target_share = col_double(),
    seen_context_share = col_double(),
    mean_log_loss = col_double(),
    perplexity = col_double(),
    selected = col_logical(),
    selection_reason = col_character()
  )
)
test_results <- read_csv(
  test_path,
  na = c("", "NA"),
  col_types = cols(
    split = col_character(),
    model = col_character(),
    context_tokens = col_integer(),
    interpolation_strength = col_double(),
    next_token_rows = col_integer(),
    oov_target_share = col_double(),
    seen_context_share = col_double(),
    mean_log_loss = col_double(),
    perplexity = col_double(),
    top_1_accuracy = col_double(),
    role = col_character()
  )
)
test_by_speech <- read_csv(
  "data/inaugural/next-token-test-by-speech.csv",
  na = character(),
  col_types = cols(
    speech_id = col_character(),
    next_token_rows = col_integer(),
    model = col_character(),
    perplexity = col_double()
  )
)
study_metadata <- read_csv(
  metadata_path,
  na = character(),
  col_types = cols(
    artifact = col_character(),
    description = col_character(),
    settings = col_character(),
    source = col_character(),
    license = col_character(),
    built_on = col_character(),
    fingerprint = col_character()
  )
)

split_counts <- split_assignments |>
  mutate(
    split = factor(
      split,
      levels = c("training", "validation", "test")
    ),
    era = factor(
      era,
      levels = c("before 1900", "1900 or later")
    )
  ) |>
  count(split, era, name = "speeches") |>
  arrange(split, era)

Speech-level assignments used for next-token modeling
Split	Era	Speeches
training	before 1900	17
training	1900 or later	20
validation	before 1900	5
validation	1900 or later	5
test	before 1900	6
test	1900 or later	7

The study uses validation speeches to select a candidate model. It uses test speeches only after that choice is final. Calling both sets “held out” would obscure their different purposes.

Build paragraph-bounded next-token rows

The study uses an explicit word-like pattern so token counts do not depend on the rendering machine’s word-boundary library. It then keeps lowercase ASCII alphabetic tokens. The 2,499 most common training token types receive separate vocabulary entries. Every other token maps to <unk>, the unknown-token entry. The model vocabulary therefore contains 2,500 tokens.

Each prediction row keeps its paragraph and speech ID. previous_1 is the token immediately before the target. previous_2 is the token before previous_1. The first two tokens in each paragraph lack a full two-token context, so the evaluation excludes those rows.

No transition joins the end of one paragraph to the beginning of another. This model therefore does not learn how to start a paragraph.

paragraphs <- inaugural_paragraphs()
token_pattern <- paste0(
  "(?:[a-z]\\.){2,}|",
  "[a-z0-9]+(?:['\\x{2019}][a-z0-9]+)*"
)

training_tokens <- paragraphs |>
  select(paragraph_id, speech_id, paragraph) |>
  inner_join(
    split_assignments |>
      filter(split == "training") |>
      select(speech_id),
    by = join_by(speech_id)
  ) |>
  mutate(
    token = str_extract_all(
      str_to_lower(paragraph, locale = "en"),
      regex(token_pattern)
    )
  ) |>
  select(-paragraph) |>
  unnest_longer(token) |>
  filter(str_detect(token, "^[a-z]+$")) |>
  mutate(position = row_number(), .by = paragraph_id)

training_vocabulary <- training_tokens |>
  count(token, sort = TRUE) |>
  slice_head(n = 2499L) |>
  pull(token)
model_vocabulary <- c(training_vocabulary, "<unk>")

training_next_tokens <- training_tokens |>
  mutate(
    model_token = if_else(
      token %in% training_vocabulary,
      token,
      "<unk>"
    )
  ) |>
  arrange(paragraph_id, position) |>
  mutate(
    previous_1 = lag(model_token, 1L),
    previous_2 = lag(model_token, 2L),
    .by = paragraph_id
  ) |>
  filter(position >= 3L) |>
  transmute(
    paragraph_id,
    speech_id,
    previous_2,
    previous_1,
    next_token = model_token
  )

training_summary <- tibble(
  measure = c(
    "training speeches",
    "training paragraphs",
    "training tokens",
    "next-token rows",
    "model vocabulary"
  ),
  value = c(
    n_distinct(training_next_tokens$speech_id),
    n_distinct(training_tokens$paragraph_id),
    nrow(training_tokens),
    nrow(training_next_tokens),
    length(model_vocabulary)
  )
)

Training data for the inspectable count model
Training item	Value
training speeches	37
training paragraphs	867
training tokens	83458
next-token rows	81724
model vocabulary	2500

Count what followed each token

A unigram ignores context and uses overall token frequency. A bigram uses one previous token. A trigram uses two. More context can help, but each longer context appears less often.

Raw counts assign probability zero to a next token that never followed the context in training. The study avoids zero probabilities through interpolation, which combines longer-context counts with the shorter model. The interpolation strength controls the weight of the shorter model. A larger value gives that model more weight. The study compares six values on validation speeches. The selected trigram uses 100.

selected_validation <- validation_results |>
  filter(selected)
interpolation_strength <- selected_validation$interpolation_strength

unigram_counts <- training_tokens |>
  mutate(
    model_token = if_else(
      token %in% training_vocabulary,
      token,
      "<unk>"
    )
  ) |>
  count(model_token, name = "unigram_count")
training_token_total <- sum(unigram_counts$unigram_count)

unigram_probabilities <- tibble(
  model_token = model_vocabulary
) |>
  left_join(unigram_counts, by = join_by(model_token)) |>
  mutate(
    unigram_count = replace_na(unigram_count, 0L),
    unigram_probability = (unigram_count + 1) /
      (training_token_total + length(model_vocabulary))
  )

bigram_counts <- training_next_tokens |>
  count(previous_1, next_token, name = "bigram_count")
trigram_counts <- training_next_tokens |>
  count(
    previous_2,
    previous_1,
    next_token,
    name = "trigram_count"
  )

predict_after_one <- function(previous_token) {
  context_counts <- bigram_counts |>
    filter(previous_1 == previous_token)

  tibble(next_token = model_vocabulary) |>
    left_join(
      unigram_probabilities,
      by = join_by(next_token == model_token)
    ) |>
    left_join(context_counts, by = join_by(next_token)) |>
    mutate(
      bigram_count = replace_na(bigram_count, 0L),
      context_count = sum(bigram_count),
      bigram_probability = (
        bigram_count +
          interpolation_strength * unigram_probability
      ) / (
        context_count +
          interpolation_strength
      )
    ) |>
    arrange(desc(bigram_probability), next_token)
}

predict_after_two <- function(previous_2_token, previous_1_token) {
  context_counts <- trigram_counts |>
    filter(
      previous_2 == previous_2_token,
      previous_1 == previous_1_token
    )

  predict_after_one(previous_1_token) |>
    left_join(context_counts, by = join_by(next_token)) |>
    mutate(
      trigram_count = replace_na(trigram_count, 0L),
      context_count = sum(trigram_count),
      probability = (
        trigram_count +
          interpolation_strength * bigram_probability
      ) / (
        context_count +
          interpolation_strength
      ),
      context = str_c(
        previous_2_token,
        previous_1_token,
        sep = " "
      )
    ) |>
    arrange(desc(probability), next_token)
}

context_predictions <- bind_rows(
  predict_after_two("we", "must") |>
    slice_head(n = 5),
  predict_after_two("the", "united") |>
    slice_head(n = 5)
) |>
  select(context, next_token, trigram_count, probability)

Highest-probability next tokens from the selected trigram
Previous two tokens	Possible next token	Training count	Probability
we must	be	9	0.1725
we must	<unk>	8	0.0864
we must	not	5	0.0474
we must	do	5	0.0380
we must	show	4	0.0288
the united	states	112	0.7189
the united	nations	6	0.0386
the united	<unk>	2	0.0275
the united	and	1	0.0172
the united	the	0	0.0139

After the united, the model assigns about 0.72 probability to states. That reflects a repeated phrase in US inaugural addresses. It does not show a general English rule. After we must, several continuations remain plausible.

The <unk> row represents any token outside the training vocabulary. Predicting that category does not recover the original word.

Select a next token

The probability table does not specify a text-generation rule. Greedy decoding takes the highest-probability token. Sampling draws according to the full distribution. The seed below makes five draws reproducible.

greedy_next <- we_must_distribution$next_token[[1]]

set.seed(6303)
sampled_next <- sample(
  we_must_distribution$next_token,
  size = 5,
  replace = TRUE,
  prob = we_must_distribution$probability
)

choice_table <- tibble(
  method = c("greedy", "five seeded samples"),
  output = c(
    greedy_next,
    str_c(sampled_next, collapse = ", ")
  )
)

Different choice rules can use one probability distribution
Choice rule	Next-token output after ‘we must’
greedy	be
five seeded samples	be, races, maintain, strive, change

Greedy choice always selects the most common continuation. Sampling can select less common tokens and will change when the seed changes. Neither rule establishes whether a continuation is factual, useful, or appropriate for a reader.

Select context length and smoothing before the final test

The builder compares six interpolation strengths for both the bigram and trigram, plus the unigram baseline, on validation speeches. Log loss penalizes probability assigned away from the observed next token. Perplexity is the exponentiated mean log loss. Lower values are better when the tokenizer, vocabulary, and evaluation rows are identical.

validation_display <- validation_results |>
  filter(
    model == "unigram" |
      (model == "bigram" & interpolation_strength == 100) |
      (
        model == "trigram" &
          interpolation_strength %in% c(10, 100, 1000)
      )
  ) |>
  arrange(context_tokens, interpolation_strength) |>
  transmute(
    model,
    context_tokens,
    interpolation_strength = if_else(
      is.na(interpolation_strength),
      "not applicable",
      as.character(interpolation_strength)
    ),
    seen_context = if_else(
      is.na(seen_context_share),
      "not applicable",
      sprintf("%.1f%%", 100 * seen_context_share)
    ),
    perplexity = round(perplexity, 1),
    selected = if_else(selected, "yes", "no")
  )

Validation settings around the selected trigram
Candidate	Context tokens	Interpolation strength	Validation contexts seen in training	Validation perplexity	Selected
unigram	0	not applicable	not applicable	238.0	no
bigram	1	100	100.0%	134.8	no
trigram	2	10	74.9%	194.4	no
trigram	2	100	74.9%	133.3	yes
trigram	2	1000	74.9%	143.6	no

The trigram with interpolation strength 100 has the lowest validation perplexity, 133.3. The best bigram is close at 134.8. The compared strengths above and below 100 score worse. The selected value is therefore not a boundary of the candidate grid.

Only 74.9 percent of the trigram’s two-token validation contexts appeared in training. When a context is unseen, its estimate relies entirely on the bigram and then the unigram.

The bigram’s 100 percent seen-context rate results from mapping every token into the training vocabulary, which contains <unk>. It does not show that the model generalizes to every word.

Open the test speeches once

The study selects the trigram and its interpolation strength before the final test. It then scores the trigram, the predeclared unigram baseline, and a uniform 2,500-token comparison model on the 13 test speeches.

test_display <- test_results |>
  transmute(
    model,
    role,
    next_token_rows,
    oov_targets = sprintf("%.1f%%", 100 * oov_target_share),
    perplexity = round(perplexity, 1),
    top_1_accuracy = if_else(
      is.na(top_1_accuracy),
      "not applicable",
      sprintf("%.1f%%", 100 * top_1_accuracy)
    )
  )

perplexity_reduction <- 1 -
  test_results$perplexity[test_results$model == "trigram"] /
    test_results$perplexity[test_results$model == "unigram"]

test_comparison_by_speech <- test_by_speech |>
  select(speech_id, model, perplexity) |>
  pivot_wider(
    names_from = model,
    values_from = perplexity
  )

speech_people <- paragraphs |>
  distinct(speech_id, president)
training_presidents <- split_assignments |>
  filter(split == "training") |>
  inner_join(speech_people, by = join_by(speech_id)) |>
  pull(president)
test_speaker_overlap <- split_assignments |>
  filter(split == "test") |>
  inner_join(speech_people, by = join_by(speech_id)) |>
  mutate(
    president_seen_in_training = president %in% training_presidents
  )

paired_test <- test_comparison_by_speech |>
  summarise(
    test_speeches = n(),
    trigram_wins = sum(trigram < unigram),
    ties = sum(trigram == unigram),
    trigram_losses = sum(trigram > unigram)
  )

speaker_overlap_test <- test_speaker_overlap |>
  summarise(
    repeated_president_speeches = sum(president_seen_in_training),
    other_speeches = sum(!president_seen_in_training)
  )
non_overlap_test <- test_comparison_by_speech |>
  inner_join(
    test_speaker_overlap |>
      select(speech_id, president_seen_in_training),
    by = join_by(speech_id)
  ) |>
  filter(!president_seen_in_training) |>
  summarise(
    speeches = n(),
    trigram_wins = sum(trigram < unigram)
  )

Untouched-test result for the selected model and baseline
Model	Role	Test next-token rows	Targets mapped to `<unk>`	Test perplexity	Top-1 accuracy
uniform	uniform vocabulary floor	24337	11.7%	2500.0	not applicable
unigram	predeclared frequency baseline	24337	11.7%	240.7	11.7%
trigram	validation-selected model	24337	11.7%	145.3	19.2%

The uniform model has perplexity 2,500 because it gives every vocabulary item the same probability. The selected trigram reduces test perplexity by about 39.6% relative to the unigram baseline. It has lower perplexity in the paired comparison for all 13 test speeches.

Its highest-probability guess is right on 19.2 percent of test rows. The unigram scores 11.7 percent because its most likely token is <unk>, the category for many rare training words. A correct unknown-token prediction does not recover the original word. Roughly one third of the trigram’s correct guesses are also unknown-token matches rather than recovered words.

Presidents who also appear in training delivered eight of the 13 test speeches. The test separates speeches, not speakers. Repeated presidents can therefore contribute similar phrasing to different splits. The trigram still has lower perplexity than the unigram on all five test speeches whose president is absent from training.

Perplexity would change under another vocabulary, tokenizer, boundary rule, or test set. The values here cannot be compared with the LDA perplexity in Lesson 55 or a published model score from another tokenization.

The result does not show that the model understands a speech, writes a good report, or produces facts. The next lesson uses these same measurements to show how a report can retain its evidence references.

What to remember

Next-token prediction assigns probabilities after a context.
Paragraph boundaries prevent invented transitions between source records.
Validation data select context length and interpolation strength. Test data measure performance after that choice.
The study selects the trigram and interpolation strength together on validation data.
The selected trigram has lower perplexity than the unigram on all 13 test speeches here. Its top choice is still wrong on about four out of five rows.
Perplexity is comparable only under the same tokens, vocabulary, and test rows.
A likely next token is not necessarily factual or useful.