docs/ARCHITECTURE.md

# Sortformer V2：双流 Speaker-Aware 说话人日志

## 一、整体思路

在原始 Sortformer 的 18 层 Transformer Encoder 基础上，引入一条与主特征流（H1）并行的 **speaker 特征流（H2）**。两条流在每一层都做 self-attention 和互相之间的 cross-attention，让说话人表征与 diarization 表征持续交互。spk 分支有自己的输出头，用 chunk 级别的 PIL loss 单独监督。

spk encoder 的表征**只在第一层之前作为输入**，此后 H2 由每层的 H2 分支自己演化，不再重新注入。

---

## 二、网络结构

```
原始音频 (16kHz, 最长300s)
    │
    ├──────────────────────────────┐
    │                              │
    ▼                              ▼
Mel Spectrogram               CAM++ Speaker Encoder（冻结backbone）
(128维, 10ms)                Fbank → 卷积backbone → 帧级特征
    │                        → Conv下采样 → 192维@80ms
    ▼                              │
FastConformer (17层, 512维)        │
8x下采样 → 512维@80ms              │
    │                              │
    ▼                              ▼
Linear 512→192                 H2 初始表征 (B,T,192)
    │                              │
    ▼                              ▼
    ┌──────────────────────────────────────────┐
    │    DualStream Transformer (18层, 192维)  │
    │                                          │
    │   每层 Block:                            │
    │     H1: self-attn → H1/H2 互attn → FFN   │
    │     H2: self-attn → H1/H2 互attn → FFN   │
    │                                          │
    └──────────────────────────────────────────┘
    │                              │
    ▼                              ▼
主 Head (预训练权重)          spk Head (随机初始化)
    │                              │
    ▼                              ▼
ATS+PIL loss                  chunk级PIL loss
```

### 每层 DualStreamTransformerBlock 细节

```
输入: H1 (B,T,192), H2 (B,T,192)

1. H1 self-attention   (first_sub_layer, 预训练权重)
   H1 = LN1(H1 + MHA(H1,H1,H1))

2. H2 self-attention   (h2_self_attn, 随机初始化)
   H2 = LN1(H2 + MHA(H2,H2,H2))

3. 互 attention（无 gate, 随机初始化）
   H1 = crossLN1(H1 + MHA(H1→H2))   # H1 attend H2
   H2 = crossLN2(H2 + MHA(H2→H1))   # H2 attend H1

4. FFN
   H1 = LN2(H1 + FFN(H1))   (second_sub_layer, 预训练权重)
   H2 = LN2(H2 + FFN(H2))   (h2_ffn, 随机初始化)
```

所有 attention 均为 **8 头**（head_dim = 192/8 = 24）。

**关键设计：H1 的参数名与原始 TransformerEncoderBlock 完全一致**（`first_sub_layer`、`layer_norm_1`、`second_sub_layer`、`layer_norm_2`），因此原始 Sortformer checkpoint 可以无缝 warm-start 到 H1 流，新增的 H2 流和互 attention 保持随机初始化。

### 权重来源总结

| 模块 | 权重来源 |
|------|---------|
| FastConformer encoder | 预训练 |
| H1 self-attn + FFN + LN | 预训练（同名加载） |
| H2 self-attn + FFN + LN | 随机初始化 |
| H1↔H2 互 attention ×2 + LN | 随机初始化 |
| 主输出 head | 预训练 |
| spk 输出 head | 随机初始化 |
| CAM++ 卷积 backbone | 预训练，**永远冻结** |
| CAM++ 下采样 conv (proj) | 随机初始化，**可训练** |

---

## 三、Speaker Encoder（CAM++）

### 提取流程

```
原始音频 (最长300s)
    │
    ▼
切分为无重叠窗口：每窗口 2 秒（chunk_dur_sec=2.0, chunk_stride_sec=2.0）
    │
    ▼  每窗口独立处理（encode_batch_size=16 批量并行）：
Kaldi Fbank (80维, 10ms) → CAM++ 卷积backbone（冻结）
    → 帧级特征 (512维, 20ms)，每帧编码周围大范围上下文
    │
    ▼
Conv1d(512→192, kernel=12, stride=4, padding=4)  ← 可训练下采样层
    → 每个输出帧融合 240ms（12帧×20ms）上下文
    → 192维 @80ms 帧率
    │
    ▼
窗口按时间顺序拼接 → (B, T_diar, 192)，与 sortformer 帧率严格对齐
```

### 设计说明

- **不做池化**：CAM++ 原始的 utterance-level TSTP 全局池化不适用于逐帧场景。改为可学习卷积下采样，让 H2 流自己的 self-attention 做时间整合。
- **kernel=12（240ms）**：每个 H2 帧融合 240ms 的 CAM++ 帧级特征，比输出帧率（80ms）宽 3 倍，提供更丰富的局部上下文。
- **proj 可训练**：下采样卷积是新参数（~98K），随训练学习如何把帧级特征映射到 H2 空间；CAM++ backbone 冻结。
- **2s 窗口**：与 CAM++ 训练时长一致，窗口间无重叠，显存可控。

---

## 四、训练算法

### 双损失

```
主分支:  loss1 = 0.5 × ATS_loss + 0.5 × PIL_loss   （原始 Sortformer 损失）
spk分支: loss2 = chunk级 PIL loss
总损失:  loss = loss1 + spk_pil_weight × loss2     （spk_pil_weight=1.0）
```

### Chunk 级 PIL Loss（spk 分支）

主分支沿用原始做法：整段拼接后一次匈牙利排列匹配。spk 分支不同——按 `chunk_len=188` 帧（15.04秒）切块：

```
spk_preds / targets 按 188 帧切成 n 个 chunk
每个 chunk 内部独立做 PIL 排列匹配 → 各自算 BCE
最后所有 chunk 的 loss 取平均
```

每个 chunk 是独立的排列问题，模型被强制在短窗口内做说话人区分，与流式推理场景一致。

### 两种训练模式

| 模式 | 配置 | 可训练参数 | 说明 |
|------|------|-----------|------|
| 只训新增 | `freeze_base_model: true` | ~13.4M | H2流 + 互attn + spk head + proj |
| 全量训练 | `freeze_base_model: false` | ~131M | 所有参数（除 CAM++ backbone） |

两种模式下 **CAM++ 卷积 backbone 均冻结**（proj 下采样层可训练）。

### 学习率分组（全量模式）

| 参数组 | 学习率 |
|--------|--------|
| 原始 Sortformer 参数 | `lr=2e-5` |
| 新增 H2 流 + 互attn + spk head + proj | `speaker_lr=1e-4` |
| CAM++ backbone | 冻结 |

---

## 五、流式模式

- 流式推理时每 chunk 处理 `[spkcache | fifo | chunk]` 拼接序列
- `spkcache_spk`/`fifo_spk` 同步缓存对应帧的 speaker embedding，压缩时用相同的 `topk_indices` gather，保证 H2 与 H1 每帧时刻严格对齐
- spk 分支（H2 → spk head）只用于训练监督，推理时只用主分支输出

---

## 六、超参数

### 新增超参数

| 参数 | 值 | 说明 |
|------|-----|------|
| `use_dual_stream` | true | 启用双流架构 |
| `spk_pil_weight` | 1.0 | spk 分支 chunk PIL loss 权重 |
| `freeze_base_model` | false/true | 只训新增 / 全量训练 |
| `speaker_lr` | 1e-4 | 新增参数学习率（全量模式） |
| `speaker_encoder.chunk_dur_sec` | 2.0 | spk encoder 窗口长度（无重叠） |
| `speaker_encoder.chunk_stride_sec` | 2.0 | 窗口步长（=窗口长，无重叠） |
| `speaker_encoder.encode_batch_size` | 24 | 多窗口批量并行 |
| `speaker_encoder.downsample_kernel` | 12 | 下采样卷积核（240ms 上下文） |
| `speaker_encoder.checkpoint_path` | null | 权重从 warm-start nemo 加载，不硬编码 |

### 沿用超参数

| 参数 | 值 |
|------|-----|
| 采样率 | 16000 |
| 最大说话人数 | 8 |
| 最大音频时长 | 300s |
| batch size | 1 |
| 优化器 | AdamW (β=[0.9,0.98]) |
| 基础学习率 | 2e-5 |
| 权重衰减 | 2e-3 |
| 最大 epoch | 10 |
| chunk_len | 188 (15.04s) |
| spkcache_len | 376 |
| 因果注意力概率 | 0.5 (rc=7) |
| dropout (attn/ffn) | 0.5 |

---

## 7. 文件结构

```
fintune_sortformer_v2/
├── docs/ARCHITECTURE.md                      ← 本文档
├── checkpoints/
│   └── dual_stream_init_8spk.nemo            ← 初始化nemo（含全部权重）
├── src/
│   ├── finetune_pipeline/
│   │   ├── conf/
│   │   │   └── streaming_sortformer_diarizer_8spk-v1-finetune.yaml
│   │   │       ← use_dual_stream, spk_pil_weight, freeze_base_model, speaker_lr
│   │   └── scripts/
│   │       ├── init_dual_stream_nemo.py       ← 生成含spk encoder权重的nemo
│   │       └── ckpt2nemo.py                  ← ckpt转nemo（含全部新权重）
│   └── third_party/nemo/collections/asr/
│       ├── models/sortformer_diar_models.py  ← 双流接入、spk head、chunk PIL loss、两种训练模式
│       └── modules/
│           ├── speaker_encoder.py            ← 冻结CAM++ backbone + 可训练Conv下采样
│           ├── sortformer_modules.py         ← spk cache 追踪（spkcache_spk/fifo_spk）
│           └── transformer/
│               ├── dual_stream_transformer.py ← 双流block + encoder
│               └── transformer_encoders.py    ← 原始单流（use_dual_stream=false时用）
└── wespeaker/                                ← CAM++ 模型
```

---

## 8. 初始化与训练流程

```bash
# 1. 生成初始化 nemo（H1=预训练, H2/互attn=随机, CAM++=冻结, proj=随机）
python src/finetune_pipeline/scripts/init_dual_stream_nemo.py \
    --pretrained /path/to/original_sortformer.nemo \
    --output /path/to/output_dual_stream.nemo

# 2. 训练时 warm-start 这个 nemo
bash run_train.sh --init_nemo_path /path/to/output_dual_stream.nemo

# 3. 训练后 ckpt 转 nemo（ckpt 已含全部权重）
python src/finetune_pipeline/scripts/ckpt2nemo.py \
    /path/to/checkpoint.ckpt /path/to/output.nemo
```