GenAID / README.md
walston's picture
Add GenAID vs WhisAID architecture comparison
1627856 verified
|
Raw
History Blame Contribute Delete
3.88 kB
---
language:
- zh
license: apache-2.0
library_name: transformers
pipeline_tag: audio-classification
tags:
- accent-recognition
- speaker-disentanglement
- wav2vec2
---
# GenAID
GenAID is a Chinese accent encoder based on `facebook/wav2vec2-large-xlsr-53`. It produces a 64-dimensional accent embedding designed to reduce speaker information. The model recognizes nine labels: north, Sichuan, Guangdong, south, Henan, Shanghai, Wuhan, Tianjin, and Singapore.
## GenAID and WhisAID Medium side by side
The comparison excludes Whisper's text decoder because accent identification uses only its audio encoder.
| Component | GenAID (this model) | WhisAID Medium |
|---|---|---|
| Model-native input | Raw 16 kHz waveform `[B, T]` | 80-bin Whisper log-Mel `[B, 80, 3000]` |
| Backbone | Wav2Vec2 XLSR-53 Large | Whisper Medium audio encoder |
| Audio frontend | 7 convolution layers, combined stride 320 | 2 convolution layers, second layer stride 2 |
| Frontend output | `[B, Tβ€², 512]` | `[B, 1500, 1024]` |
| Feature projection | `512 β†’ 1024` | None after the CNN |
| Transformer blocks | 24 | 24 |
| Transformer width | 1024 | 1024 |
| Attention heads | 16 Γ— 64 dimensions | 16 Γ— 64 dimensions |
| FFN per block | `1024 β†’ 4096 β†’ 1024` | `1024 β†’ 4096 β†’ 1024` |
| Transformer output | `[B, Tβ€², 1024]` | `[B, 1500, 1024]` |
| Pooling | Attention-mask-aware temporal mean | Temporal mean |
| Accent embedding | `1024 β†’ 64 β†’ 64` | `1024 β†’ 256` |
| Accent output | `64 β†’ 9` | `256 β†’ 9` |
| Speaker disentanglement | `64 β†’ 336`; uniform-distribution adversarial MSE | GRL + speaker-token attention, `256 β†’ 336` |
| Task-relevant parameters | **315,530,560** | 306,473,897, excluding text decoder |
| Frozen parameters | 4,210,176 (7-layer XLSR CNN) | Audio encoder is evaluated under `torch.no_grad()` in the compared recipe |
| Effectively updated parameters | **311,320,384** | 444,249 |
`Tβ€²` is the downsampled XLSR sequence length. For a 10-second, 16 kHz waveform, `T=160,000` and `Tβ€²β‰ˆ499`.
### GenAID module details
| Module | Layers | Input β†’ output | Parameters | Training behavior |
|---|---:|---|---:|---|
| XLSR feature extractor | 7 CNN layers | `[B,T] β†’ [B,Tβ€²,512]` | 4,210,176 | Frozen |
| Feature projection + XLSR encoder | Projection + 24 Transformer blocks | `[B,Tβ€²,512] β†’ [B,Tβ€²,1024]` | 311,228,544 | Updated |
| Masked mean pooling | 1 | `[B,Tβ€²,1024] β†’ [B,1024]` | 0 | No parameters |
| Information bottleneck | 2 Γ— Linear + GELU | `1024 β†’ 64 β†’ 64` | 69,760 | Updated |
| Accent classifier | Linear, no bias | `64 β†’ 9` | 576 | Updated |
| Speaker adversarial classifier | Linear, no bias | `64 β†’ 336` | 21,504 | Updated during training |
| **Total** | | | **315,530,560** | |
The 64-dimensional output of the second GELU is the reusable accent embedding. The speaker classifier is a training-only probe: its predicted speaker distribution is optimized toward a uniform distribution, discouraging the bottleneck from retaining speaker identity. The speaker branch is not required when extracting embeddings.
## Usage
```python
import librosa
import torch
from transformers import AutoFeatureExtractor, AutoModel
repo = "walston/GenAID"
processor = AutoFeatureExtractor.from_pretrained(repo)
model = AutoModel.from_pretrained(repo, trust_remote_code=True).cuda().eval()
wav, _ = librosa.load("audio.wav", sr=16000, mono=True)
inputs = processor(wav, sampling_rate=16000, return_tensors="pt")
inputs = {key: value.cuda() for key, value in inputs.items()}
with torch.inference_mode():
output = model(**inputs)
accent_embedding = output.embedding # [batch, 64]
accent_logits = output.accent_logits # [batch, 9]
```
Loading this repository requires `trust_remote_code=True` because it includes the small GenAID bottleneck and classification heads around the standard XLS-R encoder.