AURA / README.md
BryanHan's picture
Upload README.md with huggingface_hub
e47c25b verified
|
Raw History Blame Contribute Delete
10.5 kB

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

arXiv      GitHub

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 underline 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 161.04M B+P 72.28 64.55 36.64 70.98 68.54 71.35
DSpAST-BAT† (1 epoch) 161.16M B+P 80.84 60.73 37.90 75.06 58.07 47.72
OWL (released Stage-2 checkpoint) 161.04M B+P 9.82 48.17 13.36 71.68 N/R N/R
AURA 64.63M B+P 73.58 42.69 48.21 64.72 68.86 70.65
  • 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:

pip install "huggingface_hub==0.36.2"

Download the main checkpoint and shared codebooks with the CLI:

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

Or download the same files in 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:

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:

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:

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:

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

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