# TraceGuard: Fine-Tuning Pre-Trained Model by Using Stego Images to Trace Its User
## 3 Proposed Approach
### Training Steganography Network  
As shown in Figure 3, encoder \(E\), decoder \(D\), and pre-trained model \(P\) participate in the training of the steganography network. \(D_{s}=\{I_{co}^{i}\}_{i=1}^{N}\) denotes the image set for training the steganography network, where \(I_{co}^{i}\in\mathbb{R}^{C\times H\times W}\), and \(C\), \(H\), \(W\) denote the channel number, height, and width of the image, respectively. \(K_{em}\in\{0,1\}^{L}\) is the binary secret key of length \(L\). The encoder \(E\) learns to embed \(K_{em}\) into \(I_{co}\) and obtain the stego image \(I_{st}\). Next, \(I_{st}\) is fed into the pre-trained model to obtain the embedding. Finally, the decoder \(D\) learns to extract the secret key \(K_{ex}\) from the embedding. The structure of \(D\) is composed of two fully connected layers and an activation function ReLU, which can compress the embedding and transform it into \(K_{ex}\) with the same length as \(K_{em}\), in which each element is normalized between 0 and 1 using the sigmoid function.  
The loss function consists of two items: the secret key loss and the image loss. To evaluate the extraction error of the secret key, the loss is defined. To constrain the visual similarity between \(I_{st}\) and \(I_{co}\), the loss is defined. During training, the steganography network P is frozen. Then, E and D are alternately learned by using the combination of the two loss items:  
\[\begin{split}\mathcal{L}_{ste}&=\lambda_{k} \mathcal{L}_{k}(K_{ent},K_{ex})+\lambda_{I}\mathcal{L}_{I}(I_{co},I_{st})\\ \mathcal{L}_{k}(K_{em},K_{ex})&=\text{BCE}(K_{em}, K_{ex})\triangleq\|K_{em}-K_{ex}\|_{2}^{2}/L\\ \mathcal{L}_{I}(I_{co},I_{st})&=\text{MSE}(I_{co}, I_{st})\triangleq\|I_{co}-I_{st}\|_{2}^{2}/CHW\end{split} \tag{1}\]  
where \(\lambda_{k}\) and \(\lambda_{I}\) are weights to control the balance between the two loss items.  
The procedure for training the steganography network is shown in Algorithm 1.  
\begin{table}
\begin{tabular}{c c} \hline
**Notation** & **Definition** \\ \hline \(E\), \(D\) & encoder and decoder of the steganography networks \\ \hline \(P\), \(P_{sd}\), & pre-trained model/shadow pre-trained model/ \\ \(P_{f}\), \(P_{sp}\) & fingerprinted pre-trained model/suspect model \\ \hline \(K_{em}\), \(K_{ex}\) & embedded key/extracted key \\ \hline \(I_{co}\), \(I_{st}\) & cover images/stego images \\ \hline \(D_{s}\) & dataset for training steganography network/ \\ \(D_{f}\), \(D_{t}\) & dataset for fingerprinting/tracing pre-trained model \\ \hline \end{tabular}
\end{table}
Table 2: Definitions of used notation.  
Figure 3: Flowchart of training steganography network: encoder (E), decoder (D), pre-trained model (P), and loss function.  
```
0: base pre-trained model \(P\), \(D_{s}\), secret key set \(S_{k}\), training epochs \(N_{t}\), learning rate \(\eta\) Output: encoder (E), decoder (D)
1:for\(n=1\) in \(N_{t}\)do
2:while\(\{I_{co}\}^{i}\
eq NULL\)do
3: Select \(\{I_{co}\}^{i}\in D_{s}\), \(K^{i}_{em}\in S_{k}\) /* \(\{I_{co}\}^{i}\) is a batch of samples */
4:\(\{I_{st}\}^{i}=E_{\theta}\big{(}\{I_{co}\}^{i},K^{i}_{em}\big{)}\)\(\triangleright\) embed key
5:\(\{y\}^{i}=P\big{(}\{I_{st}\}^{i}\big{)}\)\(\triangleright\) obtain embedding from stego image
6:\(\{K_{ext}\}^{i}=D_{\theta}\big{(}\{y\}^{i}\big{)}\)\(\triangleright\) extract key
7:\(\mathcal{L}=\lambda_{k}\text{BCE}\big{(}\{K_{ext}\}^{i},K^{i}_{em}\big{)}+ \lambda_{l}\text{MSE}\big{(}\{I_{co}\}^{i},\{I_{st}\}^{i}\big{)}\)
8:\(\theta_{i+1}=\theta_{i}-\eta\bigtriangledown_{\theta}\mathcal{L}\)\(\triangleright\) update the parameters of encoder and decoder
9:endwhile
10:endfor
```  
**Algorithm 1** Training Steganography Network