beamgrad

Differentiable beam search for PyTorch. Exact beam search forward, surrogate gradients backward, native on CPU and CUDA.

Beam search is how sequence models decode, but it is discrete: top-k selection has no useful gradient. So models are usually trained with teacher forcing, then decoded with a search they never saw during training. beamgrad puts the search inside the computation graph. It runs exact, deterministic beam search, then backpropagates through the hypotheses it selected, so you can write losses on what the decoder actually produces.

import torch
import beamgrad

# An autoregressive model: next-token log-probabilities for each beam's prefix.
def step(beams):                                                # beams.sequences: [B, K, t] tokens so far
    return model(src, beams.sequences).log_softmax(-1)         # [B, K, V]

options = beamgrad.BeamOptions(beam_size=4, eos_token=EOS)
result = beamgrad.beam_search(step, options, max_steps=T, batch_size=B)
# result.scores: [B, K], best first, differentiable w.r.t. the model
# result.sequences: [B, K, T] tokens of each beam, -1 after it ends

# Structured margin: the reference must beat the best beam that is not the reference.
gold_lp = model.token_log_probs(src, gold)                      # [B, T] teacher-forced; gold is [B, T], -1-padded
gold_score = beamgrad.sequence_scores(gold_lp, gold_lengths, options)  # on the beams' scale
loss = beamgrad.losses.structured_margin(result, gold, gold_score, margin=1.0)
loss.backward()                                                 # through the search, into the model

The loss is defined on what beam search actually returns. When the reference already wins by the margin it is zero; otherwise it raises the reference and lowers the beam that beat it. examples/train_lm.py is a runnable version, with a GRU whose hidden states follow the beams (beams.parents reorders them, as it would a key/value cache). beamgrad.losses.minimum_risk (expected cost over the beams) and the estimators in beamgrad.estimators are the alternatives. In a controlled translation experiment, minimum-risk training improved test BLEU over continued MLE on every seed, and this margin against the reference did not (training guide).

If the next-token distributions of every beam are already in a [B, T, K, V] tensor, score and decode it directly:

scores = beamgrad.final_scores(log_probs, options)             # [B, K], differentiable
best = beamgrad.backtrack(beamgrad.decode(log_probs, options))[:, 0]   # [B, T] best sequence

Features

• Exact beam search. GNMT length penalty, EOS handling (finished beams are carried forward and keep competing), minimum length, variable-length batches, banned tokens, n-gram blocking and a repetition penalty. A strict total order on candidates makes results deterministic.
• Gradients through the search. Each final score is differentiated along the path that produced it (see how it works). The gradients match the C reference bit for bit and agree with finite differences.
• Trains in bounded memory. The backward pass hands each step only its own path gradient, so no dense [B, T, K, V] gradient is ever built. With rescore_fn, the search runs with inference memory and the gradient comes from one teacher-forced pass, which can use activation checkpointing. A Qwen2.5-0.5B training step at 8 beams × 64 steps × batch 8 drops from out of memory on 16 GB to 1.3 GiB, with the same gradient (docs/training.md).
• Losses and estimators. beamgrad.losses has a structured margin and minimum-risk training. beamgrad.estimators has smoother surrogates: softmax over the selected beams, and a relaxed top-k that also sends gradient to candidates the search pruned.
• Drives real models. beam_search runs an autoregressive model inside the search, one step at a time, reordering its cache by each beam's parent. On Qwen2.5-0.5B and Qwen3-0.6B it returns the same beams, with bit-identical scores, as transformers' generate(num_beams=...) in float32, at the same speed (benchmarks/hf_beam_search.py), and fine-tunes the model through the search.
• Native everywhere. CPU kernels are multi-threaded across the batch and pick AVX-512, AVX2, SSE4.2 or NEON at runtime. A CUDA engine runs forward and backward on the GPU for beams up to 1024, on PyTorch's stream and allocator. Every backend selects the same beams with the same scores, bit for bit.
• A good PyTorch citizen. The operators are registered with torch.library, with fake-tensor, autograd and vmap rules: torch.compile (even fullgraph=True), torch.export, torch.vmap and torch.func (grad, vjp, jacrev, per-example gradients) work.
• A stable C ABI. libdbs works from C, C++ or any FFI. It adds forced tokens and token-filter callbacks, fp16/bf16 input, incremental model-callback decoding (with each beam's parent, for KV-cache reordering), and two extra smooth surrogates: selected-beam softmax weights and a relaxed top-k pool.
• Also in JAX, through a custom VJP (beamgrad.jax.final_scores), with jit, grad and vmap.

Installation

Releases include prebuilt wheels for PyTorch 2.13 and 2.14: Linux (CPU, CUDA 12.6, CUDA 13.0), macOS arm64 and Windows. Each wheel works on every Python from 3.10. docs/installation.md has the matrix and the command for your PyTorch. Otherwise beamgrad compiles against your installed PyTorch:

pip install torch
pip install --no-build-isolation "git+https://github.com/maged15/beamgrad"

--no-build-isolation matters: the compiled operators only work with the PyTorch they were built against, and if the two differ, import beamgrad says so and gives the fix. If a CUDA toolkit (nvcc) is available, the CUDA operators are built automatically; BEAMGRAD_CUDA=1 makes them required and BEAMGRAD_CUDA=0 skips them. beamgrad.cuda_available() reports what you got. For the C library alone, use CMake (see the C API).

How it works

At each step, every live beam proposes every token. Its cumulative log-probability is ranked by raw / ((5 + length) / 6) ** alpha, and the K best candidates survive. Beams that emitted EOS are carried forward unchanged. Backward holds this selection fixed. The gradient of final beam k's score with respect to log_probs[t, p, v] is 1 / penalty(length_k) for every (t, p, v) on its path, and zero elsewhere. That is the exact derivative wherever a small perturbation would not change the selection. docs/algorithm.md gives the full definitions, including the additional C-level surrogates.

Performance

python benchmarks/benchmark.py times the PyTorch API against a beam search written with torch.topk, and checks that the scores agree. CPU results on a 4-core Intel Xeon (2.8 GHz, AVX-512) container, median milliseconds:

Backward time is dominated by writing the dense [B, T, K, V] gradient that autograd expects, so it is memory-bound. Run the script on your own hardware, including GPUs, before relying on these numbers; docs/cuda.md describes the CUDA engine.

C and C++

#include "dbs.h"

DBSOptionsC opt = {0};              /* zero fields select defaults */
opt.beam_size = 4;
opt.eos_token = -1;                 /* no EOS */

DBSDecoderHandle* decoder = NULL;
dbs_create_ex(opt, &decoder);

DBSResultHandle* result = NULL;
dbs_decode(decoder, log_probs, T, V, &result);           /* log_probs: [T, K, V] */
const float* scores = dbs_result_final_scores(result);   /* [K] */

float grad_final[4] = {1, 0, 0, 0};
DBSBackwardHandle* grad = NULL;
dbs_backward(decoder, result, NULL, NULL, grad_final, &grad);  /* sparse d scores[0] / d log_probs */

dbs_free_backward(grad);
dbs_free_result(result);
dbs_destroy(decoder);

find_package(beamgrad 2 REQUIRED)
target_link_libraries(app PRIVATE beamgrad::dbs)   # or beamgrad::dbs_cuda

The complete version, with error handling, is examples/c_api.c. The test suite builds and runs it.

Documentation

Scope and limitations

• Gradients are surrogate gradients. They are exact for a fixed beam selection and do not model how the selection itself would change.
• Through the steps (the default), the model's graph for every step is kept until the backward pass, as with any backpropagation through generation. rescore_fn avoids this at the cost of one teacher-forced pass, and the estimators, which need the rows, are only available through the steps.
• CUDA supports beams up to 1024 and about 268M candidates (K × V) per step.

Contributing

Issues and pull requests are welcome; see CONTRIBUTING.md. make test and make python-test run the suites locally.

Citing

If beamgrad helps your research, please cite it; CITATION.cff has the details.

License

MIT © Maged Amr