Connectionist temporal classification (CTC) loss

Deep dive

How it works

In: log-probabilities lp[t][c] for T frames and V symbols (blank = 0) and a target y of L labels · Out: −log p(y|x) and its gradient w.r.t. the logits · State: none (a T×(2L+1) table, only while computing) · Cost: O(T·(2L+1)) log-add-exps; 3.2-4.1 ms for T=1,000, L=100 (measured, below)

1. Collapse rule. A path π (one symbol per frame) maps to B(π): merge consecutive repeats, then delete blanks. B(a a ∅ a b b) = a a b. The blank lets the net say "no new label" and separates genuine repeats (l ∅ l).
2. Extended label. Insert blanks around every label: y = (a, b) → l' = (∅, a, ∅, b, ∅), S = 2L+1 states. Every path that collapses to y is a monotone walk through l'.
3. Transitions. From frame t−1 to t, state s is reached from s (stay), s−1 (advance), or s−2 (skip a blank) only if l'_s is not blank and l'_s ≠ l'_{s−2}. Equal neighbours may not skip, or l l would collapse to l.
4. Forward. α_t(s) = (α_{t−1}(s) + α_{t−1}(s−1) + [skip] α_{t−1}(s−2)) · p_t(l'_s), starting at s ∈ {0, 1} for t = 0. In log space the sum is logaddexp, with −∞ = probability 0.
5. Loss. p(y|x) = α_{T−1}(S−1) + α_{T−1}(S−2) (end on the last label or the final blank); loss = −log p(y|x).
6. Backward. β runs the same recursion from the end. It equals the forward table of the reversed input and reversed target, which the code uses.
7. Gradient w.r.t. logit u_t^k (softmax output y_t^k): y_t^k − (1/(p·y_t^k)) · Σ_{s: l'_s = k} α_t(s) β_t(s) (α and β both include the emission at t; Graves et al. 2006). Autodiff through the forward pass gives the same numbers.
8. Decode. Greedy best path: argmax per frame, then collapse. Prefix beam search keeps per label prefix p_blank and p_nonblank (paths ending in blank / in the last label), summing paths per prefix: it finds the best labelling, not the best path. An external LM adds α·log p_LM + β·length to each prefix score.

Code

const NEG: f64 = f64::NEG_INFINITY; // log 0
fn lse(a: f64, b: f64) -> f64 { // log(e^a + e^b), safe with -inf
    if a == NEG { return b; }
    if b == NEG { return a; }
    let m = a.max(b);
    m + ((a - m).exp() + (b - m).exp()).ln()
}
fn log_softmax(u: &[f64]) -> Vec<f64> {
    let z = u.iter().fold(NEG, |a, &v| lse(a, v));
    u.iter().map(|&v| v - z).collect()
}
/// alpha[t][s] over (blank, y0, blank, y1, ..., blank); includes the emission at frame t.
fn alphas(lp: &[Vec<f64>], y: &[usize]) -> Vec<Vec<f64>> {
    let n = 2 * y.len() + 1;
    let lab = |s: usize| if s % 2 == 0 { 0 } else { y[s / 2] };
    let mut a = vec![vec![NEG; n]; lp.len()];
    a[0][0] = lp[0][0];
    if n > 1 { a[0][1] = lp[0][y[0]]; }
    for t in 1..lp.len() {
        for s in 0..n {
            let mut v = a[t - 1][s];
            if s >= 1 { v = lse(v, a[t - 1][s - 1]); }
            if s >= 2 && lab(s) != 0 && lab(s) != lab(s - 2) { v = lse(v, a[t - 1][s - 2]); }
            a[t][s] = if v == NEG { NEG } else { v + lp[t][lab(s)] };
        }
    }
    a
}
fn log_p(a: &[Vec<f64>]) -> f64 { // end on the last label or the final blank
    let (last, n) = (&a[a.len() - 1], a[0].len());
    if n > 1 { lse(last[n - 1], last[n - 2]) } else { last[0] }
}
/// -log p(y|x); +inf when T < L + repeats.
pub fn ctc_loss(lp: &[Vec<f64>], y: &[usize]) -> f64 { -log_p(&alphas(lp, y)) }

/// Loss and dLoss/dlogits. Beta = alpha of the reversed input and reversed target.
pub fn ctc_loss_grad(u: &[Vec<f64>], y: &[usize]) -> (f64, Vec<Vec<f64>>) {
    let lp: Vec<Vec<f64>> = u.iter().map(|r| log_softmax(r)).collect();
    let a = alphas(&lp, y);
    let lpp = log_p(&a);
    let rl: Vec<Vec<f64>> = lp.iter().rev().cloned().collect();
    let b = alphas(&rl, &y.iter().rev().cloned().collect::<Vec<_>>());
    let (tl, n, v) = (lp.len(), a[0].len(), lp[0].len());
    let mut g = vec![vec![0.0; v]; tl];
    for t in 0..tl {
        let mut occ = vec![NEG; v]; // log sum of alpha*beta over states holding each symbol
        for s in 0..n {
            let c = if s % 2 == 0 { 0 } else { y[s / 2] };
            occ[c] = lse(occ[c], a[t][s] + b[tl - 1 - t][n - 1 - s]);
        }
        for c in 0..v { g[t][c] = lp[t][c].exp() - (occ[c] - lp[t][c] - lpp).exp(); }
    }
    (-lpp, g)
}
/// Best path: argmax per frame, merge repeats, drop blanks.
pub fn greedy(lp: &[Vec<f64>]) -> Vec<usize> {
    let (mut out, mut prev) = (Vec::new(), 0);
    for r in lp {
        let k = (0..r.len()).max_by(|&i, &j| r[i].partial_cmp(&r[j]).unwrap()).unwrap();
        if k != 0 && k != prev { out.push(k); }
        prev = k;
    }
    out
}
fn collapse(p: &[usize]) -> Vec<usize> {
    (0..p.len()).filter(|&i| p[i] != 0 && (i == 0 || p[i - 1] != p[i])).map(|i| p[i]).collect()
}
fn brute(lp: &[Vec<f64>], y: &[usize]) -> f64 { // sum over all V^T paths
    let (tl, v) = (lp.len(), lp[0].len());
    (0..v.pow(tl as u32)).map(|code| {
        let path: Vec<usize> = (0..tl).map(|t| code / v.pow(t as u32) % v).collect();
        if collapse(&path) == y { (0..tl).map(|t| lp[t][path[t]].exp()).product() } else { 0.0 }
    }).sum()
}
fn main() {
    let mut seed = 12345u64; // LCG, no rand crate
    let mut rnd = move || { seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407); (seed >> 11) as f64 / (1u64 << 53) as f64 };
    let mut mk = |t: usize, v: usize| -> Vec<Vec<f64>> { (0..t).map(|_| (0..v).map(|_| 3.0 * rnd()).collect()).collect() };
    let u = mk(6, 3);
    let lp: Vec<Vec<f64>> = u.iter().map(|r| log_softmax(r)).collect();
    for y in [vec![1, 2], vec![1, 1], vec![2, 1, 2], vec![]] {
        let (f, b) = (-ctc_loss(&lp, &y), brute(&lp, &y).ln());
        assert!((f - b).abs() < 1e-9, "{y:?}: {f} vs {b}");
    }
    assert_eq!(ctc_loss(&lp[..2], &[1, 1]), f64::INFINITY); // needs T >= 3
    assert_eq!(ctc_loss(&lp[..3], &[1, 2, 1, 2]), f64::INFINITY);
    assert!(ctc_loss(&lp[..3], &[1, 1]).is_finite());
    let (y, (_, g)) = ([1, 2, 1], ctc_loss_grad(&u, &[1, 2, 1]));
    for (t, c) in [(0, 0), (2, 1), (5, 2)] { // central finite differences
        let f = |d: f64| { let mut v = u.clone(); v[t][c] += d; ctc_loss_grad(&v, &y).0 };
        assert!(((f(1e-6) - f(-1e-6)) / 2e-6 - g[t][c]).abs() < 1e-6);
    }
    println!("greedy {:?}", greedy(&lp));
    let lp: Vec<Vec<f64>> = mk(1000, 30).iter().map(|r| log_softmax(r)).collect();
    let y: Vec<usize> = (0..100).map(|i| 1 + (i * 7) % 29).collect();
    let best = (0..50).map(|_| { let t0 = std::time::Instant::now(); ctc_loss(&lp, &y); t0.elapsed().as_secs_f64() }).fold(f64::MAX, f64::min);
    println!("T=1000 L=100 V=30 forward, best of 50: {:.0} us", best * 1e6);
}

Complexity and constants

• Time O(T·S), S = 2L+1; each cell is 1-2 logaddexp (two exp, one ln each). A batch is O(N·T·L), sequences independent.
• Memory: α and β are T×S floats each; the loss alone needs two rows.
• States with too few frames left (or used) to finish are unreachable: a diagonal band of cells can be skipped (order of 2× fewer when T ≈ L).
• Measured (Apple M4 Pro, one thread, --release, f64, Vec<Vec<f64>>, best of 50; the Mac was shared, spread ±20 %): forward 3.2-4.1 ms for T=1,000, L=100, V=30 (201,000 cells, 16-20 ns/cell); loss + gradient 9.5-12 ms.

Efficiency

• Cost: 201,000 cells for T=1,000, L=100; 3.2-4.1 ms forward in log space, 9.5-12 ms with gradient (measured above). Compute-bound on exp/ln, not memory: the table is 1.6 MB.
• Levers:
  • Linear-space recursion with per-frame rescaling — 0.17-0.18 ms vs 3.7-3.9 ms, about 20× (measured in a scratch variant on the same input, same loss to 3 decimals) — when frame probabilities don't underflow within one frame.
  • One max and one ln per cell, two rolling rows — 15-30 % faster (measured) — when staying in log space.
  • Skip unreachable states — up to about 2× fewer cells when T ≈ L (order of) — short-input regimes.
  • GPU CTC — Deep Speech 2 lists a fast GPU CTC implementation among its scaling optimizations (Amodei et al. 2015) — large batches.
  • Prune the beam, cache LM states — decode cost grows with beam width (order of linear) — decoding with an LM.
• Naive → standard → state of the art:
  • Enumerate all V^T paths · exponential (3^6 = 729 in the test, 30^1000 for real input).
  • Log-space forward-backward · O(T·L), 16-20 ns/cell here.
  • Rescaled linear space or fused GPU kernels · 0.9 ns/cell measured for the rescaled CPU variant.

Implementation notes

• Feed log_softmax output: PyTorch's CTCLoss takes log-probabilities (T, N, C), blank index 0 by default, and targets may not contain the blank. Pass true input_lengths and target_lengths for padded batches.
• PyTorch defaults: reduction='mean' divides each loss by its target length, then averages over the batch; zero_infinity=False leaves inf (and NaN gradients) when inputs are too short, True zeroes them.
• The frame rate must exceed the label rate: an 8× strided encoder can make T < L + r for fast speech or long words. Choose output units (characters, BPE, phones) and stride together.

Pitfalls

• Infinite loss: one too-short sample poisons the batch. Filter by T ≥ L + r or use zero_infinity.
• Greedy is not the most likely text: many paths share one labelling, so the argmax path can lose to a labelling whose probability is spread over many paths. Prefix beam search sums them.
• Peaky outputs: spikes are not boundaries, timestamps are approximate; use forced alignment for boundaries.
• No language model: pure CTC output has phonetic spelling errors. Deep Speech 2 dev word error rate for its 9-layer English model fell from 14.93 to 9.52 with an LM in the beam search (Amodei et al. 2015, Table 6).

Variants and alternatives

In production

• Speech recognition: Deep Speech 2 trains CTC networks and decodes with an LM in beam search (Amodei et al. 2015).
• Scene text and line recognition: CRNN uses a CTC-defined label probability on top of convolutional and recurrent features (Shi et al. 2015); see OCR text recognition.
• Nanopore sequencing: Oxford Nanopore's Bonito basecaller lists the Distill CTC guide in its README references.
• Wake-word and keyword detection (wake word); Distill names speech and handwriting recognition as CTC's main uses.

Implementations

Sources

• Graves, Fernández, Gomez, Schmidhuber, Connectionist Temporal Classification: Labelling Unsegmented Sequence Data with Recurrent Neural Networks, ICML 2006, pp. 369-376. paper
• Hannun, Sequence Modeling With CTC, Distill 2017. article
• Graves, Sequence Transduction with Recurrent Neural Networks, 2012. paper
• Amodei et al., Deep Speech 2: End-to-End Speech Recognition in English and Mandarin, 2015. paper
• Shi, Bai, Yao, An End-to-End Trainable Neural Network for Image-based Sequence Recognition and Its Application to Scene Text Recognition, 2015. paper
• PyTorch, torch.nn.CTCLoss. docs