Conformer + RNN-Transducer, from scratch — streaming speech recognition for Hindi, Marathi and English

A from-scratch PyTorch reimplementation of the Conformer encoder (Gulati et al., 2020) trained with the RNN-Transducer loss (Graves, 2012), with a streaming mode, the paper's Table 3 ablations as switches, and an optional auxiliary CTC loss. Data: Google FLEURS Hindi, Marathi and English (CC-BY 4.0).

Written by hand: log-mel front-end, SpecAugment, convolution subsampling, the macaron Conformer block, Transformer-XL relative-position attention, the convolution module, the prediction and joint networks, the transducer loss (anti-diagonal forward algorithm), greedy decoding and Graves' beam search, CER / WER. PyTorch is used for tensors, autograd, nn.LSTM, nn.BatchNorm1d and Adam.

Status: code complete and verified by check.py (below). Training runs are next; results will be added here from the logs.

Paper → code

Paper	What	Code
Conformer §2, eq. 1	Macaron block: ½ FFN → MHSA → Conv → ½ FFN → LayerNorm	ConformerBlock in model.py
Conformer §2.1	Multi-head self-attention with Transformer-XL relative positions (content + position terms, biases u, v)	RelPositionAttention
Conformer §2.2, Fig. 2	Pointwise conv → GLU → depthwise conv → BatchNorm → Swish → pointwise conv	ConvModule
Conformer §2.3	Pre-norm feed-forward, 4× expansion, Swish	FeedForward
Conformer §3.1	Conformer (S): 16 layers, d 144, 4 heads, kernel 32, 1-layer LSTM decoder of 320	Config defaults
Conformer §3.1	Adam (0.9, 0.98, 1e-9), warmup 10k, peak lr 0.05/√d, dropout 0.1, L2 1e-6, SpecAugment	train.py
Conformer Table 3	Ablations: no macaron, no convolution, absolute positions, ReLU instead of Swish	--no-macaron, --no-conv, --abs-pos, --relu
Graves §2	Prediction network (labels so far) and joint network	Predictor, Joint
Graves §2.2	Forward variable α(t, u) and the loss −log P(y|x)	rnnt_loss in rnnt.py
Graves §2.4, Alg. 1	Beam search over output sequences	beam_search
Park et al. 2019	SpecAugment	spec_augment

Model size with the paper's Conformer (S) settings: 9.76M parameters (encoder 8.71M, decoder and joint 1.05M). The paper's 10.3M includes a 1k word-piece vocabulary; with 112 characters the joint and embedding are 0.58M smaller, which accounts for the difference.

Streaming

--chunk 8 --left-chunks 4 trains a model that never looks ahead more than one chunk:

• attention sees the current 8-frame chunk (320 ms) and the 4 chunks before it, never later chunks
• the depthwise convolution is causal (padded on the left only)
• the transducer decodes frame by frame, so text appears while audio arrives

The subsampling convolutions read 3 mel frames (30 ms) past their output frame; that is the only look-ahead inside a chunk. check.py proves the property: changing the audio after chunk c leaves every earlier encoder output unchanged, while a full-context model's earlier outputs do change.

Where this differs from the papers

	Paper	Here	Why
Data	LibriSpeech, 970 h, English	FLEURS, ~23 h, Hindi + Marathi + English	The target languages
Output units	1k word pieces	~112 characters	Small data: every symbol must be seen often; Devanagari matras compose
Decoder regularization	variational noise	none	Not implemented
Beam search	Graves Algorithm 1 with prefix merging	without prefix merging, at most 10·beam expansions per frame	Simpler; the cap keeps an untrained model from stalling
Kernel	32	31	An odd kernel centres on the frame

Correctness checks

python check.py (CPU, under a minute):

Check	What it proves
rnnt_loss_matches_brute_force	The loss and its gradient equal a sum over every alignment path, enumerated explicitly, on a padded batch with different lengths
rnnt_loss_on_long_sequences	The loss equals a plain double-loop dynamic program on 40 × 25 and 25 × 40 grids (more labels than frames)
relative_positions	Every (query, key) score uses the encoding of distance i − j, compared with an explicit double loop
padding_is_ignored	A clip gives the same encoder output alone and padded in a batch
streaming_never_sees_the_future	The streaming encoder's past outputs ignore future audio; the full-context one's do not
ablations_build	Every Table 3 variant builds and gives the right shapes and lengths
can_overfit	A tiny transducer memorizes 4 utterances; greedy and beam search both read them back exactly

The checks were themselves tested by breaking the code on purpose: an off-by-one in the final blank, a shifted emission index, a non-causal convolution in streaming mode, and a flipped relative distance each make the matching check fail.

Run it

pip install -r requirements.txt
python check.py
python data.py                                   # FLEURS hi + mr + en: downloads, log-mels, vocabulary
python train.py --overfit 16                     # sanity check: memorize 16 clips
python train.py --name conformer-s               # Conformer (S) transducer
python train.py --name stream --chunk 8 --left-chunks 4
python train.py --name no-conv --no-conv         # Table 3, one change at a time
python evaluate.py runs/conformer-s.pt --beam 4  # test-set CER / WER per language
python evaluate.py runs/conformer-s.pt --wav clip.wav

References

• Gulati et al., 2020. Conformer: Convolution-augmented Transformer for Speech Recognition. arXiv:2005.08100
• Graves, 2012. Sequence Transduction with Recurrent Neural Networks. arXiv:1211.3711
• Dai et al., 2019. Transformer-XL. arXiv:1901.02860
• Park et al., 2019. SpecAugment. arXiv:1904.08779
• Graves et al., 2006. Connectionist Temporal Classification. ICML
• Conneau et al., 2022. FLEURS. arXiv:2205.12446