TokenChain: A Discrete Speech Chain via Semantic Token Modeling

TokenChain is a fully discrete speech chain framework that integrates semantic-token ASR with a two-stage TTS architecture. The paper, accepted to ICASSP 2026, demonstrates that speech chain learning remains effective when both the intermediate interface and the underlying models are discrete and token-based.

Abstract

Machine Speech Chain, simulating the human perception-production loop, proves effective in jointly improving ASR and TTS. We propose TokenChain, a fully discrete speech chain coupling semantic-token ASR with a two-stage TTS: an autoregressive text-to-semantic model co-trained with ASR and a masked-generative semantic-to-acoustic model for synthesis only. End-to-end feedback across the text interface is enabled with straight-through argmax/Gumbel-Softmax and balanced with supervised ASR via dynamic weight averaging. Ablations examine optimal temperature schedules for in- and cross-domain transfer. Evaluation reveals TokenChain surpasses baseline accuracy 2-6 epochs earlier and yields 5-13% lower equal-epoch error with stable T2S on LibriSpeech, and reduces relative ASR WER by 56% and T2S WER by 31% on TED-LIUM with minimal forgetting, showing that chain learning remains effective with token interfaces and models.

Key Features

• 🔄 Discrete Speech Chain: Fully discrete speech chain with semantic token interfaces
• 🎯 Semantic-Token ASR: High-performance ASR using semantic tokens from SpeechTokenizer
• 🎵 Two-Stage TTS: Autoregressive text-to-semantic + masked-generative semantic-to-acoustic
• 🔗 End-to-End Feedback: Straight-through estimation with Gumbel-Softmax for gradient flow
• ⚖️ Dynamic Weight Balancing: Balanced training between ASR supervision and chain feedback
• 🌡️ Temperature Scheduling: Optimized temperature schedules for in- and cross-domain transfer
• 📊 Comprehensive Evaluation: Extensive evaluation on LibriSpeech and TED-LIUM datasets

Architecture

TokenChain employs a novel architecture that combines:

1. Semantic-Token ASR: Encoder-decoder architecture processing semantic tokens from SpeechTokenizer
2. Text-to-Semantic (T2S) Model: Autoregressive model co-trained with ASR for semantic token generation
3. Semantic-to-Acoustic (S2A) Model: Masked-generative transformer for acoustic token synthesis
4. Straight-Through Estimation: Enables end-to-end gradient flow through discrete token interfaces

Dependencies

TokenChain builds upon two major open-source frameworks:

• ESPnet: For the ASR components and speech processing utilities
• Amphion: For the TTS components and audio generation

Installation

Prerequisites

• Python 3.9+
• CUDA 12.1+ (for GPU acceleration)

Environment Setup

1. Setup System Dependencies:

   # Create conda environment
   conda create -n tokenchain python=3.9
   conda activate tokenchain
   # Install CUDA
   conda install nvidia/label/cuda-12.1.0::cuda
   # (Optional) Change pip source for faster download if necessary 
   pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple/
   

2. Clone the ESPnet repository:

   cd <any-place>
   git clone https://github.com/espnet/espnet
   

3. Setup Python environment:

   <espnet-root>/tools/activate_python.sh must be created to specify the Python interpreter used in espnet recipes. Follow the steps below:

   cd <espnet-root>/tools
   bash setup_python.sh $(command -v python3)
   

4. Install ESPnet:

   The Makefile tries to install ESPnet and all dependencies, including PyTorch. We specify the PyTorch version to be 2.3.1 and CUDA version to be 12.1.

   cd <espnet-root>/tools
   make TH_VERSION=2.3.1 CUDA_VERSION=12.1
   

5. Install ESPnet custom tools:

   cd installers
   bash install_s3prl.sh
   bash install_sph2pipe.sh
   bash install_phonemizer.sh
   bash install_pyopenjtalk.sh
   

6. Install further dependencies:

   pip install torch==2.3.1 torchaudio==2.3.1 torchvision==0.18.1 accelerate==0.31.0 transformers==4.41.2 ruamel.yaml easydict json5 beartype timm IPython encodec vocos speechtokenizer mutagen jieba cn2an jiwer pykakasi funasr zhon zhconv
   
   # Manual install for flash-attention, Causal-Conv1d and Mamba-ssm
   wget https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.3cxx11abiFALSE-cp39-cp39-linux_x86_64.whl
   pip install flash_attn-2.7.4.post1+cu12torch2.3cxx11abiFALSE-cp39-cp39-linux_x86_64.whl
   
   wget https://github.com/Dao-AILab/causal-conv1d/releases/download/v1.4.0/causal_conv1d-1.4.0+cu122torch2.3cxx11abiFALSE-cp39-cp39-linux_x86_64.whl
   pip install causal_conv1d-1.4.0+cu122torch2.3cxx11abiFALSE-cp39-cp39-linux_x86_64.whl
   
   wget https://github.com/state-spaces/mamba/releases/download/v2.2.2/mamba_ssm-2.2.2+cu122torch2.3cxx11abiFALSE-cp39-cp39-linux_x86_64.whl
   pip install mamba_ssm-2.2.2+cu122torch2.3cxx11abiFALSE-cp39-cp39-linux_x86_64.whl
   

7. Clone this repository:

   cd <any-place>
   git clone https://github.com/Merakist/TokenChain.git
   

8. Create a symbolic link to ESPnet tools in the TokenChain root:

   cd TokenChain
   ln -s <espnet-root>/tools .
   

Project Structure

TokenChain/
├── bins/                    # Framework scripts
│   ├── asr/                 # ASR framework scripts
│   ├── tts/                 # TTS framework scripts
│   └── tokenchain/          # TokenChain framework scripts
├── ckpts/                   # Pre-trained model checkpoints
│   ├── speechtokenizer_hubert_avg/  # SpeechTokenizer models
│   ├── text2semantic/       # T2S model checkpoints
│   └── tokenchain/          # TokenChain model checkpoints
├── datasets/                # Dataset caching and preprocessing
│   ├── cache_librispeech/   # LibriSpeech cached data
│   └── cache_tedlium/       # TED-LIUM cached data
├── egs/                     # Training scripts and configurations
│   ├── asr/                 # ASR training scripts and configs
│   ├── tts/                 # TTS training scripts and configs
│   └── tokenchain/          # TokenChain training scripts and configs
├── evaluations/             # Comprehensive evaluation tools
├── models/                  # Model definitions
│   ├── asr/                 # ASR model architectures
│   ├── tts/                 # TTS model architectures
│   ├── tokenchain/          # TokenChain model architectures
│   └── codec/               # Audio codec models
├── modules/                 # Core framework modules
├── optimizers/              # Optimization algorithms
├── schedulers/              # Learning rate schedulers
├── utils/                   # Utility functions
└── tools/                   # Additional ESPnet tools

Usage

Before Running

# Add TokenChain root and ESPnet root to environment PATH
export PYTHONPATH=<espnet-root>:${PYTHONPATH}
export PYTHONPATH=<TokenChain-root>:${PYTHONPATH}

Download & Extract Datasets

LibriSpeech: https://www.openslr.org/12; TED-LIUM v2: https://www.openslr.org/19

After extracting, set up the dataset symbolic link.

cd egs/asr
mkdir downloads
cd downloads
ln -s <LibriSpeech-dir> .

Download SpeechTokenizer Checkpoint

Folder speechtokenizer_hubert_avg consists two files: SpeechTokenizer.pt, and config.json. Download the checkpoint and place the speechtokenizer_hubert_avg folder inside TokenChain/ckpts.

Training ASR with LibriSpeech-100

cd egs/asr
# Modify run.sh and asr2.sh if necessary
bash run_train_librispeech_100.sh

Training T2S with LibriSpeech-100

1. Concurrent data-prep for T2S

For ASR, data prep finishes at Stage 8 when running run_train_librispeech_100.sh. After which ASR training automatically commences. Since T2S also uses LibriSpeech-100 for training, we can start making the data cache for T2S when ASR training script runs past Stage 8.

cd egs/tts/Text2Semantic
python t2s_librispeech_100_data_prep.py

This Python script will construct .pkl cache files at TokenChain/datasets/cache_librispeech

2. Training T2S Component with LibriSpeech-100

Modify model parameters in .json files, and then run training with corresponding .sh files. For TokenChain, 100M parameter models are used.

cd egs/tts/Text2Semantic
# (Optional) modify model parameters if necessary
vi exp_config_llama_100m.json
bash run_train_llama_100m.sh

Training TokenChain

1. Prepare Data - LibriSpeech 960 full set / TED-LIUM v2 dataset

IMPORTANT! Before starting to prepare data for LibriSpeech full set or TED-LIUM v2, relevant folders within data, dump, exp or themselves need special attention:

• For LibriSpeech, back up or rename asr_stats* and asr_train* within exp. Check data and dump for potential overwrites.
• For TED-LIUM, back up or rename data, dump, exp.

This is due to the fact that ESPnet uses similar naming scheme for data and training file management, thus files will get overwritten. Double-check what files will be accessed before running data-prep / training to avoid data loss or accessing wrong training data.

# Prepare LibriSpeech dataset
cd egs/asr
bash run_data_prep_librispeech_960.sh

# Prepare TED-LIUM dataset
cd egs/asr
bash run_data_prep_tedlium.sh

2. Merge pre-trained ASR and T2S into TokenChain model

cd egs/tokenchain
python tokenchain_merge.py \
  --asr-config  "egs/asr/your-exp/config.yaml" \
  --asr-ckpt    "egs/asr/your-exp/20epoch.pth" \
  --t2s-config  "egs/tts/Text2Semantic/your-config.json" \
  --t2s-ckpt    "ckpts/text2semantic/your-exp/model.safetensors" \
  --output      "ckpts/tokenchain/your-exp/token_chain_model.pth"

3. Construct resumable TokenChain checkpoint

cd egs/tokenchain
python tokenchain_construct_checkpoint.py \
  --token-chain-ckpt "ckpts/tokenchain/your-exp/token_chain_model.pth" \
  --prev-asr-ckpt    "egs/asr/your-exp/checkpoint.pth" \
  --t2s-optim        "ckpts/text2semantic/your-exp/optimizer.bin" \
  --t2s-sched        "ckpts/text2semantic/your-exp/scheduler.bin" \
  --output           "ckpts/tokenchain/your-exp/checkpoint.pth"

4. Train TokenChain

# Train with LLaMA backbone
cd egs/tokenchain
# Modify tokenchain_(_gumbel_)config.yaml if necessary, and
bash tokenchain_llama_ste_train.sh
bash tokenchain_llama_gumbel_train.sh

# Similarly, train with Mamba backbone
bash tokenchain_mamba_ste_train.sh
bash tokenchain_mamba_gumbel_train.sh

5. Train Baseline

Copy checkpoint and epoch20.pth from asr_train_librispeech_100 to asr_train_librispeech_960, and resume training on the full set with:

cd egs/asr
bash run_baseline_resume_train_librispeech_960.sh

Evaluation

Checkpoint Decouple for ASR Inference (No need for T2S)

cd egs/tokenchain
python tokenchain_decouple.py \
  --tokenchain-ckpt "ckpts/tokenchain/your-exp/token_chain_model.pth" \
  --asr-config "egs/asr/your-exp/config.yaml" \
  --t2s-config "egs/tts/Text2Semantic/your-config.json" \
  --t2s-type llama \
  --output-asr "out/asr.pth"

ASR Inference

cd egs/asr
# Modify ASR inference model at line 43, and
bash run_eval_librispeech_960.sh

T2S Inference

cd evaluations
python tokenchain_eval_accelerate.py \
  --s2a_cfg    "egs/tts/SoundStorm/your-s2a-config.json" \
  --s2a_ckpt   "ckpts/soundstorm/your-exp/model.safetensors" \
  --bpe_model  "egs/asr/data/token_list/tgt_bpe_unigram5000_ts_en/bpe.model" \
  --ckpt_roots "ckpts/tokenchain/your-exp-1" "ckpts/tokenchain/your-exp-2" \
  --cfg_llama  "egs/tts/Text2Semantic/exp_config_llama_100m.json" \
  --cfg_mamba  "egs/tts/Text2Semantic/exp_config_mamba_100m.json" \
  --eval_base  "evaluations/results" \
  --test_jsons "evaluations/librispeech_test_100.json" "evaluations/tedlium_dev_100.json" \
  --sim_model  "ckpts/wavlm_large_finetune.pth"

The --test_jsons files should be JSON files with a "test_cases" key, each entry containing "wav_path" and "target_text" fields. The --sim_model checkpoint (wavlm_large_finetune.pth) can be downloaded from the WavLM repository.

Ablation Studies

cd evaluations
python tokenchain_ablations.py \
  --cfg_soundstorm "egs/tts/SoundStorm/your-s2a-config.json" \
  --s2a_ckpt       "ckpts/soundstorm/your-exp/model.safetensors" \
  --spk_model      "ckpts/wavlm_large_finetune.pth" \
  --json           "evaluations/librispeech_test_100.json" \
  --json           "evaluations/tedlium_dev_100.json" \
  --out_root       "evaluations/ablations"

Configuration

Key configuration files are located in egs/tokenchain/:

• tokenchain_config.yaml: Main TokenChain configuration
• tokenchain_gumbel_config.yaml: Gumbel-Softmax specific settings
• asr_config.yaml: ASR component configuration
• t2s_config.yaml: Text-to-semantic model configuration

Results

Performance on LibriSpeech

• Convergence: 2-6 epochs earlier than baselines
• Error Reduction: 5-13% lower equal-epoch error
• Stable T2S: Consistent text-to-semantic performance

Performance on TED-LIUM

• ASR WER Reduction: 56% relative improvement
• T2S WER Reduction: 31% relative improvement
• Minimal Forgetting: Maintains performance across domains

Temperature Scheduling Results

• In-domain: Annealed ST-Gumbel (τ: 2.0 → 0.1) performs best
• Cross-domain: Sharper interface (τ ≈ 0.75) favors transfer learning

Citation

If you use TokenChain in your research, please cite our paper:

@article{wang2025tokenchain,
  title={TokenChain: A Discrete Speech Chain via Semantic Token Modeling},
  author={Wang, Mingxuan and Nakamura, Satoshi},
  journal={arXiv preprint arXiv:2510.06201},
  year={2025}
}

Acknowledgments

• This work builds upon ESPnet for ASR components
• TTS components are based on Amphion
• Special thanks to the SpeechTokenizer and HuBERT teams for their foundational work
• Supported by the Guangdong Introducing Innovative and Entrepreneurial Teams Program

License

This project is licensed under the MIT License. See the LICENSE file for details.

Contributing

Contributions are welcome! Please feel free to submit issues, feature requests, or pull requests.

Contact

• Mingxuan Wang: mingxuanwang1@link.cuhk.edu.cn
• Satoshi Nakamura: snakamura@cuhk.edu.cn

For questions about the implementation or to report issues, please open a GitHub issue.