AG-REPA — Model Weights

Companion model release for the paper "AG-REPA: Causal Layer Selection for Representation Alignment in Audio Flow Matching" Pengfei Zhang, Tianxin Xie, Minghao Yang, Li Liu — ICML 2026.

💻 Code: github.com/zpforlove/AG-REPA | 📄 Paper & poster: ICML 2026 Virtual | 🌐 Language: English | 简体中文

This repository hosts the final AG-REPA checkpoints for the unified TTS+TTA Flow-Matching framework described in the paper — everything you need to run inference. The training/inference source code lives on GitHub: zpforlove/AG-REPA.

This is a slim, inference-ready release: only the best/final (epoch-3) checkpoints of the AG-REPA models are included. The no-alignment baselines, intermediate epochs, and the FoG-A / LASP diagnostic artifacts are not bundled here — they can be reproduced from the code (the attribution figures are shown in the code README and the paper).

⚠️ Third-party base models are NOT included. The frozen BEATs and CosyVoice-300M / ttsfrd models must be downloaded from their original sources — see §3.

Total size: ~15 GB.

1. The system in brief

A two-stage cascade (see the code repository for full detail):

Stage 1 — Autoregressive LLM (Qwen3-0.6B-Base, fine-tuned) predicts discrete acoustic tokens (S³ for speech, AudioSet for audio) from text + reference style.
Stage 2 — Flow-Matching DiT transforms noise into a mel-spectrogram conditioned on those tokens; Vocos decodes the mel to a 24 kHz waveform.

AG-REPA aligns the causally dominant DiT layers (found by the FoG-A probe: speech L1/L9/L5, audio L1/L21/L9) to the Whisper (speech) / BEATs (audio) teachers, with attribution-proportional weights — yielding the quality gains in §5.

2. Directory layout

AG-REPA (model repo)
├── audioset_tokenizer/                       # Stage-0: AudioSet Tokenizer (AST), ~0.34 GB
│   └── best_epoch_3_loss_0.0883.pth          #   RepCodec VQ-VAE on BEATs features (vocab 4096)
│
├── llm/                                       # Stage-1: autoregressive LLM (Qwen3-0.6B-Base)
│   ├── single_codebook/                       #   Config A
│   │   └── best_model_loss_2.1508_epoch_3.pth #   (~4.5 GB)
│   └── dual_codebook/                         #   Config B
│       └── best_model_loss_1.8048_epoch_1.pth #   (~3.4 GB)
│
└── flow_matching/                             # Stage-2: DiT Flow-Matching (AG-REPA, final ep3)
    ├── agrepa_single_codebook/                #   Config A
    │   └── best_ep3_val_loss_0.2161_step869394.pt   #   (~3.1 GB)
    └── agrepa_dual_codebook/                  #   Config B
        └── best_ep3_val_loss_0.1211_step1304091.pt  #   (~3.7 GB)

(not included — download separately, see §3)
└── pretrained_base_models/    # BEATs + CosyVoice-300M + CosyVoice-ttsfrd

Conventions

single vs dual codebook — Config A (S³ + AudioSet tokens) vs Config B (Config A + interleaved BEATs tokens). Pick whichever matches your use case.
Checkpoint filenames encode the training epoch, the Flow-Matching validation loss, and the optimizer step, e.g. best_ep3_val_loss_0.2161_step869394.pt.

val_loss is the FM regression loss used for checkpoint selection — not the FAD reported in the paper (FAD is computed post-hoc on VGGish embeddings at 16 kHz).

3. What each component is

Component	In this repo?	Source / architecture	Used by
`audioset_tokenizer/*.pth`	✅	RepCodec VQ-VAE (vocab 4096), trained on BEATs features	Discretises general audio into AudioSet tokens
`llm/*`	✅	Qwen3-0.6B-Base, fine-tuned on LibriSpeech + AudioSet	Stage-1 token prediction (TTS + TTA)
`flow_matching/agrepa_*`	✅	24-layer DiT, adaLN-Zero, trained with AG-REPA	Stage-2 velocity-field prediction → mel
BEATs (`BEATs_iter3_plus_AS2M.pt`)	❌ download	BEATs (Chen et al., 2022) — https://github.com/microsoft/unilm/tree/master/beats	Acoustic teacher; audio-token features; LLM style embeddings
CosyVoice-300M	❌ download	CosyVoice (Du et al., 2024) — https://huggingface.co/FunAudioLLM/CosyVoice-300M	S³ speech tokenizer (`speech_tokenizer_v1.onnx`) + frontend
CosyVoice-ttsfrd	❌ download	CosyVoice text frontend — https://www.modelscope.cn/models/iic/CosyVoice-ttsfrd	Text normalisation for TTS inference

The Whisper semantic teacher is loaded directly from the openai-whisper package (large-v3) at training time and is not needed for inference.

After downloading, place the base models under a local pretrained_base_models/ folder (or directly into the code variant's pretrained_models/, see §4).

4. Using the weights with the code

Clone the code repository, download this model repo, and download the base models from their upstream sources. The code (per variant directory) expects a pretrained_models/ folder and a checkpoints/{ast,llm,flow}/ tree.

Pairings published here (the AG-REPA models — use the code's REPA_* variants):

Code variant	Stage-1 LLM	Stage-2 Flow-Matching
`REPA_single_codebook`	`llm/single_codebook`	`flow_matching/agrepa_single_codebook`
`REPA_dual_codebook`	`llm/dual_codebook`	`flow_matching/agrepa_dual_codebook`

# 1) download this model repo
hf download AustinZhang/AG-REPA --local-dir AG-REPA-Model

# 2) wire it into the AG-REPA single-codebook code variant
cd AG-REPA/REPA_single_codebook
ln -s /path/to/pretrained_base_models                            pretrained_models   # BEATs / CosyVoice (see §3)
mkdir -p checkpoints
ln -s /path/to/AG-REPA-Model/audioset_tokenizer                  checkpoints/ast
ln -s /path/to/AG-REPA-Model/llm/single_codebook                 checkpoints/llm
ln -s /path/to/AG-REPA-Model/flow_matching/agrepa_single_codebook checkpoints/flow

Then run inference as described in the code README.

Want the no-alignment baselines, all training epochs, or the FoG-A/LASP diagnostics? They are not in this slim release — train them from the code (the Fusion_* variants produce the baselines and the diagnostic artifacts).

5. Headline results (from the paper)

AG-REPA vs. the best fixed-layer REPA baseline on the unified DiT, Config B (Table 2):

Method	Speech WER ↓	Speech FAD ↓	Audio FAD ↓	Speech MOS ↑	Audio MOS ↑
Baseline (None)	5.82	1.84	3.45	3.62	3.45
REPA @ L4,8,12 (best fixed)	4.93	1.45	2.88	3.92	3.77
AG-REPA (Top-3)	3.45	1.29	2.56	4.12	3.94

The method targets the Store–Contribute Dissociation: the deep layers that store the most information (Cos-SEM top-3 = L24/L18/L17) are not the shallow layers that causally contribute to the velocity field (FoG-A top-3 = L1/L9/L5 for speech, L1/L21/L9 for audio). AG-REPA aligns the latter. AG-REPA also transfers across architectures (Voicebox, CosyVoice, F5-TTS). See the paper and the code README for the full attribution figures (Figure 1 / Table 1).

6. Licensing & responsible use

The trained AG-REPA checkpoints in this repository are released under the MIT License. Note that the Stage-1 LLM checkpoints (llm/*) are fine-tuned from Qwen3-0.6B-Base (Apache-2.0); downstream use should also respect that upstream license. The third-party BEATs and CosyVoice-300M / ttsfrd base models are not redistributed here and retain their own original licenses — obtain them from the sources in §3.

As noted in the paper's Impact Statement, high-fidelity audio generation and voice cloning carry risks (deepfakes, impersonation, voice spoofing). Responsible deployment should incorporate audio watermarking, spoofing detection, and restricted access to voice-cloning capabilities.

7. Citation

@inproceedings{zhang2026agrepa,
  title     = {AG-REPA: Causal Layer Selection for Representation Alignment in Audio Flow Matching},
  author    = {Zhang, Pengfei and Xie, Tianxin and Yang, Minghao and Liu, Li},
  booktitle = {Proceedings of the International Conference on Machine Learning (ICML)},
  year      = {2026}
}

Downloads last month: -; Downloads are not tracked for this model. How to track