wav2tok 2.0: Scalable Audio Tokenization Maintaining Explicit Pairwise Token Alignment for Efficient Audio Retrieval

Abstract: Learning discrete speech representations that preserve similarity across variable-length utterances is central to query-by-example spoken term detection (QbE-STD). While wav2tok introduced CTC-based sequence alignment to enforce token consistency, its tightly coupled clustering and alignment training recipe limits scalability. We propose wav2tok 2.0, a scalable alignment-aware speech tokenizer built on the BEST-STD backbone. wav2tok 2.0 employs staged training, first learning discriminative, speaker-invariant representations via contrastive learning and vector quantization, and then enforcing pairwise token consistency using a CTC alignment loss and a novel DTW-aligned framewise prediction objective with adaptive weighting. Experiments show that wav2tok 2.0 consistently outperforms BEST-STD and general-purpose tokenizers on QbE-STD while remaining efficient and scalable.

• The paper introduces explicit pairwise token alignment with a composite CTC and DTW framework to enhance retrieval accuracy.
• It employs a two-phase training regime combining SimCLR-style pretraining and fine-grained DTW-based framewise loss for robust, speaker-invariant embeddings.
• Empirical benchmarks on LibriSpeech and TIMIT show state-of-the-art performance in token consistency and query-by-example retrieval.

Scalable Audio Tokenization with Explicit Pairwise Token Alignment for Efficient Audio Retrieval

Introduction

Discrete speech tokenization models are increasingly central for audio-based retrieval tasks, notably query-by-example spoken term detection (QbE-STD), where the objective is to retrieve utterances containing a spoken query directly from the audio. Despite considerable progress in self-supervised learning frameworks, most general-purpose speech tokenizers lack mechanisms for alignment-aware consistency, which limits their efficacy for retrieval scenarios requiring robust and speaker-invariant token alignment. The "wav2tok 2.0" framework (2606.26824) addresses this deficiency by introducing an explicit, scalable, alignment-driven training paradigm while preserving the efficiency and discriminative qualities established by previous methods such as BEST-STD.

Architectural Innovations and Learning Framework

wav2tok 2.0 adopts a staged representation learning strategy leveraging the bidirectional Mamba-enhanced speech tokenizer backbone of BEST-STD. The architecture comprises a spectrogram frontend and a bidirectional Mamba-based state-space encoder producing normalized frame-level embeddings. These are discretized via cosine-similarity-based vector quantization. Training is split into two phases: discriminative pretraining and pairwise alignment.

Discriminative Pretraining

In the first training phase, the framework employs SimCLR-style self-supervised contrastive objectives to yield discriminative, speaker-invariant embeddings. Frame-level anchor–positive pairings are constructed using monotonic DTW alignments, and a commitment penalty regularizes the encoder–centroid relationship. This ensures formation of a robust latent space conducive to explicit alignment objectives.

Explicit Pairwise Token Alignment

The second stage introduces two explicit alignment mechanisms: CTC-based pairwise sequence alignment and a novel DTW-aligned framewise token prediction loss. The CTC objective enforces likelihood maximization of deduplicated token sequences derived from paired utterances, without requiring the blank symbol typically used in CTC, improving alignment precision. In parallel, the DTW-aligned framewise loss enforces fine-grained token prediction consistency along DTW paths, leveraging cross-view agreement at frame resolution.

Figure 1: Stage II pairwise alignment framework combining CTC-based sequence alignment with a novel DTW-aligned framewise token prediction.

Adaptive loss weighting stabilizes optimization by scaling the CTC loss relative to the contrastive loss, addressing the sensitivity of the CTC path marginalization to latent space geometry.

Empirical Evaluation

Token Consistency Analysis

wav2tok 2.0 was benchmarked against BEST-STD, HuBERT, WavLM, EnCodec, and SpeechTokenizer on LibriSpeech. The framework yields the highest unigram and bigram Jaccard similarity for token sequences across all codebook sizes, with especially pronounced gains at the bigram level. This reflects enhanced preservation of both phonetic content and local sequential structure, especially in the presence of speaker and channel variation.

The paper highlights that explicitly modeled alignment—first via CTC (in wav2tok) and further advanced with DTW-aligned framewise prediction (in wav2tok 2.0)—results in a significant improvement in paired-token consistency relative to both general-purpose and discriminative, but implicitly aligned, tokenizers.

QbE-STD Retrieval Performance

On QbE-STD benchmarks (LibriSpeech and TIMIT), wav2tok 2.0 consistently achieves state-of-the-art Mean Average Precision (MAP), Mean Reciprocal Rank (MRR), and Maximum Term Weighted Value (MTWV) for both in-vocabulary (IV) and out-of-vocabulary (OOV) queries. Gains exceed those of general-purpose models and prior retrieval-oriented architectures across matching and cross-domain evaluation, with particularly robust generalization in OOV and cross-corpus setups.

The framework also demonstrates that fine-grained DTW-aligned supervision not only strengthens pairwise alignment but also enhances subword structure generalization to unseen lexical items, a critical property for continual learning and open-vocabulary speech retrieval.

Theoretical and Practical Implications

wav2tok 2.0 formalizes alignment-aware tokenization by bridging the gap between discriminative contrastive learning and explicit sequence alignment. The explicit CTC and DTW-based losses provide stronger supervision signals than implicit alignment from positive pair mining, directly optimizing token consistency. The removal of blank CTC transitions exploits the monotonicity of deduplicated token sequences, increasing model stability and training convergence. The decoupling of representation learning and alignment constraints addresses previous scalability barriers, as in wav2tok, by removing clustering–alignment dependencies.

This paradigm enables scalable, efficient inverted-index-based retrieval with high token consistency, supporting realistic audio retrieval in sizable archives. The framework is extensible to incorporate codebook balancing via optimal transport or noisy environmental augmentation, and the strong alignment signal is relevant for speech LLMs, speech understanding, and universal tokenizer design for cross-modal models (2606.26824).

Future Directions

Potential extensions involve integrating optimal transport regularization for further codebook balancing and extending alignment-aware objectives to multilingual and noisy regimes. The methodology is highly compatible with recent speech LLM pretraining frameworks, and the demonstrated gains in alignment fidelity recommend wav2tok 2.0 architectures as a stable paradigm for large-scale, context-invariant, and explainable speech discrete tokenization pipelines.

Conclusion

wav2tok 2.0 (2606.26824) establishes a robust and scalable approach to alignment-consistent discrete speech representation, outperforming both prior retrieval tokenizers and large self-supervised encoders on QbE-STD and token consistency. The introduction of a composite CTC and DTW-aligned supervision regime enables greater efficiency, generalization, and interpretability in audio tokenization, with implications for the next generation of retrieval-oriented and speech language foundation models.