FrankenSplit：用于移动边缘计算的浅层变分瓶颈注入高效神经特征压缩
A. Furutanpey, P. Raith and S. Dustdar, "FrankenSplit: Efficient Neural Feature Compression With Shallow Variational Bottleneck Injection for Mobile Edge Computing," in IEEE Transactions on Mobile Computing, vol. 23, no. 12, pp. 10770-10786, Dec. 2024, doi: 10.1109/TMC.2024.3381952.
1. 引言与研究动机
移动AI加速器的兴起使得对延迟敏感的应用能够在客户端执行轻量级深度神经网络（DNN）。然而，关键应用需要强大的模型，而边缘设备无法承载这些模型，必须因此卸载请求，其中高维数据将竞争有限的带宽。分割计算（Split Computing, SC）通过在设备之间分配DNN层来缓解资源效率低下的问题，但当前方法过于特定，仅能略微减少带宽消耗。
本文提出将重点从执行分割DNN的浅层转移到专注于为机器可解释性优化的变分压缩。我们引入了一个新颖的资源意识压缩模型框架，并在反映边缘设备和服务器之间非对称资源分布的环境中广泛评估我们的方法。
2. 相关工作
2.1 学习图像压缩
有损图像压缩的目标是在保留对人类感知至关重要的信息的同时最小化比特率。变换编码是有损压缩的基本框架，它将压缩任务分为去相关和量化。去相关减少像素的统计依赖性，允许更有效的熵编码，而量化将值表示为有限的整数集。
Ballé等人引入了因子化先验（FP）熵模型，并通过找到具有最小熵的表示来制定神经压缩问题。编码器网络将原始输入转换为潜在变量，捕获输入的统计依赖性。在后续工作中，Ballé等人和Minnen等人通过包含超先验作为先验的侧信息来扩展FP熵模型。
2.2 特征压缩
Singh等人通过在倒数第二层引入瓶颈并用图像分类的交叉熵替换失真损失，展示了信息瓶颈原理在压缩框架中的实用方法。Dubois等人将VIB泛化为多个下游任务，并首次正式描述了特征压缩任务。
最接近我们工作的是Matsubara等人提出的Entropic Student（ES），因为我们遵循相同的实时推理目标与特征压缩。然而，他们简单地应用了学习目标和来自先前工作的缩小版自动编码器。
3. 神经数据压缩的理论基础
3.1 率失真理论的扩展
根据香农的率失真理论，我们寻求一个受失真约束限制的从随机变量$X$到随机变量$U$的映射，以最小化$X$结果的比特率。更正式地说，给定失真度量$D$和失真约束$D_c$，最小比特率为：
$$\min_{P_{U|X}} I(X;U) \quad \text{s.t.} \quad D(X, U) \leq D_c$$
其中$I(X;U)$是互信息，定义为：
$$I(X;U) = \int\int p(x, u) \log \frac{p(x, u)}{p(x)p(u)} dxdu$$
在有损图像压缩中，$U$通常是原始输入$X$的重建$\tilde{X}$，失真度量是某种平方误差和$d(x, \tilde{x})$。由于率失真理论不限制我们进行图像重建，我们可以应用与机器对机器通信相关的失真度量。
为了详细说明丢弃信息对判别任务的潜力，考虑数据处理不等式（DPI）。对于形成马尔可夫链$X \leftrightarrow Y \leftrightarrow Z$的任何3个随机变量$X, Y, Z$，以下成立：
$$I(X; Y) \geq I(X; Z)$$
然后，通过将分层神经网络视为连续表示的马尔可夫链，描述$n$层顺序DNN中的信息流：
$$I(X; Y) \geq I(R_1; Y) \geq I(R_2; Y) \geq \ldots I(R_n; Y) \geq I(\tilde{Y}; Y)$$
3.2 信息瓶颈原理的应用
当任务是从联合分布$P_{X,Y}$预测真实标签$Y$时，率失真目标本质上由信息瓶颈原理给出。通过拉格朗日乘数松弛约束，目标是最大化：
$$I(Z; Y) - \beta I(Z; X)$$
具体来说，编码$Z$应该是$X$相对于$Y$的最小充分统计量。对于任意分类任务，$I(Z; Y)$的实用近似是条件交叉熵：
$$D = H(P_Y, P_{\tilde{Y}|Z})$$
4. 解决方案方法
4.1 端到端优化的损失函数
我们遵循非线性变换编码（NTC）来实现神经压缩算法。具体来说，我们嵌入一个与熵模型联合优化的随机压缩模型。
对于图像向量$x$，我们有一个参数化分析变换$g_a(x; \phi_g)$，将$x$映射到潜在向量$z$。然后，量化器$Q$将$z$离散化为$\bar{z}$，使得熵编码器可以使用熵模型将$\bar{z}$无损压缩为比特序列。
类似于变分推断，我们用参数化变分密度$q(\tilde{z}|x)$近似难处理的后验$p(\tilde{z}|x)$：
$$\mathbb{E}_{x \sim p_x} D_{KL}[q||p_{\tilde{z}|x}] = \mathbb{E}_{x \sim p_x} \mathbb{E}_{\tilde{z} \sim q}\left[-\log p(x|\tilde{z}) - \log p(\tilde{z})\right]$$
通过假设高斯分布，使得失真项的似然由下式给出：
$$P_{x|\tilde{z}}(x | z, \tilde{\theta}_g) = \mathcal{N}(x | g_s(\tilde{z}; \theta_g), 1)$$
我们可以使用$h$和$\tilde{h}$之间的平方差和作为我们的失真损失。
率项描述了压缩$\tilde{z}$的成本。类似于LIC方法，我们应用均匀量化$\bar{z} = \lfloor\tilde{z}\rceil$。由于离散化导致梯度流问题，我们通过添加均匀噪声$\eta \sim U(-\frac{1}{2}, \frac{1}{2})$来应用连续松弛。结合率和失真项，我们导出估计目标的损失函数：
$$\mathcal{L} = \|P_h(x) - g_s(g_a(x; \phi_g)+\eta; \theta_g)\|^2_2 + \beta \log p(g_a(x; \theta_g) + \eta)$$
4.2 显著性引导失真
使用HD作为失真项的次优性源于将$H$中的每个像素视为同等重要。这意味着方程中的MSE过度严格地惩罚包含冗余信息的空间位置的像素，这些信息后面的层可以安全丢弃。相反，当$\tilde{h}$在数值上接近$h$时，损失可能没有足够惩罚显著像素。
因此，我们可以通过引入额外信号来改进损失，这些信号正则化次优失真项。对于每个样本，我们需要一个向量$\mathbf{S}$，其中每个$s_i \in \mathbf{S}$是关于剩余尾层条件概率分布的空间位置显著性的权重项。然后，我们应该能够通过正则化失真项来改善率失真性能：
$$\mathcal{L}_{distortion} = \gamma_1 \cdot \mathcal{L}_1 + \gamma_2 \cdot s_i \cdot \frac{1}{N}\sum_i (h_i - \tilde{h}_i)^2$$
其中$\mathcal{L}_1$是原始失真项，$\gamma_1, \gamma_2$是和为1的非负实数。
5. 网络架构设计
5.1 模型分类法
图9：简单分类法与最小示例
该图展示了我们方法的三层分类结构。顶层Archtype反映模型的主要归纳偏置（如CNNs与Vision Transformers）。架构族描述变体（例如ResNet系列包括ResNet、Wide ResNet、ResNeXt等）。直接相关指相同架构的不同大小（如Swin-T、Swin-S、Swin-B等）。
5.2 瓶颈位置确定
考虑大多数现代DNN由初始嵌入后跟几个阶段组成。在直接相关的架构中，各个组件是相同的。变体之间的差异主要是嵌入维度或最深阶段的块比率。例如，ResNet-50的块比率是3:4:6:3，而ResNet-101的块比率是3:4:23:3。
5.3 解码器蓝图
区分原型的关键特征是基本构建块引入的归纳偏置（例如，卷积层与注意力层）。为了考虑非相关架构之间的不同表示，我们不应该通过直接重新利用神经压缩模型来忽视架构引入的偏置。
图10：FrankenSplit的参考实现
该图展示了训练后应用于两个变体的两个蓝图的FrankenSplit参考实现。灰色区域内的框是单独的实例（即只有一个编码器），具有相同名称的框共享架构。外部的圆角框组织从粗粒度到细粒度的层视图。
6. 实验评估
6.1 率失真性能
图13：ImageNet的率失真曲线
该图显示了使用Swin-B主干的率失真曲线。FrankenSplit-FP（FS-FP）和FrankenSplit-SGFP（FS-SGFP）的架构相同，但FS-SGFP使用显著性引导训练。结果显示FS-SGFP通过简单地用显著性图引导失真损失，在不影响预测强度或增加额外运行时开销的情况下实现了25%更低的比特率。
与Entropic Student（ES）相比，即使没有显著性引导，FS-FP也始终以很大的优势超过ES。具体来说，对于无损配置，FS-FP和FS-SGFP分别实现了32%和63%更低的比特率。
6.2 与图像编解码器的比较
图14：多个主干上基线的预测损失
该图显示了不同编解码器在多个主干架构上的预测损失。使用PSNR而不是bpp作为横轴，展示了PSNR不适合作为下游任务的度量。压缩模型积极丢弃信息，预测损失在主干之间是可比的，但差异可达3-5%。
6.3 蓝图泛化到任意主干
图15：各种主干的率失真曲线
该图展示了三个架构族（Swin、ResNet和ConvNeXt）的率失真性能。直接相关的架构共享相应的蓝图但具有单独训练的压缩器。跨所有架构族，我们观察到相似的率失真性能。
6.4 单编码器支持多个主干
图16：使用解码器恢复准确性的迭代次数
该图显示了通过微调解码器参数恢复准确性所需的迭代次数。实验证明了单个编码器可以通过适当的解码器蓝图服务多个主干架构。
6.5 张量维度对率失真性能的影响
图18：比较对大小的影响
左图显示增加编码器容量自然导致更低的比特率而不增加预测损失。右图显示增加输出通道会导致更高维的潜在张量（$C \times 28 \times 28$），但与压缩文件大小成反比关系。
6.6 直接优化对SVBI的局限性
图19：对比率失真性能
该图展示了在CIFAR-10和ImageNet数据集上应用SVBI-CE和SVBI-KD目标的结果。在CIFAR-10上，SVBI-CE和SVBI-KD比JPEG产生适度的性能增益。然而，它们在ImageNet上的表现明显更差，证明了直接优化对浅层瓶颈的不适用性。
7. 预测延迟和开销
7.1 计算开销
表VI：推理管道组件执行时间
该表总结了预测管道组件的执行时间。Lu-JHAP展示了没有顺序上下文组件的LIC模型明显更快，但仍比FrankenSplit慢9.3x-9.6x，尽管具有明显更差的率失真性能。
7.2 与卸载的竞争
图20：比较对大小的影响
该图绘制了使用NX客户端对手工编解码器的推理延迟。对于更强的连接（如4G LTE），它仍比使用PNG快3.3倍。在高度受限的网络（如BLE）中，FrankenSplit的速度提高达16倍。
表VII：各种无线标准的总延迟
该表总结了我们的方法在各种标准中的表现。一般来说，网络越受限，我们从减少传输大小中获益越多。
8. 结论
本工作引入了一个新颖的轻量级压缩框架，以促进依赖大型DNN的关键MEC应用。我们展示了我们设计启发式的最小实现足以超越众多基线。主要贡献包括：
- 彻底探索了浅层和深层瓶颈注入在特征压缩方面的差异
- 引入了一种新颖的显著性引导训练方法，克服了训练容量有限的轻量级编码器的挑战
- 引入了一种可泛化的设计启发式，用于将变分特征压缩模型嵌入到任意DNN架构中
尽管取得了显著改进，但仍存在几个局限性。我们强调报告结果的主要见解是充分的失真度量和正则化方法对神经特征压缩的潜力。
附录：一些推导
A. 变分推断与率失真目标的联系
从贝叶斯角度，我们可以将压缩问题视为寻找潜在表示$z$的后验分布$p(z|x)$。由于真实后验难以处理，我们使用变分推断近似它。
定义变分下界（ELBO）：
$$\log p(x) \geq \mathbb{E}_{q(z|x)}[\log p(x|z)] - D_{KL}[q(z|x)||p(z)]$$
重新排列并期望化：
$$\mathbb{E}_{x \sim p_x}[\log p(x)] \geq \mathbb{E}_{x \sim p_x}\mathbb{E}_{q(z|x)}[\log p(x|z)] - \mathbb{E}_{x \sim p_x}[D_{KL}[q(z|x)||p(z)]]$$
第一项是重建项（负失真），第二项是率项。这直接对应于率失真目标。
B. 信息瓶颈的详细推导
从联合分布$p(x,y)$开始，我们寻找压缩表示$z$。信息瓶颈目标是：
$$\max_{p(z|x)} I(Z;Y) - \beta I(Z;X)$$
使用互信息的定义：
$$I(Z;Y) = \int\int p(z,y) \log \frac{p(z,y)}{p(z)p(y)} dzdy$$
$$I(Z;X) = \int\int p(z,x) \log \frac{p(z,x)}{p(z)p(x)} dzdx$$
对于分类任务，我们可以使用：
$$I(Z;Y) \approx \mathbb{E}_{p(x,y)}\mathbb{E}_{q(z|x)}[\log p(y|z)]$$
这导致实际损失函数：
$$\mathcal{L}_{IB} = -\mathbb{E}_{p(x,y)}\mathbb{E}_{q(z|x)}[\log p(y|z)] + \beta \mathbb{E}_{p(x)}[D_{KL}[q(z|x)||p(z)]]$$
C. 头部蒸馏的信息论分析
考虑信息路径$X \to H \to Z \to \tilde{H} \to \tilde{Y}$，其中$H = P_h(X)$是原始头部输出。
根据数据处理不等式：
$$I(X;Y) \geq I(H;Y) \geq I(Z;Y) \geq I(\tilde{H};Y) \geq I(\tilde{Y};Y)$$
头部蒸馏优化：
$$\max_{q(z|h)} I(Z;H) - \beta I(Z;X)$$
由于$H$是$X$的确定性函数，我们有：
$$I(Z;X) = I(Z;H) + I(Z;X|H)$$
其中$I(Z;X|H) \geq 0$。因此：
$$I(Z;H) - \beta I(Z;X) = I(Z;H)(1-\beta) - \beta I(Z;X|H)$$
当$\beta < 1$时，我们最大化$I(Z;H)$并最小化$I(Z;X|H)$，这鼓励$Z$成为$H$的充分统计量。
D. 显著性引导失真的理论基础
设$s_i$为位置$i$的显著性权重，定义为：
$$s_i = \frac{\partial I(\tilde{Y};Y)}{\partial h_i} \Bigg|_{h_i = P_h(x)_i}$$
这表示位置$i$的特征对最终预测的贡献。使用链式法则：
$$\frac{\partial I(\tilde{Y};Y)}{\partial h_i} = \sum_j \frac{\partial I(\tilde{Y};Y)}{\partial \tilde{y}_j} \cdot \frac{\partial \tilde{y}_j}{\partial h_i}$$
在实践中，我们使用Grad-CAM近似这个量：
$$s_i \approx \text{ReLU}\left(\sum_k \alpha_k A^k_i\right)$$
其中$\alpha_k = \frac{1}{Z}\sum_i \sumj \frac{\partial y^c}{\partial A^k{ij}}$，$A^k$是第$k$个激活图。
E. 量化与熵模型
对于潜在变量$z$，我们应用均匀量化：
$$\bar{z} = \text{round}(z) = \lfloor z + 0.5 \rfloor$$
为了保持可微性，我们在训练期间使用加性均匀噪声：
$$\tilde{z} = z + \eta, \quad \eta \sim U(-0.5, 0.5)$$
熵模型假设因子化先验：
$$p(\tilde{z}) = \prod_i p(\tilde{z}_i)$$
其中每个$p(\tilde{z}_i)$是学习的单变量分布。率估计为：
$$R = \mathbb{E}_{x}\left[\sum_i -\log p(\tilde{z}_i)\right]$$
使用学习的累积分布函数（CDF）：
$$p(\bar{z}_i = k) = F(k + 0.5) - F(k - 0.5)$$
其中$F$是参数化CDF。
F. 资源不对称下的优化
考虑客户端计算能力$C{client}$和服务器计算能力$C{server}$，其中$C{client} \ll C{server}$。
总延迟：
$$T_{total} = T_{encode}(C_{client}) + T_{transfer}(R) + T_{decode}(C_{server}) + T_{predict}(C_{server})$$
我们的目标是：
$$\min_{g_a, g_s} T_{total} \quad \text{s.t.} \quad \text{Accuracy} \geq \text{Threshold}$$
这导致约束优化问题：
$$\mathcal{L}_{total} = \mathcal{L}_{distortion} + \beta \mathcal{L}_{rate} + \lambda \max(0, T_{encode} - T_{max})$$
其中最后一项确保编码时间保持在可接受的界限内。
