# XNet encoder 与 MobileMamba encoder 实现对比

_面向 `X_SSL_Net` 当前实现，比较 `lib/modules/xnet_2d.py` 与 `ref/fft_wave_pro/MobileMamba/model/mobilemamba/mobilemamba.py` 的 encoder 设计。_

---

## 1. 对比范围

本文只比较 encoder / backbone 部分，不比较 XNet 的 decoder、segmentation head，也不比较 MobileMamba 的分类 head。

对比代码入口：

| 项目 | 文件 | 主要类 / 函数 |
| --- | --- | --- |
| XNet 当前实现 | `lib/modules/xnet_2d.py` | `XNetStem2d`、`XNetDownsample2d`、`XTEB2d`、`XNetEncoder2d` |
| MobileMamba 参考实现 | `ref/fft_wave_pro/MobileMamba/model/mobilemamba/mobilemamba.py` | `MobileMamba`、`MobileMambaBlock`、`MobileMambaModule`、`MBWTConv2d`、`PatchMerging` |

结论先写清楚：两者都使用了 convolution、wavelet、SS2D/Mamba-like global modeling，但组织方式完全不同。**XNet 是为分割设计的 4 级多尺度 encoder，输出 `[E1,E2,E3,E4]` 给 decoder skip 使用；MobileMamba 是分类 backbone，patch embedding 后只向前堆叠 blocks，最后全局池化分类。**

## 2. 总体数据流

### 2.1 XNet encoder

XNet encoder 的目标是产生分割 decoder 可用的多尺度特征金字塔。

```text
Input [B, 3, H, W]
  |
  v
Stem: two stride-2 conv steps
  |
  v
E1 [B, C1, H/4,  W/4]
  |
  v
Down1 + Stage2
  |
  v
E2 [B, C2, H/8,  W/8]
  |
  v
Down2 + Stage3
  |
  v
E3 [B, C3, H/16, W/16]
  |
  v
Down3 + Stage4
  |
  v
E4 [B, C4, H/32, W/32]
```

默认配置是：

```text
encoder_channels = [32, 64, 128, 192]
encoder_depths   = [2, 2, 2, 2]
```

`XNetEncoder2d.forward()` 直接返回：

```python
[e1, e2, e3, e4]
```

这说明它天然是 segmentation encoder，而不是分类 backbone。

### 2.2 MobileMamba encoder / backbone

MobileMamba 的目标是分类。它先用 4 个 stride-2 convolution 直接把输入降到 `1/16`，然后经过 3 组 blocks，中间通过 `PatchMerging` 再降采样。

```text
Input [B, 3, H, W]
  |
  v
PatchEmbed: four stride-2 convs
  |
  v
X0 [B, C1, H/16, W/16]
  |
  v
blocks1
  |
  v
blocks2: pre refine + PatchMerging + post refine + stage blocks
  |
  v
X1 [B, C2, H/32, W/32]
  |
  v
blocks3: pre refine + PatchMerging + post refine + stage blocks
  |
  v
X2 [B, C3, H/64, W/64]
  |
  v
AdaptiveAvgPool2d + BN_Linear
```

典型配置如 `MobileMamba_B2`：

```text
img_size = 384
embed_dim = [200, 376, 448]
depth = [2, 3, 2]
global_ratio = [0.8, 0.7, 0.6]
local_ratio = [0.2, 0.2, 0.3]
kernels = [7, 5, 3]
```

它默认不返回多尺度特征，只返回分类 logits。因此如果要把它直接用于分割，必须额外改 `forward_features` 或在 `blocks1/2/3` 后取中间特征。

## 3. Stem / patch embedding 差异

| 维度 | XNet | MobileMamba |
| --- | --- | --- |
| 初始降采样倍率 | `1/4` | `1/16` |
| 结构 | `Conv s2 -> DWConv -> PWConv -> Conv s2` | 4 个连续 `Conv s2` |
| 输出通道 | 默认 `32` | 默认 `embed_dim[0]`，如 `200` |
| 设计目的 | 保留较高分辨率给分割 skip | 快速进入低分辨率分类 token/feature stage |
| 对小病灶/边界 | 更友好，因为 `E1` 是 `H/4` | 早期细节丢失更多，原始实现不适合直接当分割浅层 skip |

这一点是最重要的结构差异。XNet 的 stem 只降到 `H/4`，因为 decoder 需要浅层边界、纹理和定位信息。MobileMamba 一开始就降到 `H/16`，这对 ImageNet 分类是合理的，但对医学分割尤其是超声小病灶不一定合适。

## 4. Stage 组织方式差异

### 4.1 XNet 是显式 4 级 feature pyramid

XNet 写法很直接：

```python
e1 = self.stage1(self.stem(x))
e2 = self.stage2(self.down1(e1))
e3 = self.stage3(self.down2(e2))
e4 = self.stage4(self.down3(e3))
return [e1, e2, e3, e4]
```

每一级都是：

```text
Downsample if needed -> XTEB2d x depth
```

它的边界非常清楚：stage 负责表征提取，downsample 负责降尺度，返回值负责 decoder skip。

### 4.2 MobileMamba 把 downsample 插入下一组 blocks

MobileMamba 在构造 blocks 时把 stage 的降采样模块 append 到下一组 blocks 里：

```text
blocks1 = stage0 blocks
blocks2 = downsample 0->1 + stage1 blocks
blocks3 = downsample 1->2 + stage2 blocks
```

`PatchMerging` 前后还各有一段 residual depthwise conv + FFN refine：

```text
Residual(DWConv)
Residual(FFN)
PatchMerging
Residual(DWConv)
Residual(FFN)
```

这对分类 backbone 来说很紧凑，但如果要做 segmentation backbone，中间特征取点需要更小心。因为 `blocks2` 不是纯 stage2，它包含从 stage1 到 stage2 的 transition。

## 5. Block 内部信息组织差异

### 5.1 XNet 的 XTEB2d：全通道并行三分支

`XTEB2d` 对同一个输入特征做三条并行路径：

```text
pre_norm(x)
  |
  +-- local_branch(x)   : 3x3 depthwise + 5x5 depthwise
  |
  +-- wavelet_branch(x) : DWT -> LL/high conv -> IDWT
  |
  +-- global_branch(x)  : VMamba SS2D
  |
  v
concat -> 1x1 fuse -> channel gate -> post conv -> residual -> FFN
```

关键点：

- local、wavelet、global 三个分支都处理完整 `C` 个通道
- 三个输出 concat 成 `3C`，再压回 `C`
- `XBranchFusion2d` 额外用 `GAP -> MLP -> sigmoid` 做通道 gate
- 分支融合后还有 `post` residual 和 `ffn` residual

这是一种“全通道多视角建模”：每个通道组都能同时被 local、wavelet、global 路径处理，代价是计算和显存更高。

### 5.2 MobileMambaModule：按通道切分三路

MobileMamba 不把完整通道送进每个分支，而是先按比例切分：

```python
x1, x2, x3 = torch.split(
    x,
    [global_channels, local_channels, identity_channels],
    dim=1,
)
```

三路分别是：

```text
global channels  -> MBWTConv2d
local channels   -> DWConv2d_BN_ReLU
identity channels -> Identity
```

然后 concat 回来，再用：

```text
ReLU -> Conv2d_BN(dim, dim, bn_weight_init=0)
```

做输出投影。

这是一种“通道分工式建模”：部分通道做重型 global+wavelet，部分通道做轻型 local，剩余通道直接保留 identity。它比 XNet 的全通道并行三分支更省算力，也更符合 MobileNet 风格。

## 6. Wavelet 实现差异

| 维度 | XNet | MobileMamba |
| --- | --- | --- |
| wavelet 库 | `ptwt.wavedec2 / waverec2` | `pywt` 生成固定 filter，再用 `conv2d / conv_transpose2d` |
| 默认类型 | 只允许 `haar` | 默认 `db1`，本质上与 Haar 很接近 |
| 分解层数 | 当前只支持 `level=1` | `wt_levels` 可扩展，当前调用为 `1` |
| 高频处理 | `LH/HL/HH` concat 后 depthwise + pointwise conv | `LL/LH/HL/HH` reshape 成 `4C` 后 depthwise conv |
| 低频处理 | `LL` 单独 `3x3 Conv2dBN + ReLU` | `LL` 和 high bands 一起进入 wavelet conv 后再拆分 |
| 逆变换 | `ptwt.waverec2`，显式 crop | `conv_transpose2d`，按记录 shape crop |
| AMP 处理 | 显式关闭 autocast，内部 float32 | 没有单独 autocast 保护 |

XNet 的 wavelet branch 更像一个独立的“频率专家分支”：它把低频和高频拆开处理，再重建到原空间尺度。MobileMamba 的 wavelet 更深地嵌在 `MBWTConv2d` 中，和 SS2D global attention 共同构成 global branch。

## 7. SS2D / Mamba 使用差异

| 维度 | XNet | MobileMamba |
| --- | --- | --- |
| 使用位置 | `XGlobalBranch2d`，作为 XTEB 三分支之一 | `MBWTConv2d.global_atten`，只作用于 global channel split |
| 作用通道 | 默认 stage2-4 全通道，stage1 可关闭 | 只作用于 `global_channels` |
| `d_state` | 默认 `16` | 固定 `1` |
| `ssm_ratio` | `global_ratio`，默认 `2.0`，且至少 `1.0` | 配置 `ssm_ratio`，常见为 `2` |
| `initialize` | `"v0"` | `"v2"` |
| `forward_type` | 默认 `"v3"` | 默认 `"v052d"` |
| backend 控制 | `ssm_backend=auto/oflex/torch`，运行时切换 | 无显式 backend 切换 |
| `k_group` | 未传 | `k_group=2` |

这里不能简单说谁更强。XNet 的 SS2D 更重，默认 `d_state=16`，且在较深 stage 全通道使用，表达能力更强但显存和时间开销更高。MobileMamba 的 SS2D 更轻，`d_state=1`，且只对一部分 global channels 使用，更适合移动端分类骨干。

## 8. Local branch 差异

XNet local branch：

```text
branch3: DWConv3x3 -> ReLU -> PWConv1x1
branch5: DWConv5x5 -> ReLU -> PWConv1x1
output = branch3 + branch5
```

MobileMamba local branch：

```text
DWConv kxk -> BN -> ReLU -> grouped 1x1 conv -> BN
```

MobileMamba 的 kernel 随 stage 改变：

```text
kernels = [7, 5, 3]
```

XNet 每个 stage 的 local branch 都固定使用 `3x3 + 5x5` 双尺度局部分支。对超声分割来说，这个设计更稳定、更直接；MobileMamba 的大核到小核递减更像分类 backbone 的经验设计，浅层更大感受野，深层更省计算。

## 9. Fusion / residual 差异

| 维度 | XNet | MobileMamba |
| --- | --- | --- |
| 分支融合 | concat 三个完整分支输出，然后 `1x1 fuse + channel gate` | split 后各自处理，concat 后 `ReLU + Conv2d_BN` |
| 残差结构 | `x_in + post(fusion(...))`，再 `x + ffn(x)` | block 内 `dw0 -> ffn0 -> mixer -> dw1 -> ffn1` 多重 residual，最后 `shortcut + drop_path(x)` |
| stochastic depth | 无 | 有 `DropPath` |
| BN 初始化 | 多处最后 BN weight 初始化为 `0` | FFN / proj 也使用 `bn_weight_init=0` |
| 部署优化 | 未提供 BN fuse 接口 | 多个模块有 `fuse()`，并提供 `replace_batchnorm()` |

MobileMamba 的 block 工程上更偏部署友好：`Conv2d_BN`、`DWConv2d_BN_ReLU`、`BN_Linear` 都实现了 fuse。XNet 当前更偏研究实验清晰性，没有做推理融合优化。

## 10. 输出形式和任务适配差异

这是两者最根本的任务差异。

| 维度 | XNet encoder | MobileMamba backbone |
| --- | --- | --- |
| 原任务 | 2D segmentation | image classification |
| 输出 | `[E1,E2,E3,E4]` | class logits |
| 中间特征 | 显式返回四级特征 | 默认不返回 |
| 最浅层分辨率 | `H/4` | patch embed 后为 `H/16` |
| 最深层分辨率 | `H/32` | 通常到 `H/64` |
| 是否直接可接 U-Net decoder | 是 | 否，需要改 forward 并补浅层特征 |

如果把 MobileMamba 直接替换为 XNet encoder，会遇到两个问题：

1. 缺少 `H/4` 和 `H/8` 的浅层 skip，decoder 的边界恢复会受影响。
2. 默认 forward 返回分类 logits，不返回 segmentation decoder 需要的多尺度特征。

## 11. 相同点总结

两者的共同点主要有：

1. 都是 channel-first 2D feature map 实现，不是 token-first Transformer 实现。
2. 都使用 depthwise convolution 建模局部空间关系。
3. 都引入 wavelet transform，把频率/尺度信息纳入 feature mixing。
4. 都使用 SS2D/Mamba-like 模块建模长程依赖。
5. 都通过 residual + FFN 稳定训练。
6. 都使用 Conv-BN-ReLU 风格的轻量卷积组件。

这些共同点说明 XNet encoder 的方向和 MobileMamba 有同源性：都是把 CNN、wavelet、SSM/Mamba 混合起来，而不是单纯 CNN 或单纯 Transformer。

## 12. 核心差异总结

| 类别 | XNet encoder | MobileMamba encoder |
| --- | --- | --- |
| 任务导向 | 分割 | 分类 |
| 多尺度输出 | 显式 4 级输出 | 默认不输出中间层 |
| 初始分辨率保留 | 保留到 `H/4` | 快速降到 `H/16` |
| 分支组织 | 全通道并行 local / wavelet / global | 通道切分 global / local / identity |
| Wavelet 位置 | 独立 branch | 嵌入 global branch 的 `MBWTConv2d` |
| SS2D 范围 | 多数 stage 全通道 | global split 部分通道 |
| 计算风格 | 更重、更适合精细分割 | 更轻、更适合分类/移动端 |
| 下采样 | 简单 stride-2 conv | `PatchMerging` + SE + pre/post refine |
| 工程部署 | 偏研究清晰 | 提供 BN fuse，偏部署 |

一句话概括：**XNet encoder 是把 MobileMamba 类思想改造成分割特征金字塔后的实现，但它没有照搬 MobileMamba 的通道切分和移动端压缩策略，而是选择了更重的全通道三分支融合。**

## 13. 对当前 XNet 的启发

MobileMamba 中值得借鉴的点：

1. **通道切分策略**  
   XNet 当前 local、wavelet、global 都处理全通道，计算开销大。可以做一个轻量变体，把通道拆成 global/local/identity 三组，作为显存友好的 ablation。

2. **PatchMerging 的 SE transition**  
   XNet 的 `XNetDownsample2d` 是简单 `3x3 stride-2 conv`。MobileMamba 的 `PatchMerging` 有 `1x1 expand -> depthwise stride-2 -> SE -> 1x1 project`，可能改善 stage transition 的表达能力。

3. **DropPath**  
   XNet 当前没有 stochastic depth。较深配置如 `[2,2,3,2]` 或更宽通道时，可以考虑加入轻量 DropPath。

4. **BN fuse / deployment path**  
   如果后续关心推理速度，MobileMamba 的 `fuse()` 设计值得参考。

不建议直接照搬的点：

1. **不要直接使用 `1/16` patch embedding**  
   超声分割需要边界和小目标细节，直接降到 `H/16` 会削弱浅层 skip。

2. **不要直接改成分类式 forward**  
   XNet 的训练链路需要 `[E1,E2,E3,E4]`，不能只输出 pooled feature。

3. **不要盲目把 `d_state` 降到 1**  
   MobileMamba 的 `d_state=1` 是轻量分类选择。XNet 用于医学分割，是否能降需要单独消融。

4. **不要把所有 stage 都按 MobileMamba 的 global_ratio 复制**  
   XNet 的 stage1 默认关闭 global branch 是合理的，因为高分辨率 SS2D 显存开销大，且浅层更需要局部纹理。

## 14. 建议后续消融

如果要系统证明 XNet encoder 相对 MobileMamba-style encoder 的改动价值，建议做以下消融：

| 消融项 | 目的 |
| --- | --- |
| XNet full-channel tri-branch vs channel-split branch | 验证全通道三分支是否值得额外显存 |
| Simple downsample vs MobileMamba PatchMerging | 验证更强 transition 是否提升 Dice/IoU |
| `d_state=16` vs `d_state=1/4/8` | 找到 SS2D 表达能力和显存的平衡点 |
| stage1 global off vs on | 验证高分辨率 global branch 是否必要 |
| XNet wavelet branch vs MBWTConv2d-style branch | 比较独立 wavelet 分支和 wavelet+SS2D 耦合分支 |

优先级最高的是第一项和第二项。第一项直接回答“XNet 为什么不用 MobileMamba 的轻量通道切分”；第二项可能以较低风险提升 encoder transition。

## 15. 最终判断

XNet encoder 和 MobileMamba encoder 的关系不是简单的“谁抄谁”或“谁替代谁”。更准确的判断是：

```text
MobileMamba = classification-oriented, mobile-style, channel-split wavelet-SS2D backbone
XNet encoder = segmentation-oriented, full-channel tri-branch wavelet-SS2D feature pyramid
```

XNet 当前实现更适合医学分割，因为它保留了 `H/4` 到 `H/32` 的四级特征，并显式返回给 decoder。MobileMamba 更适合轻量分类，因为它快速降采样、通道切分、提供部署融合路径。若要融合两者，合理方向不是直接替换 encoder，而是把 MobileMamba 的 channel split、PatchMerging、DropPath 和 fuse 思路作为 XNet 的轻量化 ablation。