forced_align returns a sub-optimal path when two predecessor states tie #4221

🐛 Describe the bug

forced_align can return an alignment that is not a most-likely path. It happens when two of a cell's three predecessor states have exactly equal scores.

The per-cell choice in forced_align_impl (cpu/compute.cpp#L113-L122) is:

if (x2 > x1 && x2 > x0) {          // skip the blank
  result = x2; backPtr = 2;
} else if (x1 > x0 && x1 > x2) {   // advance one
  result = x1; backPtr = 1;
} else {                           // stay
  result = x0; backPtr = 0;
}

When x1 == x2 and both are greater than x0, the first condition fails (x2 > x1 is false) and the second fails too (x1 > x2 is false), so the else branch takes x0 — the strictly worst of the three. The DP then propagates a value that is not the maximum, and the backtrace follows a sub-optimal path.

The GPU kernel carries the identical chain (gpu/compute.cu#L90-L99).

Minimal reproduction

Three frames, a vocabulary of three, two target tokens — the whole search space is 27 paths, so the optimum can be checked by hand:

import torch, torchaudio.functional as F

log_probs = torch.tensor([[[0., -1., -1.],
                           [0., -1., -2.],
                           [-2., 0., 0.]]])
targets = torch.tensor([[1, 2]], dtype=torch.int32)

path, scores = F.forced_align(log_probs, targets, blank=0)
print(path[0].tolist(), scores.sum().item())

Actual: [1, 2, 2], total score -3.0
Expected: total score -1.0 — reached by [1, 0, 2] and by [0, 1, 2], both of which collapse to the target [1, 2]

The five paths that spell the target score -1.0, -1.0, -2.0, -3.0, -5.0. The kernel returns the fourth-best.

How often it fires

Exhaustive check, comparing the kernel's total score against the best score over all enumerated valid paths, on 3243 randomly generated small cases (T ≤ 5, vocabulary ≤ 4, integer log-probabilities so that ties are exact):

On continuous-valued emissions exact ties are rare, which is presumably why this went unnoticed, and it also means the fix is not disruptive:

• 43 recorded reference alignments (seeded random log-softmax emissions, including heavy-repeat targets and tight T == L + repeats fits) are bit-identical before and after, and the existing test_forced_align expectations are unaffected.
• On a real workload — a 737 s speech recording aligned against its transcript with a wav2vec2 CTC model, 36 820 frames and 10 044 target tokens — the tie condition fires 55 921 times out of 740 M cells but never on the winning path, and all 2 253 word timestamps come out identical.

So the defect is latent on ordinary emissions and bites on quantised or otherwise degenerate ones, where exact ties are common.

Two independent cross-checks that the fixed behaviour is the correct one: a scalar Python port of this exact C++ loop reproduces torchaudio on 234/234 cases with the clause and reproduces the fixed kernel on 234/234 without it; and an independent NumPy Viterbi agrees with the fixed kernel on 277/277 cases (it agreed with stock on only 264).

The fix

Reaching the second branch already implies that x2 is not strictly greater than both others, so x1 > x0 alone is correct and complete there:

• branch 1 gives x2, and x2 > x1, x2 > x0 — the maximum;
• branch 2 gives x1 with x1 > x0; if x2 > x1 held, then x2 > x1 > x0 would have taken branch 1, so x2 ≤ x1 — the maximum;
• branch 3 gives x0 with x0 ≥ x1; if x2 > x0 held, then x2 > x0 ≥ x1 would have taken branch 1, so x2 ≤ x0 — the maximum.

Deleting && x1 > x2 therefore changes nothing except routing exact ties to the larger value. PR to follow.

Versions

The clause dates back to #3354 (2023-05-22) and is unchanged on main today. The reproduction above is from the released wheel; the fix was additionally built and verified from main at 4e3e282.

torchaudio : 2.11.0
torch      : 2.13.0
Python     : 3.12.8
OS         : macOS-27.0-arm64-arm-64bit
machine    : arm64
numpy      : 2.2.6
CUDA       : False

The defect is in device-independent code and the CUDA kernel carries the identical chain, so it is not specific to this platform.

cc @pearu @NicolasHug