Blame for README.md

# AURA: Auditory Spatial Intelligence via a Unified Representation from a General-Audio Encoder through Complete Alignment

<p align="center">
  <a href="https://arxiv.org/abs/2611.13687"><img src="https://img.shields.io/badge/arXiv-Paper-b31b1b?logo=arxiv&logoColor=white" alt="arXiv"></a>
  &nbsp;&nbsp;&nbsp;&nbsp;
  <a href="https://github.com/burgerIO-823/AURA"><img src="https://img.shields.io/badge/GitHub-Code-181717?logo=github&logoColor=white" alt="GitHub"></a>
</p>

## Model introduction

AURA is a spatial audio question-answering model built on **Qwen3-8B** and a frozen
**X-Codec** general-audio encoder, with **64.6M trainable parameters** in LoRA adapters
and an audio projector. Given binaural audio and a text question, it identifies sound
sources and reasons about their directions, distances, and spatial relationships.
AURA is trained on **SpatialSoundQA** with a curriculum spanning single-source perception,
two-source perception, and spatial reasoning, using scene descriptions and geometric
supervision. On SpatialSoundQA, the main model achieves **48.21% dual-source direction
accuracy**, outperforming the strongest dedicated spatial audio baseline by **10.31
percentage points**, while retaining general-audio question-answering capabilities.

Code, data synthesis and training recipes: **https://github.com/burgerIO-823/AURA**

## Model Summary

| Component | Configuration |
|---|---|
| Language backbone | Qwen3-8B; base weights frozen |
| Audio encoder | X-Codec general-audio v2; encoder, RVQ, and codebooks frozen |
| Trainable components | LoRA adapters and a two-layer audio projector; 64,626,688 parameters in total |
| LoRA configuration | Rank 16, alpha 32, dropout 0.05, no bias |
| LoRA target modules | `q_proj`, `k_proj`, `v_proj`, `o_proj`, `gate_proj`, `up_proj`, `down_proj` |
| Audio projector | Linear (1024 → 4096) → GELU → Linear (4096 → 4096) |
| Audio input | Binaural left/right channels; 16 kHz; up to 10 seconds |
| Codec representation | 8 RVQ layers at 50 frames/s per ear; 1024-dimensional codebook embeddings |
| RVQ bands | Content: sum of layers 0–2; sound: sum of layers 3–7 |
| Channel representation | Left, right, and side (`S = (L − R) / 2`), each split into content and sound bands |
| Temporal pooling | Factor 4; max pooling for content, mean pooling for sound |
| Audio tokens | Up to 750 tokens, prepended to the text prompt embeddings |
| Training objective | Language modeling loss on target text, with scene descriptions and geometric supervision |
| Output format | Scene description followed by the final answer; geometric analysis for relational questions |

## Evaluation Results

Spatial performance on the full **SpatialSoundQA** evaluation sets. Scores are percentages;
**bold** marks the best result and <u>underline</u> the second best, excluding Random.

| Model | Trainable params | Input | Single DoA ↑ | Single DER ↓ | Dual DoA ↑ | Dual DER ↓ | Direction BA ↑ | Distance BA ↑ |
|---|---:|---|---:|---:|---:|---:|---:|---:|
| Random (BAT Table 4) | — | P | 12.57 | 67.33 | 12.41 | 67.46 | 50.00 | 50.00 |
| Qwen2-Audio (zero-shot) | — | M+P | 8.85 | 80.03 | 9.74 | 80.27 | 49.32 | 46.14 |
| Qwen2.5-Omni (zero-shot) | — | M+P | 10.61 | 69.50 | 12.02 | 79.26 | 49.55 | 54.30 |
| BAT | <u>161.04M</u> | B+P | 72.28 | 64.55 | 36.64 | <u>70.98</u> | <u>68.54</u> | **71.35** |
| DSpAST-BAT† (1 epoch) | 161.16M | B+P | **80.84** | 60.73 | <u>37.90</u> | 75.06 | 58.07 | 47.72 |
| OWL (released Stage-2 checkpoint) | <u>161.04M</u> | B+P | 9.82 | <u>48.17</u> | 13.36 | 71.68 | N/R | N/R |
| **AURA** | **64.63M** | B+P | <u>73.58</u> | **42.69** | **48.21** | **64.72** | **68.86** | <u>70.65</u> |

- **DoA**: exact accuracy over the eight direction octants (left/right × front/behind × above/below).
- **DER**: distance error rate; an answer is incorrect when its absolute distance error exceeds 0.5 m.
- **BA**: binary accuracy on direction or distance relationship questions.

Evaluation uses raw audio amplitudes without per-source RMS normalization and greedy decoding.
The single-source, dual-source, and spatial-reasoning sets contain 29,656, 21,280, and
32,298 items, respectively. AURA uses 192 new tokens for single-/dual-source tasks and
384 for spatial reasoning.

`B` = binaural audio, `M` = mono audio formed with a left-30° beam, and `P` = text prompt.
The zero-shot Qwen baselines use `M+P`. `N/R` indicates unavailable results: OWL has no
public Stage-3 checkpoint. BAT and OWL results use their released checkpoints under
the evaluation protocol above; DSpAST-BAT† uses the released encoder with a Q-Former
and Llama LoRA trained for one epoch.

Trainable parameter counts exclude frozen backbones.

AURA achieves the best result on four of the six spatial metrics. Its dual-source DoA
accuracy exceeds the strongest baseline by **10.31 percentage points**, using about
**40%** as many trainable parameters.

## Model Use

### Download the model

Install the Hugging Face Hub client:

```bash
pip install "huggingface_hub==0.36.2"
```

Download the main checkpoint and shared codebooks with the CLI:

```bash
hf download BryanHan/AURA --local-dir checkpoints \
  --include "stage3_c3/*" "codebooks.pt" "README.md"
```

Or download the same files in Python:

```python
from huggingface_hub import snapshot_download

snapshot_download(
    repo_id="BryanHan/AURA",
    local_dir="checkpoints",
    allow_patterns=["stage3_c3/*", "codebooks.pt", "README.md"],
)
```

To download an earlier stage, replace `stage3_c3` with `stage1_c1` or `stage2_c2`. Each stage contains all
its trained components and can be downloaded independently.

### Run inference

Set up the AURA code repository and download the frozen Qwen3-8B and X-Codec components:

```bash
git clone https://github.com/burgerIO-823/AURA AURA
cd AURA
pip install -r requirements.txt -r requirements-encode.txt
bash setup_models.sh
```

Run the model download command above from the AURA directory. Following the GitHub
example, prepare SpatialSoundQA at `data/SpatialSoundQA` (including the AudioSet clips
and binaural room impulse responses; see the GitHub data preparation instructions),
then render the example binaural scene and run inference:

```bash
python examples/make_example.py --ssqa data/SpatialSoundQA

python inference.py --ckpt checkpoints/stage3_c3 \
  --codebooks checkpoints/codebooks.pt \
  --audio examples/example.wav \
  --question "Does the sound from the hammer emanate from the front or behind side of the music's sound?"
```

The rendering script writes `examples/example.wav` and prints the question and dataset
reference answer (`front`). The reference answer is not a guaranteed model prediction.
To use your own recording, replace `examples/example.wav` and the question accordingly.

Inference requires a CUDA GPU. The input must have two channels (left/right); it is
resampled to 16 kHz and truncated to 10 seconds. The model generates a scene
description followed by a final answer. The provided AURA loader reads the LoRA
weights, projector, and architecture configuration together.

For audio already encoded by `data_preparation/encode_audio.py`, replace `--audio`
with `--codes`:

```bash
python inference.py --ckpt checkpoints/stage3_c3 \
  --codebooks checkpoints/codebooks.pt \
  --codes my_binaural_codes.pt \
  --question "What sound is to the left of the listener?"
```

The `.pt` file must contain `codes_left` and `codes_right`. This path uses Qwen3-8B
and the frozen codebooks without running the waveform encoder.

### Verify a complete download

To download all three stages and verify their weights and configurations:

```bash
hf download BryanHan/AURA --local-dir checkpoints
cd checkpoints
sha256sum -c SHA256SUMS
```

## Checkpoints in this repository

| Folder | Stage | Training data (cumulative) | Notes |
|---|---|---|---|
| `stage1_c1/` | Stage 1 (curriculum C1) | 55K alignment + 279K single-source scenes | single-source perception |
| `stage2_c2/` | Stage 2 (curriculum C2) | + 236K dual-source scenes (515K) | two-source perception and binding |
| `stage3_c3/` | Stage 3 (C3), **main model** | + 358K geometric-relation questions | spatial reasoning |
| `codebooks.pt` | – | – | the 8 frozen X-Codec RVQ codebooks (8×1024×1024, fp32) |

Each stage folder contains the trained components and their configurations. The frozen
Qwen3-8B backbone and the X-Codec encoder checkpoint are downloaded separately.
There are **64,626,688 trainable parameters** per stage: **43,646,976 LoRA** parameters
and **20,979,712 projector** parameters.

- `aura_model.safetensors`: LoRA weights (r=16, α=32, dropout 0.05) on the q/k/v/o/gate/up/down
  projections of Qwen3-8B; `adapter_config.json` is the matching PEFT configuration;
- `proj.pt`: 2-layer GELU MLP projector (1024 → 4096 → 4096);
- `aura_config.json`: input representation (six L/R/S slots, RVQ cut at layer 3, time pool 4).

All three stage folders have the same four-file layout shown above. Each checkpoint
contains all its trained components; loading Stage 3 does not require Stages 1 and 2.
The shared `codebooks.pt` contains frozen lookup embeddings, rather than the waveform encoder.

## Training

Trained on 4× NVIDIA H100 80GB (Python 3.10, torch 2.13.0+cu130, transformers 4.51.0, peft 0.20.0),
1 epoch per stage, bf16, cosine schedule, seed 42, pure LM loss on the target text:

| Stage | Init | Rows | Global batch | LR | Steps |
|---|---|---:|---:|---:|---:|
| C1 align (projector only) | – | 55,008 | 32 | 1e-4 | 1,719 |
| **Stage 1** C1 joint | C1 align | 278,784 | 16 | 1e-4 | 17,424 |
| **Stage 2** C2 joint | Stage 1 | 514,784 | 16 | 1e-4 | 32,174 |
| **Stage 3** C3 joint | Stage 2 | 872,312 | 16 | 1e-4 | 54,520 |

## License

AURA's LoRA adapters and audio projector weights are released under the
[Apache License 2.0](https://www.apache.org/licenses/LICENSE-2.0).
Third-party components, including Qwen3-8B, X-Codec, and its codebooks, retain their
respective licenses.

## Citation

If you use AURA in your research, please cite:

```bibtex
@misc{han2026auraauditoryspatialintelligence,
  title={AURA: Auditory Spatial Intelligence via a Unified Representation from a General-Audio Encoder through Complete Alignment},
  author={Boyan Han and Xiaoyan Wang and Zhe Sun and Ziqi Wang and Yujun Cai and Hongli Cao and Yiwei Wang and Ming-Hsuan Yang and Liang An},
  year={2026},
  eprint={2611.13687},
  archivePrefix={arXiv},
  primaryClass={cs.SD},
  url={https://arxiv.org/abs/2611.13687},
}
```