RNN transducer (RNN-T)

Deep dive

How it works

Parts:

• Encoder — maps T input frames to states f_t; causal or limited look-ahead (chunked attention).
• Prediction network — maps previous non-blank labels to g_u; a label-side language model (recurrent or short-context).
• Joint network — combines f_t and g_u into logits over V+1 classes per cell (t,u); blank = 0.
• Lattice loss — T×(U+1) grid; −log of the summed path probability.
• Decoder — greedy or beam search over the same grid.

1. Run the encoder, and the prediction network on [start, y_1..y_U] (teacher forcing, all g_u in parallel).
2. Run the joint on every (t,u), log_softmax. Two values per cell matter: b(t,u) = log P(blank), e(t,u) = log P(y_{u+1}).
3. A path runs (0,0) → (T−1,U): blank moves right, a label moves up, the last step is a blank. There are C(T−1+U, U) paths.
4. Forward: α(0,0)=0, α(t,u) = logaddexp(α(t−1,u)+b(t−1,u), α(t,u−1)+e(t,u−1)); log p(y|x) = α(T−1,U)+b(T−1,U); loss = −log p.
5. Backward: β(T−1,U)=b(T−1,U), β(t,u) = logaddexp(β(t+1,u)+b(t,u), β(t,u+1)+e(t,u)); β(0,0) equals the forward total. α + transition + β gives transition posteriors, hence the gradient.
6. Greedy decode: blank moves to frame t+1; a label is appended, the prediction network advances, the frame stays. Cap labels per frame.
7. Beam decode: hypotheses keyed by label prefix; equal prefixes merge by summing probabilities.

       t=0    t=1    t=2    t=3
u=2    o ───► o ───► o ───► o ──► final blank
       ▲      ▲      ▲      ▲
u=1    o ───► o ───► o ───► o       ───► blank: next frame (b)
       ▲      ▲      ▲      ▲       ▲    label: same frame (e)
u=0    o ───► o ───► o ───► o

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()
}
/// lp[t][u][k] = log P(k | t,u), blank = 0. The lattice needs only two values per cell:
/// blank[t*(U+1)+u] and emit[t*(U+1)+u] = log P(y[u] | t,u).
pub fn gather(lp: &[Vec<Vec<f64>>], y: &[usize]) -> (Vec<f64>, Vec<f64>) {
    let b = lp.iter().flatten().map(|c| c[0]).collect();
    let e = lp.iter().flat_map(|r| r.iter().enumerate().map(|(u, c)| if u < y.len() { c[y[u]] } else { NEG })).collect();
    (b, e)
}
/// log p(y|x), forward recursion with one rolling row over u: row[u] = alpha(t,u).
pub fn forward(blank: &[f64], emit: &[f64], tl: usize, ul: usize) -> f64 {
    let (w, mut row) = (ul + 1, vec![NEG; ul + 1]);
    for t in 0..tl {
        for u in 0..w {
            let down = if t > 0 { row[u] + blank[(t - 1) * w + u] } else if u == 0 { 0.0 } else { NEG };
            let right = if u > 0 { row[u - 1] + emit[t * w + u - 1] } else { NEG };
            row[u] = lse(down, right);
        }
    }
    row[ul] + blank[(tl - 1) * w + ul]
}
/// Same total from the far end: row[u] = beta(t,u), the log prob of finishing from (t,u).
pub fn backward(blank: &[f64], emit: &[f64], tl: usize, ul: usize) -> f64 {
    let (w, mut row) = (ul + 1, vec![NEG; ul + 1]);
    for t in (0..tl).rev() {
        for u in (0..w).rev() {
            let down = if t + 1 < tl { row[u] + blank[t * w + u] } else if u == ul { blank[t * w + u] } else { NEG };
            let right = if u < ul { row[u + 1] + emit[t * w + u] } else { NEG };
            row[u] = lse(down, right);
        }
    }
    row[0]
}
/// Brute force: sum over every monotone path (blank = right, label = up).
pub fn brute(b: &[f64], e: &[f64], tl: usize, ul: usize, t: usize, u: usize) -> f64 {
    let w = ul + 1;
    if t == tl - 1 && u == ul { return b[t * w + u].exp(); }
    let right = if t + 1 < tl { b[t * w + u].exp() * brute(b, e, tl, ul, t + 1, u) } else { 0.0 };
    let up = if u < ul { e[t * w + u].exp() * brute(b, e, tl, ul, t, u + 1) } else { 0.0 };
    right + up
}
/// Greedy search; joint(t, hyp) = log-probs over blank + labels given the labels so far.
pub fn greedy(tl: usize, max_sym: usize, joint: impl Fn(usize, &[usize]) -> Vec<f64>) -> Vec<usize> {
    let mut hyp = Vec::new();
    for t in 0..tl {
        for _ in 0..max_sym { // cap labels per frame, then move on
            let p = joint(t, &hyp);
            let k = (0..p.len()).max_by(|&i, &j| p[i].partial_cmp(&p[j]).unwrap()).unwrap();
            if k == 0 { break; } // blank: next frame
            hyp.push(k);
        }
    }
    hyp
}
fn main() {
    let y = [2usize, 1, 3];
    let w = |t: usize, u: usize, k: usize| 1.0 + ((t * 31 + u * 17 + k * 7) % 9) as f64; // deterministic weights
    let lp: Vec<Vec<Vec<f64>>> = (0..4).map(|t| (0..=3).map(|u| {
        let z: f64 = (0..4).map(|k| w(t, u, k)).sum();
        (0..4).map(|k| (w(t, u, k) / z).ln()).collect()
    }).collect()).collect();
    let (b, e) = gather(&lp, &y);
    let (f, bw, bf) = (forward(&b, &e, 4, 3), backward(&b, &e, 4, 3), brute(&b, &e, 4, 3, 0, 0).ln());
    assert!((f - bf).abs() < 1e-9 && (f - bw).abs() < 1e-9);
    // Peaked joint: label u at frame 2u, blank elsewhere.
    let peak = |t: usize, hyp: &[usize]| {
        let u = hyp.len();
        let want = if u < 3 && t == 2 * u { y[u] } else { 0 };
        (0..4).map(|c| if c == want { -0.06 } else { -3.9 }).collect()
    };
    assert_eq!(greedy(8, 2, peak), y);
    let (tl, ul) = (500, 100); // speed: only the two gathered values per cell
    let n = tl * (ul + 1);
    let b: Vec<f64> = (0..n).map(|i| -(1.0 + (i * 37 % 101) as f64) / 40.0).collect();
    let e: Vec<f64> = b.iter().rev().cloned().collect();
    let best = (0..50).map(|_| { let t0 = std::time::Instant::now(); std::hint::black_box(forward(&b, &e, tl, ul)); t0.elapsed().as_secs_f64() }).fold(f64::MAX, f64::min);
    println!("forward {f:.9} backward {bw:.9} brute {bf:.9}; T={tl} U={ul}: {n} cells, {:.0} us", best * 1e6);
}

Prints forward -8.973647397 backward -8.973647397 brute -8.973647397 (T=4, U=3: 20 paths).

Complexity and constants

• Lattice O(T·U), one logaddexp per cell; the joint adds a D→V projection and softmax (2·D·V FLOPs) per cell.
• Memory (arithmetic): B=32, T=250, U=100, V=500, f32 → 32·250·101·500·4 B ≈ 1.6 GB of logits, the gradient the same again.
• Standard RNN-T allows U > T; the one-label-per-frame variant needs T ≥ U.
• Measured (Apple M4 Pro, one thread, --release, f64, best of 50, T=500, U=100, 50,500 cells): forward 1.2-1.4 ms over two input sets (24-28 ns/cell), backward about 1.3 ms. The loop is serial along u.

Efficiency

• Cost: 50,500 cells in 1.2-1.4 ms on one M4 Pro core (measured). The lattice is cheap; the joint logits saturate GPU memory.
• Levers:
  • Prune the lattice — joint on s_range labels per frame, not U+1, bounds from a cheap linear joiner (Kuang et al. 2022; k2 rnnt_loss_pruned) — large V, long targets.
  • Gather and fuse — keep blank and target logits; fuse log_softmax into the loss (torchaudio fused_log_softmax=True) or run joint plus loss on sub-batches (NeMo fused_batch_size) — memory-bound training.
  • Linear-space rescaled recursion — 0.27 ms vs 1.4 ms, about 5× (measured in a scratch variant, totals agree to about 1e-12) — when per-row rescaling avoids underflow.
  • Short-context prediction network — outputs depend on the last k labels, so they can be cached (HAT, Variani et al. 2020) — decoding.
  • Emission regulariser — FastEmit cut 90th-percentile latency from 210 ms to 30 ms on LibriSpeech (Yu et al. 2020) — late emissions.
• Naive → standard → state of the art:
  • Enumerate all paths · 20 in the T=4, U=3 test, astronomically many in practice.
  • Full B×T×U×V logits, log-space forward-backward · O(T·U) time, 1.6 GB in the example.
  • Pruned or fused loss · memory scales with kept cells (no figure quoted).

Implementation notes

• Log space; logaddexp with −∞ returns the other argument (lse).
• Blank index differs: here 0, torchaudio's rnnt_loss defaults to blank=-1; logits (batch, T, U+1, class) plus logit_lengths, target_lengths.
• Mask padding by true T_b, U_b: the terminal cell is (T_b−1, U_b).
• Keep log-softmax and recursion in f32 even if the joint runs in fp16 (general practice; overflow shows there first).
• Decode: a blank penalty (subtract a constant from the blank logit) counters deletions; tune it. An endpointer (VAD, end-of-turn detection) says when to finalise.

Pitfalls

• No label cap in greedy search: a frame that never emits blank loops forever.
• Late emissions: the loss ignores when a label appears among valid paths; regularise (FastEmit, k2 delay_penalty).
• Gradient mismatch: optimized_transducer's README reports different CPU and CUDA gradients in warp-transducer; test against brute force.

Variants and alternatives

In production

• On-device speech recognition: Google's streaming RNN-T beat a CTC-based model on latency and accuracy (He et al. 2018).
• Production speech at Microsoft: an RNN-T trained on 30 thousand hours matched a 5,120 MB server hybrid model's WER at 216 MB (Li et al. 2019).
• Toolkits: NeMo; k2 with icefall recipes; sherpa-onnx for streaming inference.
• Streaming translation: transducers are common in industrial streaming generation but need an extension for reordering (Ma et al. 2024).

Implementations

Sources

• Graves, Sequence Transduction with Recurrent Neural Networks, 2012. paper
• Graves, Mohamed, Hinton, Speech Recognition with Deep Recurrent Neural Networks, ICASSP 2013. paper
• He et al., Streaming End-to-end Speech Recognition For Mobile Devices, 2018. paper
• Li, Zhao, Hu, Gong, Improving RNN Transducer Modeling for End-to-End Speech Recognition, 2019. paper
• Yu et al., FastEmit: Low-latency Streaming ASR with Sequence-level Emission Regularization, 2020. paper
• Variani, Rybach, Allauzen et al., Hybrid Autoregressive Transducer (HAT), 2020. paper
• Kuang et al., Pruned RNN-T for fast, memory-efficient ASR training, 2022. paper · code
• Ma, Feng, Zhang, Overcoming Non-monotonicity in Transducer-based Streaming Generation, 2024. paper