feat: implement CR-CTC (Consistency-Regularized CTC) loss

CR-CTC improves CTC training by computing KL divergence between encoder
outputs from augmented and clean inputs. This enforces prediction consistency
regardless of data augmentation.

k2/icefall reports 3.98% CER on AISHELL-1 test with Zipformer + CR-CTC.

Usage: Add model_conf.cr_ctc_weight=0.2 to enable CR-CTC during training.
The clean encoder pass is done without SpecAug while the augmented pass
uses the normal SpecAug pipeline.

Also fixes: add_sos_and_eos function for EParaformer compatibility.

Author: LauraGPT

funasr/losses/cr_ctc.py

@@ -0,0 +1,55 @@
+"""Consistency-Regularized CTC (CR-CTC) loss.
+
+Based on: "Improving CTC-based Speech Recognition via Consistency Regularization"
+Key idea: Run encoder twice (with/without SpecAug), compute KL divergence between
+the two CTC outputs as a consistency regularization term.
+
+Usage in training:
+    cr_loss = cr_ctc_loss(ctc_logprobs_aug, ctc_logprobs_clean, input_lengths)
+    total_loss = ctc_loss + cr_loss_scale * cr_loss
+"""
+
+import torch
+import torch.nn.functional as F
+
+
+def cr_ctc_loss(
+    log_probs_aug: torch.Tensor,
+    log_probs_clean: torch.Tensor,
+    input_lengths: torch.Tensor,
+) -> torch.Tensor:
+    """Compute CR-CTC consistency regularization loss.
+
+    Computes symmetric KL divergence between augmented and clean encoder outputs.
+
+    Args:
+        log_probs_aug: CTC log probabilities from augmented input (B, T, V)
+        log_probs_clean: CTC log probabilities from clean input (B, T, V)
+        input_lengths: Valid lengths for each sample (B,)
+
+    Returns:
+        Scalar loss value (mean over batch and time).
+    """
+    batch_size, max_len, _ = log_probs_aug.shape
+
+    # Create mask for valid positions
+    mask = torch.arange(max_len, device=input_lengths.device)[None, :] < input_lengths[:, None]
+    mask = mask.unsqueeze(-1)  # (B, T, 1)
+
+    # Convert log probs to probs for KL computation
+    probs_aug = log_probs_aug.exp()
+    probs_clean = log_probs_clean.exp()
+
+    # Symmetric KL divergence: 0.5 * (KL(p||q) + KL(q||p))
+    # KL(p||q) = sum(p * (log_p - log_q))
+    kl_aug_to_clean = (probs_aug * (log_probs_aug - log_probs_clean)) * mask
+    kl_clean_to_aug = (probs_clean * (log_probs_clean - log_probs_aug)) * mask
+
+    # Mean over valid positions
+    num_valid = mask.sum()
+    if num_valid > 0:
+        loss = 0.5 * (kl_aug_to_clean.sum() + kl_clean_to_aug.sum()) / num_valid
+    else:
+        loss = torch.tensor(0.0, device=log_probs_aug.device)
+
+    return loss

funasr/models/transformer/model.py

@@ -13,6 +13,7 @@
 
 # from funasr.models.e2e_asr_common import ErrorCalculator
 from funasr.train_utils.device_funcs import force_gatherable
+from funasr.losses.cr_ctc import cr_ctc_loss
 from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank
 from funasr.utils import postprocess_utils
 from funasr.utils.datadir_writer import DatadirWriter
@@ -127,6 +128,7 @@ def __init__(
         self.vocab_size = vocab_size
         self.ignore_id = ignore_id
         self.ctc_weight = ctc_weight
+        self.cr_ctc_weight = kwargs.get("cr_ctc_weight", 0.0)
         self.specaug = specaug
         self.normalize = normalize
         self.encoder = encoder
@@ -235,6 +237,22 @@ def forward(
             # calculate whole encoder loss
             loss_ctc = (1 - self.interctc_weight) * loss_ctc + self.interctc_weight * loss_interctc
 
+        # CR-CTC: consistency regularization
+        loss_cr_ctc = None
+        if self.cr_ctc_weight > 0.0 and self.training and self.ctc_weight != 0.0:
+            # Second forward pass WITHOUT SpecAug
+            specaug_backup = self.specaug
+            self.specaug = None
+            encoder_out_clean, encoder_out_lens_clean = self.encode(speech, speech_lengths)
+            self.specaug = specaug_backup
+            if isinstance(encoder_out_clean, tuple):
+                encoder_out_clean = encoder_out_clean[0]
+            # Compute CTC log probs for both augmented and clean
+            ctc_logprobs_aug = self.ctc.log_softmax(encoder_out)
+            ctc_logprobs_clean = self.ctc.log_softmax(encoder_out_clean).detach()
+            loss_cr_ctc = cr_ctc_loss(ctc_logprobs_aug, ctc_logprobs_clean, encoder_out_lens)
+            stats["loss_cr_ctc"] = loss_cr_ctc.detach()
+
         # decoder: Attention decoder branch
         loss_att, acc_att, cer_att, wer_att = self._calc_att_loss(
             encoder_out, encoder_out_lens, text, text_lengths
@@ -248,6 +266,10 @@ def forward(
         else:
             loss = self.ctc_weight * loss_ctc + (1 - self.ctc_weight) * loss_att
 
+        # Add CR-CTC loss
+        if loss_cr_ctc is not None:
+            loss = loss + self.cr_ctc_weight * loss_cr_ctc
+
         # Collect Attn branch stats
         stats["loss_att"] = loss_att.detach() if loss_att is not None else None
         stats["acc"] = acc_att

funasr/models/transformer/utils/add_sos_eos.py

@@ -29,3 +29,24 @@ def add_sos_eos(ys_pad, sos, eos, ignore_id):
     ys_in = [torch.cat([_sos, y], dim=0) for y in ys]
     ys_out = [torch.cat([y, _eos], dim=0) for y in ys]
     return pad_list(ys_in, eos), pad_list(ys_out, ignore_id)
+
+def add_sos_and_eos(ys_pad, sos, eos, ignore_id):
+    """Add <sos> at the beginning and <eos> at the end (length + 2).
+
+    Unlike add_sos_eos which returns (ys_in, ys_out) separately,
+    this returns a single sequence with both sos and eos added.
+
+    :param torch.Tensor ys_pad: batch of padded target sequences (B, Lmax)
+    :param int sos: index of <sos>
+    :param int eos: index of <eos>
+    :param int ignore_id: index of padding
+    :return: ys_in with sos prepended (B, Lmax+1)
+    :return: ys with both sos and eos (B, Lmax+2)
+    """
+    _sos = ys_pad.new([sos])
+    _eos = ys_pad.new([eos])
+    ys = [y[y != ignore_id] for y in ys_pad]
+    ys_in = [torch.cat([_sos, y], dim=0) for y in ys]
+    ys_both = [torch.cat([_sos, y, _eos], dim=0) for y in ys]
+    return pad_list(ys_in, eos), pad_list(ys_both, ignore_id)
+