AG-REPA

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):

1. Stage 1 — Autoregressive LLM (Qwen3-0.6B-Base, fine-tuned) predicts discrete acoustic tokens (S³ for speech, AudioSet for audio) from text + reference style.
2. 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

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):

# 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):

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}
}