Dual-mode self-supervised speech models (S3Ms), which jointly pre-trained in the offline and online mode, suffer from attention mismatch in streaming scenarios due to missing future context. To address this challenge, we proposed online registers, learnable tokens appended to each chunk in online mode. These tokens act as virtual placeholders for unseen future frames, enabling the model to compensate for missing context without introducing additional latency. Furthermore, we introduce a future prediction loss that explicitly guides the registers to capture predictive cues, thereby enriching their ability to retain future information. Experiments on LibriSpeech, and out-of-domain benchmarks demonstrate that online registers consistently reduce the performance gap between offline and online modes, achieving a 3.4% relative improvement on LibriSpeech with 160 ms chunks, especially in low-latency settings.
Figures & tables
Figure 1: Overview of our proposed pre-training framework with online registers . As an example, we illustrate the case where the feature length is 4 frames, the chunk size is 2, and the number of online registers per chunk is 1. The dual-mode Transformer encoder processes offline input with full-context attention and online input with chunk-wise attention, where online registers are appended. The model is trained to predict quantized targets for masked frames (dotted line boxes), with an additional future prediction loss encouraging the online registers to store future context.
Figure 2: Attention mask design for the online mode. As an example, we illustrate the case where the feature length is 6 frames, the chunk size is 2, the look-ahead size is 1, and the number of online registers per chunk is 1. During attention computation, the model attends only to the past and current chunks, the look-ahead, and the online registers, while all other attention weights (white boxes) are filled with −∞
Pre-train Method
dev-clean
dev-other
test-clean
test-other
FLEURS
Offline
Online
Offline
Online
Offline
Online
Offline
Online
Offline
Online
Dual-mode (Baseline)
2.11
2.98
6.50
9.89
2.73
3.65
6.63
10.15
21.95
32.65
+ Online Registers
2.08
2.87
6.41
9.38
2.70
3.50
6.52
9.80
21.72
32.02
+ Future Prediction
2.02
2.82
6.47
9.60
2.67
3.51
6.65
10.16
22.35
32.36
Table 1: Word error rate (WER, %) of baseline and proposed models on LibriSpeech test-clean , test-other , and FLEURS. All models are fine-tuned only on LibriSpeech 960h. All online decoding is performed with a chunk size of 160 ms and without look-ahead. We use a single online register per chunk.
Method
test-clean
test-other
Offline
Online
Offline
Online
wav2vec 2.0 [ 4 ]
2.6
-
6.1
-
UFO2 [ 11 ]
3.0
3.8
7.1
9.4
Ours
2.7
3.5
6.5
8.5
Table 2: WER (%) of different methods on LibriSpeech test-clean and test-other . All online decoding is performed with a 640 ms chunk size and without look-ahead. UFO2 does not use an external LM but instead applies attention-based rescoring.
#Registers
test-clean
test-other
Offline
Online
Offline
Online
R=0
2.73
3.65
6.63
10.15
R=1
2.70
3.50
6.52
9.80
R=2
2.73
3.48
6.70
10.01
R=3
2.69
3.44
6.64
9.95
R=4
2.80
3.56
6.72
10.19
Table 3: WER (%) for different numbers of Online Registers on LibriSpeech test-clean and test-other . All online decoding is performed with a 160 ms chunk size without look-ahead. All models were pre-trained without future prediction loss.
Pre-train Method
640 ms
320 ms
160 ms
test-clean
Dual-mode (Baseline)
3.21
3.37
3.65
+ Online Registers
3.12
3.25
3.50
+ Future Prediction
3.05
3.27
3.51
test-other
Dual-mode (Baseline)
8.68
9.40
10.15
Table 4: WER (%) for test-clean and test-other with different chunk sizes in online scenario. All decoding was performed without look-ahead. We use a single online register per chunk.