# KeyNet: An Asymmetric Key-Style Framework for Watermarking Deep Learning Models
## 4 The _KeyNet_ Framework
### Watermark Embedding  
To successfully embed the WM in the original model \(h\) without compromising the accuracy of the original task, we jointly train both \(h\) and the private models \(f\) simultaneously. As a large amount of the carrier set samples have been signed, we first randomly select one-fifth of \(D_{wm}^{signed}\) for training. The random selection allows for the representation of all the possible states while reducing the carrier set size. The signed samples from other distributions \(D_{dif}^{signed}\) are combined with those randomly selected and assigned to \(D_{2}\). In the end, we add \(D_{2}\) to the original task data \(D_{org}\). The resulting combined data set \(D=D_{org}\cup D_{2}\) are used in the training step as specified next.  
During joint training, a batch \(b\) is taken from \(D\). Then, \(b\) is separated into two sub-batches: \(\{x,y\}\in D_{org},\{c,l\}\in D_{2}\). \(\{x\}\) is passed to the original model \(h\) and the loss \(L_{f}(h(x),y)\) is calculated. On the other hand, \(\{c\}\) is first passed to \(h\), and then the predictions of the original model \(h(c)\) are passed to the private model \(f\); the loss \(L_{f}(f(h(c)),l)\) is afterwards calculated. As we deal with two classification tasks, the cross-entropy loss function is used to calculate the loss for both tasks.  
We use the parameter \(\alpha\) to balance the weight of \(L_{f}(h(x),y)\) and \(L_{f}(f(h(c)),l)\) before we add them up in the joint loss \(L\). Parameter \(\alpha\) allows us to choose the best combination of the weighted loss that preserves the accuracy of \(T_{org}\) while embedding WM successfully. Then, the parameters of \(h,f\) are optimized to minimize \(L\).  
Reducing \(L_{f}(h(x),y)\) forces \(h\) to predict the correct class for \(x\), while reducing the watermarking task loss \(L_{f}(f(h(c)),l)\) forces the private model to distinguish the distribution of the WM carrier set \(D_{wm}\), and predict the location of the owner's signature in its samples. The original model, in addition to performing the original task, also executes the first part of the watermarking task by outputting the features needed to find the position of the signature. By using these features as input, the private model performs the second part of the watermarking task, which consists in identifying the signature position.  
Regarding the architecture of the private model \(f\), the number of inputs corresponds to the size of \(h\) predictions, whereas the number of outputs corresponds to the number of classes of the WM task \(z+1\). We also add at least one hidden layer in between. The hidden layer enriches the information coming from the original model before passing it to the output layer of the private model.  
Algorithm 2 summarizes the process of embedding the WM. It takes an unmarked model \(h\) that might be pretrained or be trained from scratch, the private model \(f\), the original data set \(D_{org}\), the signed WM carrier set \(D_{wm}^{signed}\), signed samples from other distributions \(D_{dif}^{signed}\), and the joint loss balancing parameter \(\alpha\). The output of the embedding phase is a marked model \(h^{*}\) along with its corresponding private model \(f\).  
```
0: Unmarked DL model \(h\), private model \(f\), original data set \(D_{org}\), signed WM carrier set \(D_{wm}^{signed}\), signed samples from other distributions \(D_{dif}^{signed}\), batch size \(BS\), weighted loss parameter \(\alpha\).
```
0: Marked model \(h^{*}\), corresponding private model \(f\).
1:\(s=size(D_{wm}^{signed})/5\)
2:\(D_{2}\gets randomSample(D_{2},s)\)
3:\(D_{2}\gets D_{2}\cup D_{dif}^{signed}\)
4:\(D\gets D_{org}\cup D_{2}\)
5:\(L_{f}=crossEntropy()\) //Loss function
6:for each batch \(b\) of size \(BS\) in \(D\)do
7:\(\{x,y\},\{c,l\}\leftarrowsplit(b)\), with \(\{x,y\}\in D_{org}\) and \(\{c,l\}\in D_{2}\)
8:\(L\leftarrow\alpha L_{f}(h(x),y)+(1-\alpha)L_{f}(f(h(c)),l)\)
9:\(optimize(L)\)
10:endfor return\(h^{*},f\) ```
**Algorithm 2** Watermark Embedding