On this page
The block as PyTorch classes, CTC training, and chunked attention for live audio
The Conformer is a speech recognition encoder. Each block passes every frame through a half-strength feed-forward layer, self-attention for long-range context, a convolution module for local detail, a second half feed-forward layer and a LayerNorm, the two halves forming a macaron around the middle; The Conformer Block explains why. Here those five lines become four PyTorch classes, and a CTC head and loss train them on every batch. A chunk mask in the attention then lets the same classes run on live audio at a chosen latency.
The block piece by piece, starting with the macaron FFN. Each half is a two-layer MLP whose residual is scaled by 0.5, so two of them add up to one full FFN.
import torch.nn as nn
P = 0.1 # dropout rate
class FFN(nn.Module):
def __init__(self, d):
super().__init__()
self.net = nn.Sequential(
nn.LayerNorm(d),
nn.Linear(d, 4 * d),
nn.SiLU(), # Swish
nn.Dropout(P),
nn.Linear(4 * d, d),
nn.Dropout(P))
def forward(self, x):
# half-step residual
return x + self.net(x) / 2P of values in training only; a higher rate fights overfitting, a lower one fits faster.Next, self-attention. MHSA wraps PyTorch's built-in nn.MultiheadAttention with a pre-norm and a residual. It leaves out the relative-position bonus, which is in Attention with relative positions (The Conformer Block). Its mask m is True where a key is hidden, PyTorch's convention.
MHA = nn.MultiheadAttention
class MHSA(nn.Module):
def __init__(self, d, n):
super().__init__()
self.ln = nn.LayerNorm(d)
self.attn = MHA(
d, n, dropout=P,
batch_first=True)
# m: True = key hidden
def forward(self, x, m=None):
q = self.ln(x)
kw = dict(attn_mask=m)
att = self.attn
o, _ = att(q, q, q, **kw)
return x + oq, which makes it self-attention; m hides masked keys, and the output is added back as a residual.Now the convolution module, the class Conv, named after the Conv term of the block. Conv1d wants channels first, hence the transposes. The padding keeps the length only for an odd kernel k, so PyTorch ports use 31 where the paper lists 32. At the common 4× subsampling (40 ms frames), a 32-tap kernel spans 32 × 40 ms ≈ 1.3 s of audio; the paper found 32 best on the dev sets and 65 worse on dev and test-other (the paper's Table 7).
class Conv(nn.Module):
def __init__(self, d, k):
super().__init__()
self.ln = nn.LayerNorm(d)
pad = (k - 1) // 2
self.net = nn.Sequential(
nn.Conv1d(d, 2*d, 1),
nn.GLU(dim=1),
# depthwise conv
nn.Conv1d(
d, d, k,
groups=d,
padding=pad),
nn.BatchNorm1d(d),
nn.SiLU(),
nn.Conv1d(d, d, 1),
nn.Dropout(P))
def forward(self, x):
# Conv1d wants (B, d, T)
h = self.ln(x)
h = h.transpose(1, 2)
h = self.net(h)
h = h.transpose(1, 2)
return x + hgroups=d is what keeps the channels apart.ConformerBlock puts the three classes together; its forward runs the five lines from Inside the block (The Conformer Block) in order.
class ConformerBlock(nn.Module):
def __init__(self, d, n, k):
super().__init__()
self.ffn1 = FFN(d)
self.mhsa = MHSA(d, n)
self.conv = Conv(d, k)
self.ffn2 = FFN(d)
self.ln = nn.LayerNorm(d)
def forward(self, x, m=None):
x = self.ffn1(x)
x = self.mhsa(x, m)
x = self.conv(x)
x = self.ffn2(x)
return self.ln(x)
# Conformer (M) sizes, odd kernel
blk = ConformerBlock(256, 4, 31)One linear layer gives each frame a score per vocabulary item, and softmax turns those scores (logits) into probabilities. A CTC loss (connectionist temporal classification) then aligns them with the text without anyone marking which frame belongs to which character: it lets frames emit a blank, collapses repeats, and adds up the probability of every frame-by-frame path that collapses to the right text.
import torch, torch.nn as nn
B, T, V = 8, 100, 32
enc_out = torch.randn(B, T, 256)
head = nn.Linear(256, V)
# (B, T, V) log-probabilities
lp = head(enc_out).log_softmax(-1)
# toy targets; 0 is the blank
tgt = torch.randint(1, V, (B, 20))
n = torch.full((B,), T)
tn = torch.full((B,), 20)
ctc = nn.CTCLoss(blank=0)
# CTCLoss wants (T, B, V)
lp = lp.transpose(0, 1)
loss = ctc(lp, tgt, n, tn)V tokens.n and the target lengths tn.Training teaches the encoder to transcribe: every batch runs filterbank frames through the front end and the encoder, scores the text with CTC, and updates every weight. torchaudio ships the block stack ready-made as torchaudio.models.Conformer. It takes frames that are already subsampled and projected to the model width, plus each utterance's valid length, and returns the encoded frames and their lengths; padded frames past a length are masked out of attention.
import torch, torchaudio
enc = torchaudio.models.Conformer(
input_dim=256, num_heads=4,
ffn_dim=1024, num_layers=16,
depthwise_conv_kernel_size=31,
dropout=0.1)
# 8 utterances, up to 400 frames
n = torch.randint(200, 401, (8,))
T = int(n.max())
x = torch.randn(8, T, 256)
out, n_out = enc(x, n)In a training loop the encoder sits between the front end and the loss, and the whole path runs every batch (Figure 2); a transducer loss (RNN-T, which also feeds the tokens already output back in, see Beyond CTC) can take the place of CTC.
One batch runs the front end, the encoder, and the CTC head and loss from Section 1, then backward, the optimizer step and the schedule step.
import torch.nn as nn
# 80 bands to width 256, 4x fewer
front = nn.Conv1d(
80, 256, 4, stride=4)
def train_step(fb, n, tgt, tn):
# Conv1d wants (B, 80, T)
x = fb.transpose(1, 2)
x = front(x).transpose(1, 2)
n = n // 4
out, n = enc(x, n)
lp = head(out).log_softmax(-1)
lp = lp.transpose(0, 1)
loss = ctc(lp, tgt, n, tn)
opt.zero_grad()
loss.backward()
opt.step()
sched.step()
return loss.item()zero_grad, backward, step: one optimizer update, then one schedule tick per batch.The optimizer covers every trained module, the front end, the encoder and the head. train_loader is a DataLoader that yields filterbank batches padded to their longest utterance, with their lengths and padded targets.
params = [*front.parameters(),
*enc.parameters(),
*head.parameters()]
opt = torch.optim.Adam(params,
lr=0.05 / 256**0.5,
betas=(0.9, 0.98), eps=1e-9)
def warmup(s, w=10000):
up = (s + 1) / w
down = (w / (s + 1)) ** 0.5
return min(up, down)
lrs = torch.optim.lr_scheduler
sched = lrs.LambdaLR(opt, warmup)
for batch in train_loader:
train_step(*batch)(fb, n, tgt, tn), the filterbank frames, their lengths, the padded targets and their lengths; one training step per batch.dropout defaults to 0.0, so set it to the paper's 0.1.enc.eval() before inference so BatchNorm uses its running statistics.To run in real time, attention becomes streaming attention: a mask limits what each query can see, so no frame waits on the future. Speech is ambiguous without a little right context, so each frame may peek a fixed number of frames ahead, paid for in latency.
In chunked emission, audio arrives in chunks of C frames, a frame reads every earlier frame plus keys up to the end of its own chunk and L look-ahead frames beyond it, and a chunk's tokens are emitted together once its last frame and those L frames arrive. A frame here is an encoder frame: the mask works inside the blocks, after the common 4× subsampling, so each frame is 40 ms of audio, four 10 ms filterbank frames. With chunks of 4 and look-ahead 2, a chunk's first frame waits for 3 more frames plus 2 look-ahead frames, 200 ms, and its last frame for the look-ahead only, 80 ms.
chunk_mask returns True where a query may read a key. MHSA takes the opposite, so the block gets ~allowed.
import torch
def chunk_mask(T, C, L):
t = torch.arange(T)
i, j = t[:, None], t[None, :]
# end of i's chunk, plus L
end = (i // C + 1) * C + L
return j < end
allowed = chunk_mask(12, 4, 2)
m = ~allowedm that attention expects.Look-ahead 0 is block-causal: no frame sees past its own chunk, yet tokens still wait up to C − 1 frames (Figure 3). Chunk 1 with look-ahead 0 is the plain causal mask, with no wait. A chunk as long as the utterance is ordinary offline attention.
Figure 3. ① Look-ahead 2 waits 80 to 200 ms; ② look-ahead 0 saves 80 ms.
Figure 4. Orange is the frame arriving now; green rows have emitted.
Streaming a Conformer needs three adjustments.
The third adjustment in code reuses chunk_mask and the Section 1 block blk.
import random, torch
T, L = 100, 0
x = torch.randn(8, T, 256)
# new chunk size per batch
sizes = [1, 2, 4, 8, 16, T]
C = random.choice(sizes)
allowed = chunk_mask(T, C, L)
out = blk(x, ~allowed)m that MHSA expects. blk keeps Section 1's BatchNorm and centered padding; a streaming build swaps both, as the first two adjustments say.What did later models change in the block itself? The Conformer Block ends on Branchformer and Squeezeformer. The CTC loss is taught in full in CTC Loss, the cost of attention in Transformer Cost and Efficiency, and the filterbank front end in MFCC and Other Audio Features.