Shiki

4.1 Overview

如图5所示，CausalCode以模型 $M$ 和训练集 $D_{1}$ 作为输入，输出增强后的模型 $M^{+}$。与 $M$ 相比，$M^{+}$ 展现出针对对抗攻击的增强鲁棒性。这一转化通过两个关键步骤实现：
（i）利用因果干预扩展训练数据；
（ii）从包含原始训练集 $D_{1}$ 和干预样本的数据集中学习不变表征。为使CausalCode有效运行，需要解决两个核心技术问题：

• CausalCode如何生成干预样本？这将在CausalCode的步骤1（第4.2节）中解答。
• CausalCode如何训练能够学习不变表征 $\Phi$ 的鲁棒模型？这将在CausalCode的步骤2（第4.3节）中解答。

![[file-20250914184557347.png]]

图5. CausalCode总体框架，包含扩展训练数据集（步骤1）和通过因果特征学习不变表征（步骤2）两个阶段。矩阵中，具有相同因果特征的样本用相同形状表示，颜色则代表不同类型的伪特征。总体而言，CausalCode通过提升因果特征的重要性，使具有相同功能的样本在学习的表征空间中相互靠近，同时降低伪特征的重要性。

4.2 步骤一：扩展训练数据

在第3.3节中，我们已经阐述了因果数据增强方法的推导过程。本节将描述如何运用该增强方法扩展训练数据。算法1与图4(c)展示了如何从给定训练集$D_{1}={{x_{1},x_{2},\ldots,x_{n}}}$生成干预样本，其中输入参数包括：最大迭代次数$iter_{max}$、包含$m$个元素的替换候选集$\mathcal{R}$，以及待增强的非鲁棒模型$M$。输出结果为干预样本构成的新训练集$D_{new}$。

该算法对每个样本$x_{i}\in D_{1}$（$1\leq i\leq n$）进行迭代处理。首先随机确定介于1和$iter_{max}$之间的停止点$iter_{stop}$（对应图4(c)中箭头数量）（第4行），以覆盖所有可能的干预情况。对于每个样本$x_{i}$，算法识别可进行迭代变换而不改变代码语义的潜在干预点$I$（第5行）。将$x_{i}^{new}$初始化为原始样本并设置迭代计数器为零后，算法进入循环流程，在此期间对候选集$\mathcal{R}$在$I$处施加干预以生成变换样本集$\mathcal{S}$（第7-14行）。

在每次迭代中，选取与模型梯度下降方向最契合的top-$k$样本（第10行）。契合度衡量基于样本嵌入向量与干预样本嵌入向量的相似度分数，计算公式为：

$$SC(\mathbf{x}{i}^{s{j}},\mathbf{x}{i})=\frac{v(\mathbf{x}{i}^{s_{j}})-v( \mathbf{x}{i})}{|v(\mathbf{x}{i}^{s_{j}})-v(\mathbf{x}{i})|{2}}\frac{ \partial L(y,M(\mathbf{x}{i}))}{\partial v(\mathbf{x}{i})}(4)$$

其中$v(\mathbf{x}{i}^{s{j}})$和$v(\mathbf{x}{i})$分别为$\mathbf{x}{i}^{s_{j}}$（即$\mathcal{S}$中一个样本）和$\mathbf{x}{i}$的嵌入向量，$M(\mathbf{x}{i})$是模型对$\mathbf{x}{i}$的预测结果，$\frac{\partial L(y,M(\mathbf{x}{i}))}{\partial v(\mathbf{x}{i})}$是$\mathbf{x}{i}$的嵌入层级梯度向量。该算法无需为每个候选样本计算梯度，仅需为原始标识符计算一次梯度，因此总体开销与生成对抗样本相当。

从这些顶级样本中随机选取一个作为新干预样本，对应图4(c)中梯度下降方向（红色区域）候选集的随机选择。当处理完$D_{1}$中所有样本后算法终止。最终，算法输出干预样本以构成扩展训练集$D_{new}$（第17行）。

4.2 步骤二：学习不变特征

在通过因果数据增强扩展训练数据后（步骤1），我们现专注于学习对伪特征具有鲁棒性的不变表示。如第3.3节所述，我们的目标是学习一个与虚假特征 $F$ 独立的特征表示 $\Phi(\mathbf{x})$ 。为实现这一目标，我们强制要求原始样本 $\mathbf{x}{i}$ 与其对应干预样本 $\mathbf{x}^{\prime}{i}$（二者共享相同语义特征 $X_{S}$）的特征表示 $\Phi(\mathbf{x})$ 在表示空间中距离为零。该条件可形式化表示为：

$$\sum_{j_{1}\neq j_{2}}\text{dist}\left(\Phi\left(\mathbf{x}^{(j_{1})}{i}\right) ,\Phi\left(\mathbf{x}^{(j{2})}_{i}\right)\right)=0$$ (6)

其中$\mathbf{x}^{(j_{1})}{i}\in D{j_{1}}$和$\mathbf{x}^{(j_{2})}{i}\in D{j_{2}}$（$j_{1}\neq j_{2}$）表示共享相同程序语义的样本对，$\Phi$为特征提取函数，$\text{dist}$为距离度量（如$l_{2}$距离）。

此外，我们引入$\mathcal{L}{CausalCode}$作为_正则化项，以帮助最小化原始样本与其干预样本之间的差异，即：

$$\mathcal{L}{Causal}=\sum{1\leq i\leq n;1\leq j_{1},j_{2}\leq v;j_{1}\neq j_{2 }}\text{dist}\left(\Phi\left(\mathbf{x}^{(j_{1})}{i}\right),\Phi\left(\mathbf{x} ^{(j{2})}_{i}\right)\right)$$ (7)

最小化$\mathcal{L}{Causal}$可促使模型学习捕获代码语义的因果特征。但仅最小化$\mathcal{L}{Causal}$可能不足，因为处理空代码时会导致$\mathcal{L}{Causal}=0$。为解决此问题，我们引入额外条件以确保$X{S}$能有效传递语义及$Y$的信息。为此，我们加入以下标准交叉熵损失函数$\mathcal{L}_{Task}$：

$$\mathcal{L}{Task}=\sum{1<=i<=n;1\leq j\leq v;\quad I\left(G\left(\Phi\left( \mathbf{x}^{(j)}{i}\right)\right),y^{(j)}{i}\right)}$$ (8)

其中$I$为适当损失函数（如交叉熵损失），$G$为模型，$\Phi$为特征提取器，$\mathbf{x}^{(j)}{i}$为第$i$个样本的特征向量，$y^{(j)}{i}$为第$i$个样本的标签，$n$为数据集中样本总数。

为优化总体目标，我们将 $\mathcal{L}{Causal}$ 和 $\mathcal{L}{Task}$ 结合，得到损失函数 $\mathcal{L}_{CausalCode}$：

$$\mathcal{L}{CausalCode}=\lambda\mathcal{L}{Causal}+\mathcal{L}_{Task}$$ (9)

其中 $\lambda$ 作为超参数。其具体取值根据代码任务和所用模型确定。通过在学习过程中最小化 $\mathcal{L}_{CausalCode}$，CausalCode 实现了图5强调的目标，确保具有相同语义的样本在表示空间中的表征彼此接近。

算法2概述了支持CausalCode的模型训练流程。算法初始设置 $\lambda=0$（第3行），$N_{s}$ 取较小值（第5行），并采用非鲁棒模型 $M$ 进行数据特征的初始学习。训练开始时最小化交叉熵损失（第6行）。数据集 $\mathcal{D}$ 包含原始训练集和 $v-1$ 个干预样本数据集，以批次方式使用（第9行）。为优先学习因果特征 $X_{S}$ 并削弱伪特征 $X_{F}$ 的影响，随着CausalCode损失的降低，算法逐步增大因果不变损失的比例直至达到1，从而最小化公式(9)所述的CausalCode损失（第8-17行）。每经过 $r$ 个周期，会更新扩展后的训练集 $\mathcal{D}$（第19行）。当损失收敛时算法终止（第4行）。定期更新随机扰动训练集 使模型能够学习多样的代码变体，提取本质特征，并获得不变表征，最终得到一个鲁棒模型 M*。