add README.md

README.md

@@ -12,88 +12,60 @@ language:
   - en
 ---
 
-# JEPA-Q2D2: A Low-Bitrate Speech Codec with Emergent Cross-Lingual Structure
-
-A **1.6 kbps neural speech codec** that pairs a JEPA (Joint-Embedding Predictive
-Architecture) encoder with **Q2D2**, a geometry-aware 2-D rhombic-lattice
-quantizer, and a HiFi-GAN decoder. Trained **without adversarial losses on the
-codec objective** (a frozen WavLM perceptual loss is used during decoder
-training).
-
-This repository accompanies the APSIPA ASC 2026 paper *"JEPA-Q2D2: A
-Low-Bitrate Speech Codec with Emergent Cross-Lingual Structure."*
-Code: https://github.com/anant-004/jepa-q2d2
-
-## Models
-
-| Subfolder | Operating point | Quantizer | Bitrate | Reported quality |
-|---|---|---|---|---|
-| [`jepa-q2d2-cd64-12.5hz`](./jepa-q2d2-cd64-12.5hz) | 12.5 Hz, code-dim 64 (**main model**) | Q2D2 | 1.6 kbps (100 tok/s) | PESQ **2.53**, ESTOI **0.80** |
-| [`jepa-q2d2-sigreg-cd32-25hz`](./jepa-q2d2-sigreg-cd32-25hz) | 25 Hz, code-dim 32, SIGReg co-design | Q2D2 | 1.6 kbps (100 tok/s) | ESTOI **0.79** |
-| [`teacher-cd128-fsq-12.5hz`](./teacher-cd128-fsq-12.5hz) | 12.5 Hz, code-dim 128 (**teacher**) | FSQ | ~2.85 kbps (237.5 tok/s) | PESQ ~**2.91** |
-
-All metrics are on the paper's fixed 50-utterance LibriLight protocol. ESTOI is
-**extended** STOI (systematically lower than vanilla STOI; comparable only
-within this paper's identical pipeline).
-
-### Main codec (cd64, 12.5 Hz)
-The headline reconstruction system. At 1.6 kbps it exceeds EnCodec at a
-comparable operating point (+0.91 PESQ, +0.15 ESTOI vs EnCodec-1.5 kbps) and
-beats Mimi on PESQ (+0.25) while trailing it on ESTOI. Under a matched internal
-ablation, Q2D2 improves over finite scalar quantization (FSQ) by **+0.37 PESQ**
-at the same bitrate.
-
-### SIGReg co-design model (cd32, 25 Hz)
-Demonstrates the paper's central finding: at the aggressive 25 Hz / 32-dim
-operating point the codec **collapses** (ESTOI -0.004) unless the encoder's
-latent distribution is Gaussianized with **SIGReg** (lambda = 0.05), which
-restores normal training (ESTOI 0.79). Two checkpoints are identical except for
-this term.
-
-### Teacher codec (cd128, FSQ, 12.5 Hz)
-The original higher-rate codec (code dim 128, FSQ quantizer, 237.5 tok/s,
-≈2.85 kbps, PESQ ~2.91). Architecturally distinct from the Q2D2 models — it uses
-finite scalar quantization rather than the Q2D2 lattice. Released because it
-serves as the **distillation teacher** for the 1.6 kbps cd64 student and is a
-useful higher-quality reference point.
-
-## Each model folder contains
-- `pytorch_model.pt` — inference checkpoint (`ckpt["state_dict"]` = full
-  encoder + Q2D2 quantizer + HiFi-GAN decoder; optimizer / discriminator state
-  stripped).
-- `model.safetensors` — the same weights in safetensors format.
-- `config.json` — strides, code dim, frame rate, bitrate, sample rate (24 kHz),
-  training step, and the checkpoint's eval metric.
-
-## Usage
+# jepa-q2d2
+
+a 1.6 kbps neural speech codec. a jepa (joint-embedding predictive architecture) encoder feeds q2d2, a geometry-aware 2d rhombic-lattice quantizer, and a hifi-gan decoder. the codec objective is trained without any adversarial loss (a frozen wavlm perceptual loss is used while training the decoder).
+
+- code: https://github.com/anant-004/jepa-q2d2
+- paper: apsipa asc 2026
+
+## models
+
+each folder has `pytorch_model.pt`, `model.safetensors`, and `config.json`.
+
+| folder | operating point | quantizer | bitrate | quality |
+|--------|-----------------|-----------|---------|---------|
+| `jepa-q2d2-cd64-12.5hz` | 12.5 hz, code dim 64 (main) | q2d2 | 1.6 kbps (100 tok/s) | pesq 2.53, estoi 0.80 |
+| `jepa-q2d2-sigreg-cd32-25hz` | 25 hz, code dim 32, sigreg | q2d2 | 1.6 kbps (100 tok/s) | estoi 0.79 |
+| `teacher-cd128-fsq-12.5hz` | 12.5 hz, code dim 128 | fsq | ~2.85 kbps (237.5 tok/s) | pesq ~2.91 |
+
+quality is on a fixed 50-utterance librilight set. estoi is extended stoi, comparable only within this paper's pipeline. the main cd64 model is ahead of encodec at a comparable rate and ahead of mimi on pesq. the sigreg model is the co-design point: at the aggressive 25 hz / cd32 setting the codec collapses unless the encoder distribution is gaussianized with sigreg. the teacher is the higher-rate fsq codec used to distill the cd64 student.
+
+## emergent properties
+
+the codec is trained on english read speech only and never sees language or emotion labels. the encoder features still pick up structure it was not supervised on.
+
+cross-lingual: a 5-nn probe on frozen pre-quantization features separates 6 fleurs languages at 0.85 (chance 0.167), ahead of mimi 0.63, encodec 0.61, dac 0.45, and close to multilingual ssl models. cluster nmi against language is higher too (0.26 vs encodec 0.10, dac 0.07). the utterance-disjoint number is the same (0.85).
+
+![cross-lingual separability](assets/crosslingual_summary.png)
+
+style and emotion: utterance embeddings separate by speaking style on a hindi 5-class set (angry, excited, neutral, sad, whisper). whisper splits off cleanly and the rest form distinct regions for the jepa encoders, while encodec and dac mix them.
+
+![emotion separation](assets/emotion_tsne.png)
+
+## usage
 
 ```python
-import torch
-# model definition lives in the companion repo:
-#   git clone https://github.com/anant-004/jepa-q2d2
-from koe.fast.benchmark_codecs import build_v2_model   # see repo for exact entry point
-
-ckpt = torch.load("jepa-q2d2-cd64-12.5hz/pytorch_model.pt",
-                  map_location="cpu", weights_only=False)
-model = build_v2_model(ckpt["config"])
-model.load_state_dict(ckpt["state_dict"], strict=False)
-model.eval()
-
-# wav: (1, T) mono @ 24 kHz
-tokens = model.encode(wav)        # discrete Q2D2 tokens
-recon  = model.decode(tokens)     # reconstructed 24 kHz waveform
+import torch, soundfile as sf
+from koe.fast.hf_codec import load_codec_from_hf   # from the github repo
+
+model, info = load_codec_from_hf("main", device="cpu")   # downloads weights from here
+print(info)
+
+wav, sr = sf.read("in.wav")                              # mono, resampled to 24 khz internally
+x = torch.from_numpy(wav).float().view(1, 1, -1)
+z_q = model.encode(x)[0]                                 # discrete q2d2 tokens
+recon = model.decode(z_q)                                # 24 khz waveform
 ```
 
-See the GitHub repo for runnable encode/decode scripts and the evaluation
-harness.
+model code lives in the github repo. the released checkpoints already include the fine-tuned encoder, so inference needs only these files plus the repo code.
+
+## training data
 
-## Training data
-LibriLight (English read speech), 24 kHz. The codec is an English-trained
-reconstruction system; the cross-lingual structure reported in the paper is an
-**emergent** property of the JEPA encoder features, evaluated zero-shot on
-FLEURS.
+librilight, english read speech, 24 khz. the cross-lingual structure above is evaluated zero-shot on fleurs.
+
+## citation
 
-## Citation
 ```bibtex
 @inproceedings{shukla2026jepaq2d2,
   title     = {JEPA-Q2D2: A Low-Bitrate Speech Codec with Emergent Cross-Lingual Structure},
@@ -103,5 +75,6 @@ FLEURS.
 }
 ```
 
-## License
-CC-BY-4.0. Weights are released for research use.
+## license
+
+cc-by-4.0, released for research use.