Visualization

Two views of the same decode: what the model looked at, and what the search considered and discarded.

Question	Functions
What did the decoder attend to?	plot_attention, format_attention
What did the search consider and prune?	plot_beam_search, format_beam_search

The second is the one students find least intuitive, because pruning is invisible in the output. A translation tells you what won and never what lost — yet the entire argument for beam search is about paths greedy never explores. See Beam search below.

Attention

Overview

Each row of an attention matrix is a probability distribution saying "while producing this target token, here is how much I looked at each source token." Rendering that turns an abstraction into something you can check against your own intuition about the sentence pair.

Two renderers are provided:

Function	Output	Needs
plot_attention	matplotlib heatmap	matplotlib
format_attention	shaded text grid	nothing extra

format_attention exists because a figure window is not always available — over SSH, in a log file, in CI, in a doctest.

Quick Start

from torchlingo.visualization import format_attention, plot_attention

logits, weights = model(src, tgt, return_attention=True)
src_tokens = src_vocab.indices_to_tokens(src_ids)
tgt_tokens = tgt_vocab.indices_to_tokens(tgt_ids)

print(format_attention(weights, src_tokens, tgt_tokens))
plot_attention(weights, src_tokens, tgt_tokens, title="Luong dot attention")

Both take a single sentence: either [tgt_len, src_len], or [1, tgt_len, src_len], which is unwrapped for you. A real batch raises, rather than silently showing you sentence zero — index it yourself with weights[i].

Works on both architectures

return_attention=True means the same thing on SimpleSeq2SeqLSTM and SimpleTransformer, and returns the same (batch, tgt_len, src_len) shape, so the code above does not care which model you hand it.

On the Transformer the weights come from the last decoder layer with heads averaged. That is the conventional choice for alignment plots. For every layer, or per-head maps, use capture_cross_attention directly.

Call eval() first

In training mode, attention dropout randomly zeroes weights and rescales the survivors, so rows will not sum to 1 and the picture is not the one inference uses. This is easy to miss because the plot still looks plausible. Check it:

assert torch.allclose(weights.sum(-1), torch.ones(weights.shape[:-1]), atol=1e-5)

Why the Transformer needs extra machinery and the LSTM does not

The LSTM computes attention in TorchLingo's own code, so returning the weights is a matter of not throwing them away.

The Transformer uses PyTorch's nn.Transformer, and TransformerDecoderLayer._mha_block calls attention with need_weights=False hardcoded. That is a deliberate performance choice: it allows a fused kernel that never materializes the attention matrix at all. The weights are not hidden behind an option, they are genuinely never computed, which is why a forward hook on multihead_attn comes back with None.

capture_cross_attention temporarily replaces each layer's multihead_attn.forward to force need_weights=True, records what comes back, and restores the original afterwards — including if the forward pass raises. The speed cost is why this is opt-in rather than always on.

This is worth knowing beyond TorchLingo. Fast paths that discard intermediate values are common in deep learning libraries, and "the framework will not give me X" often means "X is never computed on the path you are taking."

Attention for the translation the model actually produced

model(src, tgt, return_attention=True) tells you where attention went while processing a translation you supplied. Hand it the reference translation and you learn where the model would have looked had it produced the right answer — which is a different question from what it did.

The decoders answer the second one:

from torchlingo.inference import beam_search_decode, greedy_decode

tokens, weights = greedy_decode(model, src, return_attention=True)
tokens, weights = beam_search_decode(model, src, beam_size=5, return_attention=True)

Greedy decodes a batch, so it returns a list of (tgt_len, src_len) tensors, one per sentence — a list rather than a stacked tensor because decoded sequences differ in length, and padding them would invent attention rows that were never computed. Beam search takes one sentence and returns one tensor, for the winning hypothesis.

Why re-running is exact, and why beam search does not keep weights

Both decoders discard attention as they go. Rather than thread weights through the search, attention_for_sequence re-runs the finished sequence in one teacher-forced pass.

That is exact, not an approximation. The decoder is causally masked, so the state at target position t depends only on tokens up to t. Feeding the whole sequence at once reproduces each row exactly as the incremental decode computed it — asserted to floating-point noise in tests/test_attention_from_decoding.py::ExactnessTests.

For beam search there is a second reason. Weights belong to a hypothesis, and hypotheses get pruned, so most of what a beam search computes belongs to candidates that lost. Keeping all of it to discard all but one costs memory proportional to beam_size for no benefit. This mirrors what the search already does with scores: it re-scores prefixes rather than caching every partial result.

A model that got it right

EN      I want to help you.
BEAM-5  Quiero ayudar.

        <sos>    ▁I ▁want   ▁to ▁help  ▁you     . <eos>
▁Quiero   ···   ░░░   ▒▒▒   ░░░   ···   ···   ···   ···
▁ayud     ···   ···   ···   ···   ▒▒▒   ░░░   ···   ···
ar        ···   ···   ···   ░░░   ···   ░░░   ░░░   ···
.         ···   ···   ···   ···   ░░░   ░░░   ░░░   ···
<eos>     ▒▒▒   ···   ···   ···   ···   ···   ···   ▒▒▒

Read the two dark cells. ▁Quiero attends most to ▁want, and ▁ayud attends most to ▁help. The model has learned that Spanish fuses "I want" into a single inflected verb and that the alignment is not monotonic — Quiero covers source positions 1 and 2 at once.

Note also what it dropped. The translation omits "you" entirely, and the ▁you column is correspondingly faint everywhere. The map is not just showing you a correct alignment; it is showing you which source word the model never really used.

A model that got it wrong

Same checkpoint, a sentence it handles badly. This is the more common case, and the more instructive one:

EN  The cat sleeps on the mat.
ES  El código de la catura.

      <sos>  ▁The    ▁c    at    ▁s    le    ep     s   ▁on  ▁the  ▁mat     . <eos>
▁El     ···   ▓▓▓   ···   ···   ···   ···   ···   ···   ···   ···   ···   ···   ···
▁c      ···   ···   ░░░   ···   ░░░   ···   ···   ···   ···   ···   ···   ···   ···
ó       ···   ···   ░░░   ░░░   ░░░   ···   ···   ···   ···   ···   ░░░   ···   ···
d       ···   ···   ···   ···   ···   ···   ░░░   ░░░   ···   ···   ░░░   ···   ···
igo     ···   ···   ···   ···   ···   ░░░   ░░░   ░░░   ···   ···   ░░░   ···   ···
▁de     ···   ···   ···   ···   ···   ···   ···   ···   ░░░   ···   ···   ░░░   ···
▁la     ···   ···   ···   ···   ···   ···   ···   ···   ···   ░░░   ···   ░░░   ···
▁c      ···   ···   ░░░   ░░░   ░░░   ···   ···   ···   ···   ···   ···   ···   ···
at      ···   ···   ░░░   ░░░   ░░░   ···   ···   ···   ···   ···   ░░░   ···   ···
ura     ···   ···   ···   ···   ···   ░░░   ░░░   ░░░   ···   ···   ░░░   ···   ···

The translation is wrong, and the map shows how it is wrong, which a correct translation would not.

▁El attends sharply to ▁The — the one word it got right, and the one place the grid is dark. ▁c and at attend to the ▁c/at pieces of "cat", so the model half-recognized the word and produced catura. Everywhere else the row is a flat wash of ░░░: attention spread thinly across the whole source, which is what "has not learned what to look at" looks like.

Compare that to Tutorial 4's model, trained on a synthetic task with a known correct alignment, where the map is a clean diagonal. The contrast is the lesson: a sharp attention map is evidence the model learned something, and a diffuse one is evidence it did not.

Reading the text grid

Rows are target tokens, columns are source tokens, and shading runs · ░ ▒ ▓ █ from no attention to full attention. Here is a model trained on the reversal task from examples/attention_alignment.py, where the target is the source translated word-for-word and reversed:

       <sos>  bird wants     a   cat   old   dog <eos>
<sos>    ···   ···   ···   ···   ···   ···   ███   ···
perro    ···   ···   ···   ···   ···   ▓▓▓   ░░░   ···
viejo    ···   ···   ···   ···   ███   ···   ···   ···
gato     ···   ···   ···   ███   ···   ···   ···   ···
un       ···   ···   ▓▓▓   ···   ···   ···   ···   ···
quiere   ···   ███   ···   ···   ···   ···   ···   ···
pajaro   ▓▓▓   ░░░   ···   ···   ···   ···   ···   ···

The anti-diagonal is exactly right: to emit the first target word the decoder looks at the last source word. Each row is offset by one from the token it names, because row n is the state that predicts token n + 1.

Pass show_values=True for rounded percentages instead of shading when you need the actual numbers.

Beam search

Beam search keeps several hypotheses alive and discards the rest at every step. The output tells you which one won; it never tells you what was thrown away, and that is exactly where the interesting behaviour is.

Pass a list as trace and the search records every candidate it scored:

from torchlingo.inference import beam_search_decode
from torchlingo.visualization import format_beam_search

trace = []
tokens = beam_search_decode(model, src, beam_size=3, trace=trace)
print(format_beam_search(trace, itos=tgt_vocab.idx2token, winner=tokens))

Tracing is observation only. It never changes the result, and costs nothing when omitted.

Reading the output

step 0
  + -1.953  <s> w23
  > -2.007  <s> w19
  + -2.214  <s> w4
step 1
  + -3.245  <s> w19 w16
  > -3.383  <s> w19 w21
  + -3.501  <s> w4 w16
  . -3.543  <s> w23 w19
    ... 5 more considered

Marker	Meaning
>	kept, and on the path that eventually won
+	kept into the next step
.	pruned here

Look for a > sitting below a +. That is the whole lesson. At step 0 above, the eventual winner ranked second: greedy decoding would have committed to w23 and never recovered. Beam search found the better translation only because the beam was wide enough to carry a hypothesis that did not look best at the time.

If the > is always at the top, beam search did nothing greedy would not have done — and on an easy sentence that is the common case, which is worth seeing too.

plot_beam_search shows the same thing as scores over steps: kept candidates filled, pruned ones hollow, and the winning path drawn as a line. The visual question is whether that line ever dips below other filled points.

API Reference

visualization

Visualizing what a model attends to, and what a search considered.

Two things here, answering two different questions about the same decode:

• Attention shows what the decoder looked at while producing each token.
• Beam search shows what it considered and discarded on the way there.

The second is the one students find least intuitive, because pruning is invisible in the output: a translation tells you what won, never what lost, and the whole argument for beam search is about the paths greedy never explores.

Attention weights are the most directly inspectable quantity in an NMT model: each row is a probability distribution saying "while producing this target token, here is how much I looked at each source token." Plotting that matrix turns an abstraction into something a student can check against their own intuition about the sentence pair.

Two renderers are provided:

• :func:plot_attention draws the familiar alignment heatmap with matplotlib.
• :func:format_attention returns a shaded text grid instead. It needs nothing beyond the standard library, so it works over SSH, in a plain terminal, in log files, and in doctests -- anywhere a figure window is not available.

Typical usage

import torch weights = torch.tensor([[0.8, 0.1, 0.1], [0.1, 0.2, 0.7]]) print(format_attention(weights, ["el", "gato", "duerme"], ["the", "sleeps"])) el gato duerme the ▓▓▓ ··· ··· sleeps ··· ░░░ ▓▓▓

format_attention

format_attention(weights: Tensor, src_tokens: list[str], tgt_tokens: list[str], show_values: bool = False) -> str

Render an attention matrix as a shaded text grid.

Rows are target tokens, columns are source tokens, and each row sums to 1. Darker cells mean the decoder looked harder at that source token while producing that target token.

Parameters:

Name	Type	Description	Default
weights	Tensor	Attention weights of shape (tgt_len, src_len), or (1, tgt_len, src_len).	required
src_tokens	list[str]	Source token strings, one per source position.	required
tgt_tokens	list[str]	Target token strings, one per target position.	required
show_values	bool	Print two-digit percentages instead of shading. Useful when you need the exact numbers. Defaults to False.	False

Returns:

Name	Type	Description
str	str	A multi-line string suitable for printing.

Raises:

Type	Description
ValueError	If the weights are not a single 2-D matrix, or the label counts do not match its shape.

Examples:

>>> import torch
>>> w = torch.tensor([[0.9, 0.1]])
>>> print(format_attention(w, ["gato", "duerme"], ["cat"], show_values=True))
      gato duerme
cat     90     10

Source code in src/torchlingo/visualization.py

def format_attention(
    weights: torch.Tensor,
    src_tokens: list[str],
    tgt_tokens: list[str],
    show_values: bool = False,
) -> str:
    """Render an attention matrix as a shaded text grid.

    Rows are target tokens, columns are source tokens, and each row sums to 1.
    Darker cells mean the decoder looked harder at that source token while
    producing that target token.

    Args:
        weights (torch.Tensor): Attention weights of shape (tgt_len, src_len),
            or (1, tgt_len, src_len).
        src_tokens (list[str]): Source token strings, one per source position.
        tgt_tokens (list[str]): Target token strings, one per target position.
        show_values (bool, optional): Print two-digit percentages instead of
            shading. Useful when you need the exact numbers. Defaults to False.

    Returns:
        str: A multi-line string suitable for printing.

    Raises:
        ValueError: If the weights are not a single 2-D matrix, or the label
            counts do not match its shape.

    Examples:
        >>> import torch
        >>> w = torch.tensor([[0.9, 0.1]])
        >>> print(format_attention(w, ["gato", "duerme"], ["cat"], show_values=True))
              gato duerme
        cat     90     10
    """
    matrix = _as_matrix(weights)
    _check_labels(matrix, src_tokens, tgt_tokens)

    label_w = max((len(t) for t in tgt_tokens), default=0)
    cell_w = max(6, max((len(t) for t in src_tokens), default=0) + 1)

    # Header and cells are both right-justified in the same column width, so
    # each column of shading sits under its source token.
    header = " " * label_w + "".join(t.rjust(cell_w) for t in src_tokens)
    lines = [header]
    for row_label, row in zip(tgt_tokens, matrix.tolist()):
        if show_values:
            cells = "".join(f"{round(w * 100):d}".rjust(cell_w) for w in row)
        else:
            cells = "".join((_shade(w) * 3).rjust(cell_w) for w in row)
        lines.append(row_label.ljust(label_w) + cells)
    return "\n".join(lines)

plot_attention

plot_attention(weights: Tensor, src_tokens: list[str], tgt_tokens: list[str], title: str | None = None, cmap: str = 'viridis', ax: Axes | None = None) -> Axes

Draw an attention alignment heatmap with matplotlib.

Parameters:

Name	Type	Description	Default
weights	Tensor	Attention weights of shape (tgt_len, src_len), or (1, tgt_len, src_len).	required
src_tokens	list[str]	Source token strings, one per source position.	required
tgt_tokens	list[str]	Target token strings, one per target position.	required
title	str	Title for the plot.	None
cmap	str	Matplotlib colormap name. Defaults to "viridis".	'viridis'
ax	Axes	Existing axes to draw into. A new figure and axes are created when omitted.	None

Returns:

Type	Description
Axes	matplotlib.axes.Axes: The axes the heatmap was drawn into.

Raises:

Type	Description
ImportError	If matplotlib is unavailable. It ships as a TorchLingo dependency, so this normally cannot happen; the guard exists for stripped-down installs, and points at :func:format_attention.
ValueError	If the weights are not a single 2-D matrix, or the label counts do not match its shape.

Examples:

>>> import torch
>>> ax = plot_attention(torch.rand(3, 4).softmax(-1),
...                     ["a", "b", "c", "d"], ["x", "y", "z"])
...

Source code in src/torchlingo/visualization.py

def plot_attention(
    weights: torch.Tensor,
    src_tokens: list[str],
    tgt_tokens: list[str],
    title: str | None = None,
    cmap: str = "viridis",
    ax: "matplotlib.axes.Axes | None" = None,  # noqa: F821
) -> "matplotlib.axes.Axes":  # noqa: F821
    """Draw an attention alignment heatmap with matplotlib.

    Args:
        weights (torch.Tensor): Attention weights of shape (tgt_len, src_len),
            or (1, tgt_len, src_len).
        src_tokens (list[str]): Source token strings, one per source position.
        tgt_tokens (list[str]): Target token strings, one per target position.
        title (str, optional): Title for the plot.
        cmap (str, optional): Matplotlib colormap name. Defaults to "viridis".
        ax (matplotlib.axes.Axes, optional): Existing axes to draw into. A new
            figure and axes are created when omitted.

    Returns:
        matplotlib.axes.Axes: The axes the heatmap was drawn into.

    Raises:
        ImportError: If matplotlib is unavailable. It ships as a TorchLingo
            dependency, so this normally cannot happen; the guard exists for
            stripped-down installs, and points at :func:`format_attention`.
        ValueError: If the weights are not a single 2-D matrix, or the label
            counts do not match its shape.

    Examples:
        >>> import torch
        >>> ax = plot_attention(torch.rand(3, 4).softmax(-1),
        ...                     ["a", "b", "c", "d"], ["x", "y", "z"])
        ... # doctest: +SKIP
    """
    try:
        import matplotlib.pyplot as plt
    except ImportError as exc:
        raise ImportError(
            "plot_attention requires matplotlib. It normally ships with "
            "TorchLingo; reinstall with 'pip install matplotlib', or use "
            "format_attention for a text rendering that needs no extra packages."
        ) from exc

    matrix = _as_matrix(weights)
    _check_labels(matrix, src_tokens, tgt_tokens)

    if ax is None:
        _fig, ax = plt.subplots(
            figsize=(max(4, len(src_tokens) * 0.6), max(3, len(tgt_tokens) * 0.5))
        )

    image = ax.imshow(matrix.numpy(), aspect="auto", cmap=cmap, vmin=0.0, vmax=1.0)
    ax.set_xticks(range(len(src_tokens)), src_tokens, rotation=45, ha="right")
    ax.set_yticks(range(len(tgt_tokens)), tgt_tokens)
    ax.set_xlabel("source")
    ax.set_ylabel("target")
    if title:
        ax.set_title(title)
    ax.figure.colorbar(image, ax=ax, label="attention weight")
    ax.figure.tight_layout()
    return ax

format_beam_search

format_beam_search(trace: list[BeamStep], itos: list[str] | None = None, winner: list[int] | None = None, top: int = 6, max_steps: int | None = None) -> str

Render a beam search as text: what was considered, kept, and discarded.

Produced from the trace argument of :func:torchlingo.inference.beam_search_decode.

Each step lists candidates in the order the search ranked them. The marker in the first column is the point of the whole display:

========== ========================================================== > kept, and a prefix of the hypothesis that eventually won + kept into the next step . pruned here ========== ==========================================================

The steps worth looking at are the ones where a > sits below a +: the eventual winner ranked below another hypothesis at that moment and survived only because the beam was wide enough to carry it. That is exactly the situation greedy decoding cannot recover from, and it is otherwise invisible in the output.

Parameters:

Name	Type	Description	Default
trace	list[BeamStep]	Steps recorded during decoding.	required
itos	list[str]	Index-to-token mapping. Raw ids are shown when omitted.	None
winner	list[int]	The returned token sequence, used to mark which candidates were on the winning path.	None
top	int	Candidates to show per step. Defaults to 6.	6
max_steps	int	Steps to show. All of them when omitted.	None

Returns:

Name	Type	Description
str	str	A multi-line string suitable for printing.

Raises:

Type	Description
ValueError	If the trace is empty.

Examples:

>>> from torchlingo.inference import BeamCandidate, BeamStep
>>> step = BeamStep(0, [BeamCandidate([2, 7], -0.2, -0.2, True)])
>>> print(format_beam_search([step], itos=["<pad>", "<unk>", "<s>", "</s>", "", "", "", "hola"]))
step 0
  + -0.200  <s> hola

Source code in src/torchlingo/visualization.py

def format_beam_search(
    trace: list[BeamStep],
    itos: list[str] | None = None,
    winner: list[int] | None = None,
    top: int = 6,
    max_steps: int | None = None,
) -> str:
    """Render a beam search as text: what was considered, kept, and discarded.

    Produced from the ``trace`` argument of
    :func:`torchlingo.inference.beam_search_decode`.

    Each step lists candidates in the order the search ranked them. The marker
    in the first column is the point of the whole display:

    ==========  ==========================================================
    ``>``       kept, and a prefix of the hypothesis that eventually won
    ``+``       kept into the next step
    ``.``       pruned here
    ==========  ==========================================================

    The steps worth looking at are the ones where a ``>`` sits below a ``+``:
    the eventual winner ranked *below* another hypothesis at that moment and
    survived only because the beam was wide enough to carry it. That is exactly
    the situation greedy decoding cannot recover from, and it is otherwise
    invisible in the output.

    Args:
        trace (list[BeamStep]): Steps recorded during decoding.
        itos (list[str], optional): Index-to-token mapping. Raw ids are shown
            when omitted.
        winner (list[int], optional): The returned token sequence, used to mark
            which candidates were on the winning path.
        top (int, optional): Candidates to show per step. Defaults to 6.
        max_steps (int, optional): Steps to show. All of them when omitted.

    Returns:
        str: A multi-line string suitable for printing.

    Raises:
        ValueError: If the trace is empty.

    Examples:
        >>> from torchlingo.inference import BeamCandidate, BeamStep
        >>> step = BeamStep(0, [BeamCandidate([2, 7], -0.2, -0.2, True)])
        >>> print(format_beam_search([step], itos=["<pad>", "<unk>", "<s>", "</s>", "", "", "", "hola"]))
        step 0
          + -0.200  <s> hola
    """
    if not trace:
        raise ValueError("Empty trace: pass trace=[] to beam_search_decode first.")

    steps = trace if max_steps is None else trace[:max_steps]
    winning_prefixes = set()
    if winner is not None:
        winning_prefixes = {tuple(winner[: i + 1]) for i in range(len(winner))}

    lines: list[str] = []
    for record in steps:
        lines.append(f"step {record.step}")
        for candidate in record.candidates[:top]:
            if candidate.kept and tuple(candidate.tokens) in winning_prefixes:
                marker = ">"
            elif candidate.kept:
                marker = "+"
            else:
                marker = "."
            lines.append(
                f"  {marker} {candidate.normalized:6.3f}  "
                f"{_decode_tokens(candidate.tokens, itos)}"
            )
        pruned = len(record.candidates) - top
        if pruned > 0:
            lines.append(f"    ... {pruned} more considered")
    return "\n".join(lines)

plot_beam_search

plot_beam_search(trace: list[BeamStep], winner: list[int] | None = None, top: int = 6, title: str | None = None, ax: Axes | None = None) -> Axes

Plot candidate scores per step, separating survivors from pruned paths.

Kept candidates are drawn filled, pruned ones hollow, and the winning path is connected by a line. The visual question the plot answers is whether the winning line ever dips below other kept points — which is precisely when beam search earns its cost over greedy.

Parameters:

Name	Type	Description	Default
trace	list[BeamStep]	Steps recorded during decoding.	required
winner	list[int]	The returned token sequence.	None
top	int	Candidates to plot per step.	6
title	str	Title for the plot.	None
ax	Axes	Existing axes to draw into.	None

Returns:

Type	Description
Axes	matplotlib.axes.Axes: The axes drawn into.

Raises:

Type	Description
ImportError	If matplotlib is unavailable.
ValueError	If the trace is empty.

Source code in src/torchlingo/visualization.py

def plot_beam_search(
    trace: list[BeamStep],
    winner: list[int] | None = None,
    top: int = 6,
    title: str | None = None,
    ax: "matplotlib.axes.Axes | None" = None,  # noqa: F821
) -> "matplotlib.axes.Axes":  # noqa: F821
    """Plot candidate scores per step, separating survivors from pruned paths.

    Kept candidates are drawn filled, pruned ones hollow, and the winning path
    is connected by a line. The visual question the plot answers is whether the
    winning line ever dips below other kept points — which is precisely when
    beam search earns its cost over greedy.

    Args:
        trace (list[BeamStep]): Steps recorded during decoding.
        winner (list[int], optional): The returned token sequence.
        top (int, optional): Candidates to plot per step.
        title (str, optional): Title for the plot.
        ax (matplotlib.axes.Axes, optional): Existing axes to draw into.

    Returns:
        matplotlib.axes.Axes: The axes drawn into.

    Raises:
        ImportError: If matplotlib is unavailable.
        ValueError: If the trace is empty.
    """
    try:
        import matplotlib.pyplot as plt
    except ImportError as exc:
        raise ImportError(
            "plot_beam_search requires matplotlib. Use format_beam_search for a "
            "text rendering that needs no extra packages."
        ) from exc

    if not trace:
        raise ValueError("Empty trace: pass trace=[] to beam_search_decode first.")

    if ax is None:
        _fig, ax = plt.subplots(figsize=(max(5, len(trace) * 0.7), 4))

    winning_prefixes = set()
    if winner is not None:
        winning_prefixes = {tuple(winner[: i + 1]) for i in range(len(winner))}

    winning_x: list[int] = []
    winning_y: list[float] = []
    for record in trace:
        for candidate in record.candidates[:top]:
            on_winning_path = tuple(candidate.tokens) in winning_prefixes
            ax.scatter(
                record.step,
                candidate.normalized,
                facecolors="tab:blue" if candidate.kept else "none",
                edgecolors="tab:blue" if candidate.kept else "tab:grey",
                zorder=3 if candidate.kept else 2,
            )
            if on_winning_path:
                winning_x.append(record.step)
                winning_y.append(candidate.normalized)

    if winning_x:
        ax.plot(
            winning_x, winning_y, color="tab:red", linewidth=2, zorder=4, label="winner"
        )
        ax.legend(loc="best")

    ax.set_xlabel("step")
    ax.set_ylabel("length-normalized score")
    if title:
        ax.set_title(title)
    ax.figure.tight_layout()
    return ax