Weakly Supervised Transducer (WST):

Weakly supervised transducer is a training technique introduced in this paper: https://arxiv.org/pdf/2511.04035. It introduces skip paths in the transducer’s lattice to tolerate errors in the training transcripts. This document summarises its mathematical formulation and gives the pseudocode.

1. Mathematical formulation

Notation

$\mathbf{x} = (x_1, \dots, x_T)$: acoustic input, length $T$.
$\mathbf{y} = (y_1, \dots, y_U)$: label sequence, vocabulary $V$, $|V|$ symbols.
$\phi$: blank symbol.
$\star$: garbage/uncertainty token, not a real vocabulary index.
$\lambda_\phi, \lambda_y \ge 0$: fixed bypass-arc penalties, see parametrisation.
At lattice node $(t, u)$, the joint network produces logits $h_{t,u}$ with softmax outputs:
- $S_\phi(t, u) = \mathrm{softmax}(h_{t,u})[\phi]$
- $S_y(t, u) = \mathrm{softmax}(h_{t,u})[y_{u+1}]$
- $S_k(t, u) = \mathrm{softmax}(h_{t,u})[k]$
$\star$-arc base weight:

S_\star(t, u) = \frac{1 - S_\phi(t,u)}{|V| - 1}

Aggregated transition weights:

W_\rightarrow(t,u) = S_\phi(t,u) + \lambda_\phi\, S_\star(t,u) \qquad\text{(blank + blank-bypass)}

W_\uparrow(t,u) = S_y(t,u) + \lambda_y\, S_\star(t,u) \qquad\text{(label + token-bypass)}

$P = P(\mathbf{y} \mid \mathbf{x})$, $\mathcal{L} = -\log P$.

Transitions out of $`(t, u)`$

From	Symbol	Weight	To
$`(t, u)`$	$`\phi`$	$`S_\phi`$	$`(t+1, u)`$
$`(t, u)`$	$`\star`$ (blank-bypass)	$`\lambda_\phi S_\star`$	$`(t+1, u)`$
$`(t, u)`$	$`y_{u+1}`$	$`S_y`$	$`(t, u+1)`$
$`(t, u)`$	$`\star`$ (token-bypass)	$`\lambda_y S_\star`$	$`(t, u+1)`$

Forward-backward recurrences

\alpha(t, u) = \alpha(t{-}1, u)\, W_\rightarrow(t{-}1, u) \;+\; \alpha(t, u{-}1)\, W_\uparrow(t, u{-}1)

\beta(t, u) = \beta(t{+}1, u)\, W_\rightarrow(t, u) \;+\; \beta(t, u{+}1)\, W_\uparrow(t, u)

P = \sum_{(t,u):\, t+u = c} \alpha(t, u)\, \beta(t, u) \quad \text{for any constant } c

Gradient

Define the bypass leak at node $(t,u)$:

B_\lambda(t,u) \;=\; \frac{\lambda_\phi\,\beta(t{+}1,u) + \lambda_y\,\beta(t,u{+}1)}{|V|-1}

Then:

\frac{\partial \mathcal{L}}{\partial x_{t,u,k}} = \frac{\alpha(t,u)}{P} \Big[\,
S_k\big(\beta(t, u) - B_\lambda(t,u)\big)
\;-\; \delta_{k,\,y_{u+1}}\, S_y\, \beta(t, u{+}1)
\;-\; \delta_{k,\phi}\,S_\phi\big(\beta(t{+}1, u) - B_\lambda(t,u)\big)
\,\Big]

Structural comparison to RNN-T:

Term	RNN-T	WST
Everywhere $`k`$	$`S_k\,\beta(t,u)`$	$`S_k\,(\beta(t,u) - B_\lambda)`$
$`\delta_{k,\,y_{u+1}}`$	$`-S_y\,\beta(t,u{+}1)`$	$`-S_y\,\beta(t,u{+}1)`$ (unchanged)
$`\delta_{k,\phi}`$	$`-S_\phi\,\beta(t{+}1,u)`$	$`-S_\phi\,(\beta(t{+}1,u) - B_\lambda)`$

WST = RNN-T with $\beta$ replaced by $(\beta - B_\lambda)$ in the everywhere-$k$ and blank-delta terms. Reduces to RNN-T exactly when $\lambda_\phi = \lambda_y = 0$ (then $B_\lambda = 0$).

Numerical stability notes

$\log S_\star$ is needed in the forward to form $\log W_\rightarrow$ and $\log W_\uparrow$. Computed as $\log S_\star = \mathrm{log1mexp}(\log S_\phi) - \log(|V|-1)$, where $\mathrm{log1mexp}(a) = \log(1 - e^a)$ for $a \le 0$ (use the standard two-branch stable form).
Subtractions $\beta(t,u) - B_\lambda$ and $S_\phi\,\beta(t{+}1,u) - B_\lambda$ are done in linear space at the very last step of the gradient assembly. Each subtracted pair shares a common log-prefactor (alpha/P * loss_grad * S_k or alpha/P * loss_grad * S_phi), which is multiplied in via log-addition before the final exp. This is the same numerical discipline as the existing RNN-T backward.

2. Pseudocode (log-space)

Conventions

All probabilities and partial sums in log-space.
lse(a, b) = log(exp(a) + exp(b)), computed stably as max(a,b) + log1p(exp(-|a-b|)).
log1mexp(a) = log(1 - exp(a)) for a <= 0, standard two-branch form.
log_lam_blank, log_lam_emit = log(lam_phi), log(lam_y); -inf if the corresponding lam is 0.
log_Vm1 = log(|V| - 1).
denom[t,u] = logsumexp(x[t,u,:]) is precomputed.

Per-cell log helpers

def lbx(t, u):  return x[t, u, blank_idx] - denom[t, u]  # log S_phi
def lsx(v):     return log1mexp(v) - log_Vm1               # log S_star, given log S_phi
def lex(t, u):  return x[t, u, label[u]] - denom[t, u]   # log S_y  (label[u] = y_{u+1})

2.1 Forward: alpha and beta

Structurally identical to RNN-T; only the transition weights change.

def forward(x, denom, label, T, U):
    alpha = full((T, U), -inf);  alpha[0, 0] = 0.0
    for step in range(1, T + U - 1):
        for u in valid_u_on_diagonal(step, T, U):
            t = step - u
            lbx_ = lbx(t-1, u)
            log_W_blank = lse(lbx_, log_lam_blank + lsx(lbx_))
            from_blank  = alpha[t-1, u] + log_W_blank  if t > 0 else -inf

            lbx_ = lbx(t, u-1)
            log_W_emit = lse(lex(t, u-1), log_lam_emit + lsx(lbx_))
            from_emit  = alpha[t, u-1]  + log_W_emit   if u > 0 else -inf

            alpha[t, u] = lse(from_blank, from_emit)

    beta = full((T, U), -inf);  beta[T-1, U-1] = lbx(T-1, U-1)
    for step in range(T + U - 3, -1, -1):
        for u in valid_u_on_diagonal(step, T, U):
            t = step - u
            lbx_ = lbx(t, u);  lsx_ = lsx(lbx_)
            log_W_blank = lse(lbx_, log_lam_blank + lsx_)
            from_blank  = log_W_blank + beta[t+1, u]  if t < T-1 else -inf

            log_W_emit = lse(lex(t, u), log_lam_emit + lsx_)
            from_emit  = log_W_emit   + beta[t, u+1]  if u < U-1 else -inf

            beta[t, u] = lse(from_blank, from_emit)

    return alpha, beta, loss=-beta[0, 0]

2.2 Backward: gradient w.r.t. logits

def backward(x, denom, alpha, beta, label, T, U, loss_grad):
    x_grad = zeros_like(x)
    log_P  = beta[0, 0]
    for t in range(T):
        for u in range(U):
            lbx_TU    = lbx(t, u)
            common    = log(loss_grad) + alpha[t, u] - log_P
            beta_TU   = beta[t, u]
            beta_TUp1 = beta[t, u+1] if u < U-1 else -inf
            beta_Tp1U = beta[t+1, u] if t < T-1 else -inf

            from_skip_blank = log_lam_blank + beta_Tp1U
            from_skip_emit  = log_lam_emit  + beta_TUp1
            log_B_lam = -log_Vm1 + lse(from_skip_blank, from_skip_emit)

            for k in range(V):
                grad    = common + x[t, u, k] - denom[t, u]
                my_grad = exp(grad + beta_TU) - exp(grad + log_B_lam)

                if u < U-1 and k == label[u]:
                    my_grad -= exp(grad + beta_TUp1)

                if k == blank_idx:
                    if t == T-1 and u == U-1:  # terminal
                        my_grad -= exp(grad)
                    else:
                        my_grad -= (exp(common + lbx_TU + beta_Tp1U)
                                    - exp(common + lbx_TU + log_B_lam))

                x_grad[t, u, k] = my_grad
    return x_grad

2.3 Reduction to RNN-T

When lam_blank = lam_emit = 0, log_lam_blank = log_lam_emit = -inf, so log_B_lam = -inf and exp(log_B_lam) = 0. log_W_blank collapses to lbx; log_W_emit collapses to lex. Both forward and backward reduce exactly to standard RNN-T.

3. Parametrisation

$\lambda_\phi$ and $\lambda_y$ are set by --wst_skip_blank and --wst_skip_emit, in two ways that differ from the paper.

Vocab-normalized. The bypass weight $\lambda S_\star = \lambda (1 - S_\phi)/(|V|-1)$ shrinks with $1/|V|$ at fixed $\lambda$, while nothing else in the lattice depends on $|V|$, so a tuned $\lambda$ does not transfer between the testing (1024), base (8704) and large (17408) vocabularies. The CLI takes $\tilde\lambda$ and the loss is given $\lambda = \tilde\lambda |V|$, so $\tilde\lambda$ is approximately the fraction of the non-blank mass carried by the bypass arc. Tuned on the testing model: $\tilde\lambda_\phi = 0.05$ (English-only), $0.03$ (multilingual). This does not correct for tokenization granularity — a larger vocabulary gives fewer tokens, hence fewer bypass opportunities, per utterance — so expect to retune somewhat per model.

Ramped. Bypass arcs on an untrained model just license the model to skip everything, so each $\tilde\lambda$ starts at 0 and is released only once validation WER drops below --wst_{sb,se}_wer_threshold (default 0.6), then ramps linearly to its final value over --wst_{sb,se}_ramp_steps steps (0, the default, jumps straight there).

CAIMAN-ASR