# KeyNet: An Asymmetric Key-Style Framework for Watermarking Deep Learning Models
## 4 The _KeyNet_ Framework
### Problem Formulation  
The key idea of our framework is to perform two tasks at the same time: the original classification task \(T_{org}\) and the watermarking task \(T_{num}\). To do so, _KeyNet_ leverages the multi-task learning (MTL) approach to achieve high accuracy in both tasks by sharing the parameters of the original model between the two tasks. _KeyNet_ adds a private model to the original model. The original model's objective is to correctly classify the original data samples into their corresponding labels, while the private model's objective is to correctly predict the position of the owner's signature in a sample of the WM carrier set using the original model predictions. We can formally represent as follows the problem being tackled:  
* **Representation of the original, private, and combined models. ** Let \(D_{org}=\{(x_{i},y_{i})\}_{i=1}^{n}\) be the original task data and \(D_{wm}=\{c_{j}\}_{j=1}^{m}\) be the WM carrier set data. Let \(h\) be the function of the original model and \(f\) be the function of the private model. Let \(\theta_{1}\) the parameters of \(h\) and \(\theta_{2}\) be the parameters of \(f\). Let \(signature\) be the owner's signature and \(PL=\{(p_{k},l_{k})\}_{k=1}^{z}\) be the set of predefined position-label pairs (e.g., position: top left, label: 1 and position: bottom right, label: 4) where \(z\) is the total number of positions at which the signature can be located on a WM carrier set sample. Let \(sign\) be the function that puts a \(signature\) on a carrier set sample \(c\) and returns the signed sample \(c^{p_{k}}\) and its corresponding label \(l_{k}\) as: \[(c^{p_{k}},l_{k})=sign(signature,c,p_{k}).\] Let \(D_{wm}^{signed}\) be the signed carrier set samples that contain all the \((c^{p},l)\) pairs. We use \(D_{org}\) and \(D_{wm}^{signed}\) to train both \(h\) and \(f\) to perform \(T_{org}\) and \(T_{wm}\). Typically, the function \(h\) tries to map each \(x_{i}\in D_{org}\) to its corresponding \(y_{i}\), that is, \(h(x_{i})=y_{i}\). Let \(f(h)\) be the composite function that aims at mapping each \(c_{j}^{p_{k}}\) to its corresponding \(l_{k}\), that is, \(f(h(c^{p_{k}}))=l_{k}\). * **Embedding phase. ** We formulate the embedding phase as an MTL problem where we jointly learn two tasks that share some parameters in the early layers and then have task-specific parameters in the later layers. The shared parameters in our case are \(\theta_{1}\), while \(\theta_{2}\) are the WM task-specific parameters. We compute the weighted combined loss \(L\) as \[L=\alpha Loss(h(x),y)+(1-\alpha)Loss(f(h(c^{p})),l),\] where \(h(x)\) represents the predictions on the original task samples, \(f(h(c^{p})\) represents the predictions on \(D_{wm}^{signed}\) samples, \(Loss(h(x),y))\) is the loss function that penalizesthe difference between the original model outputs \(h(x)\) and the original data targets \(y\), \(Loss(f(h(c^{p})),l)\) is the loss function that penalizes the difference between the composite model outputs \(f(h(c^{p}))\) and the signed WM carrier set's target \(l\), and \(\alpha\) is the combined weighted loss parameter. Then, we seek \(\theta_{1}\) and \(\theta_{2}\) that make \(L\) small enough to get acceptable accuracy on both \(T_{org}\) and \(T_{um}\). Once this is done, the WM has been successfully embedded while preserving the accuracy of the original task \(T_{org}\). * **Verification phase. ** The verification function \(V\) checks whether a claimer (also known as the _prover_), who has delivered her \(signature\) and WM carrier set \(D_{um}\) to the authority (also known as the _verifier_), is the legitimate owner of a remote model \(h^{\prime}\). If the prover is the legitimate owner of \(h^{\prime}\), he/she will be able to pass the verification process and thus prove her ownership of \(h^{\prime}\). That is, because he/she possesses the private model \(f\), which was trained to decode \(h^{\prime}\) predictions on her signed \(D_{um}\).