Model Provenance via Model DNA

Consider a dataset $D_{s}=\{(x_{i},y_{i})\}_{i=1}^{{|D_{s}|}}$ where each $x_{i}\in\mathbb{R}^{d}$ is a data instance and $y_{i}\in Y=\{1,2,\ldots,c\}$ is its associated class label, a machine learning model $M_{s}$ is learned from $D_{s}$ by using an algorithm architecture $\mathcal{A}_{s}$. Let $D_{t}=\{(x_{j},y_{j})\}_{j=1}^{{|D_{t}|}}$ be a dataset used to learn a machine learning model $M_{t}$. If $M_{t}$ is learned using the pre-training model $M_{s}$ with the same algorithm architecture $\mathcal{A}_{s}$, we refer to $M_{s}$ as the source model and $M_{t}$ as the homologous model of $M_{s}$. Similarly, let $\bar{M}_{t}$ be a machine learning model learned from the dataset $D_{t}={(x_{j},y_{j})}_{j=1}^{{|D_{t}|}}$ but based on random initialization or pre-training without using $M_{s}$. We refer $\bar{M}_{t}$ as non-homologous to $M_{s}$.

The problem of Model Provenance (MP) is to find a function $f$ such that for a given source model $M_{s}$ and any target model $M_{t}$,

When $f(M_{s},M_{t})=1$, we say $M_{s}$ is the provenance of $M_{t}$.

Remark: The problem of Model Provenance (MP) can have various variations based on the conditions applied to the target model $M_{t}$. For example, there would be a lot of downstream models obtained by fine-tuning, distillation, pruning, etc. In this paper, our primary focus is on the MP scenario within the fine-tuning setting, specifically when $M_{t}$ shares the same structure as $M_{s}$, and its training solely depends on the pre-training model $M_{s}$. Our exploration of MP is initially focused on this easier condition, and we intend to explore more complex scenarios in future research.

In the field of continual learning, there have been several studies focused on quantifying the relationship between homologous models, such as the analysis of catastrophic forgetting [21, 11, 20]. For example, [21] provided a theoretical bound on forgetting in sequential learning. Let $L(w)$ be the loss on the training dataset with model parameter $w$, and $w_{1}$ and $w_{2}$ be the optimal or convergent parameters after training the source and homologous models sequentially. Specifically, they showed that

where $L_{1}(w_{1})$ and $L_{1}(w_{2})$, respectively, represent the losses on source training dataset with parameters $w_{1}$ and $w_{2}$, $\Delta w=w_{2}-w_{1}$, and $\lambda_{1}^{max}$ is the largest eigenvalue of the Hessian matrix $\nabla^{2}L_{1}(w_{1})$. The eigenvalues of the Hessian matrix are indicative of the curvature of the loss function [11], where smaller eigenvalues imply a flatter loss function. Thus, $\lambda_{1}^{max}$ is considered as a proxy for the flatness of the loss function (lower is flatter). Meanwhile, empirical studies conducted by [20] and [21] show that initializing models with pre-trained weights results in a relatively flat task minima, and flatter models tend to have smaller $\Delta w$ than less flat ones.

These existing works have motivated us to a straightforward solution for model provenance, i.e., using the difference of $\Delta w$ to distinguish between homologous and non-homologous models. We can expect that $\Delta w=w_{2}-w_{1}$ would be smaller for homologous models compared to non-homologous models based on the above analysis. However, we found that this solution did not work well in practice. The main issue was that it is impossible to know how flat the target model is, which makes it difficult to use the difference of $\Delta w$ as a reliable indicator of model provenance.

Furthermore, we conducted experiments using the standard continual learning setup on the image classification task [32, 24] to show the relationship between homologous models. Specifically, we examined two model architectures, ResNet18 [9] and AlexNet [15], on the CIFAR10 and CIFAR100 datasets. A source model $M_{s}$ is first trained on CIFAR10. Then we randomly select ten classes from CIFAR100 to train target models $M_{t}$. Note that $M_{t}$ is learned by initializing its parameters using those of the source model $M_{s}$. In Figure 1, we present the results of the experiments conducted on ResNet18, and the results for AlexNet can be found in [mu2023model]. The background color indicates the different training phases. The red color represents the training of source model $M_{s}$ on the CIFAR10 dataset, while the blue color represents the training of target model $M_{t}$ on the CIFAR100 dataset. The red line corresponds to the accuracy of testing the model on the training data of CIFAR10, while the blue line represents the accuracy of testing the model on the training data of CIFAR100. In the blue background part, we train the model based on the parameters of the model obtained in the red background part.

Figure 1 (a) demonstrates that the predictive accuracy of $M_{t}$ is significantly reduced when evaluated on the CIFAR10 dataset. The red section represents a model (called source model) trained on the source domain (CIFAR10), while the blue section represents this model (called target model) continuously trained on the target domain (CIFAR100). The red line indicates model evaluation on CIFAR10, and the blue line represents model evaluation on CIFAR100. This observation suggests that the training data of the source model may have been severely forgotten in the homologous model (i.e., large $||\Delta w||$). Figures 1 (b)-(d) provide further insight into the relationship between the source model $M_{s}$ and the homologous model $M_{t}$. They show the predictive accuracy of $M_{t}$ on the CIFAR10 dataset when different numbers of the last layers of $M_{t}$ are replaced with corresponding layers in $M_{s}$. In Figure 1 (b), within the blue section, we replace the model’s last layer with the model from the red section. Consequently, the model exhibits reduced performance when evaluated on CIFAR10 (red line). Similarly, in Figures 1 (c) and (d), we extended this by replacing the model’s last two/three layers with the ones from the red section, respectively. We can observe that as we replace more layers in $M_{t}$ with those from $M_{s}$, the predictive accuracy of $M_{t}$ on the CIFAR10 dataset increases. In Figure 1 (e), we consider a scenario where in the blue section, the (target) model is not initialized by the (source) model from the red section. Instead, the (target) model is trained with random initialization. In this setting, we replace the model’s last three layers with the corresponding layers from the (source) model in the red section (as depicted in Figure 1 (a)). Figure 1 (e) shows that even with the last layers replaced, the predictive accuracy of $\bar{M}_{t}$ is still poor. The green line in Figure 1 (e) represents the training accuracy of $\bar{M}_{t}$. Figures 1 (b) and (e) illustrate the presence of a relationship between the source model and its homologous model. However, non-homologous models, do not exhibit any discernible relationship even by replacing the layers between them.

Overall, our discussion has demonstrated the existence of a relationship between the source model and its homologous model, as shown by the improvement in predictive accuracy when replacing the last layers of the target model with those of the source model. These results led us to pose the question: “Can we establish a relationship between a source model and its homologous or non-homologous model based on the training data of the source model?” and motivated us to introduce the “Model DNA” and Model Provenance framework.

In the field of biology, deoxyribonucleic acid (DNA) is known as the molecule responsible for carrying genetic information essential for the growth and operation of organisms [25]. In the context of Machine Learning (ML), we introduce the concept of model DNA as a form of representation for an ML model. Drawing parallels to prior research in representation learning [2], this model representation aids in identifying differences among various model iterations across diverse tasks. Furthermore, it facilitates comparisons and assessments of similarity between different models.

In the domain of machine learning, a typical process involves training an ML model using a dataset $D$ and an algorithm architecture $\mathcal{A}_{s}$. With this understanding, our approach to DNA generation factors in both the impact of the dataset $D$ and the model’s input-output relationship. Thus, we present the model DNA definition:

Through the conceptualization of model DNA, we create a latent space of model representations encompassing the diverse DNA of various ML models. In this space, we can quantify the relationships between different ML models. Here, we assume that a DNA fragment $o_{i}$ is a vector by $o_{i}\leftarrow g(x_{i},M)$, where $x_{i}\in D$. It is important to note that we generally assume that the DNA of homologous models will be positioned closer, while the DNA of non-homologous models will be relatively distant from each other. This latent space leads to several key properties of model DNA, as described below.

Let $\mathbb{O}_{i}$ be the model DNA of any model $M_{i}$ and $\mathbb{D}$ represent the distance function such that $\mathbb{D}(o^{s}_{i},o^{t}_{j})$ is the distance between DNA fragments of any two models $M_{s}$ and $M_{t}$. The following properties of DNA latent space are desired:

• •

  Different ML models have different DNA: If $M_{s}\neq M_{t}$, $\mathbb{O}_{s}\neq\mathbb{O}_{t}$.

• •

  The homologous models should have similar DNA fragments and vice versa: $\mathbb{D}(o^{s}_{i},o^{t}_{i})<\mathbb{D}(o^{s}_{i},\bar{o}^{t}_{i})$, where $o^{s}_{i}\leftarrow g(x_{i},M_{s})$, $o^{t}_{i}\leftarrow g(x_{i},M_{t})$ and $\bar{o}^{t}_{i}\leftarrow g(x_{i},\bar{M}_{t})$. $M_{t}$ is trained based on the pre-training model $M_{s}$, whereas $\bar{M}_{t}$ is trained without using $M_{s}$.

• •

  The DNA fragments from the same ML model should be similar: $\forall i,j\ ,i\neq j$, $\mathbb{D}(o_{i},o_{j})$ should be minimized for $o_{i},o_{j}\in\mathbb{O}$.

In the following, we present our framework aimed at tackling the MP problem. This framework has been meticulously crafted to adhere to the aforementioned model DNA properties. Our approach takes into account not only the training data of the source model but also the model’s input-output to produce a comprehensive model DNA representation.

We introduce a framework termed Model DNA Generation and Model Provenance identification (MGMP), which operates on the principles of deep representation learning. The MGMP framework comprises three primary components: DNA generation, DNA similarity loss, and provenance classifier, as depicted in Figure 2. The process commences with the source model’s training data, which is fed into the DNA generator alongside the source and target models. This concatenation produces outputs from the DNA generator, source model, and target model, which are then amalgamated to generate the model DNA, exemplified by the integration of DNA generator outcomes and model predictions. Subsequently, we formulate the model DNA for both the source and target models. To preserve information relevant to homologous and non-homologous models, a DNA similarity loss is incorporated. Finally, we exploit the model DNA to train a binary prediction network, allowing us to infer provenance outcomes. Each component of the MGMP is described in subsequent sections.

All experiments were conducted using Python programming language on a machine equipped with an Intel Core CPU, 64 GB of memory, and an NVIDIA Corporation GV100GL [Tesla V100 SXM2 32GB] GPU. We show the details of our proposed framework MGMP in Table 3. In the Computer Vision (CV) task, the generator’s structure is set using the ResNet50 architecture. ResNet50 is a popular Convolutional Neural Network (CNN) architecture commonly used for image classification tasks. It consists of 50 layers, including convolutional layers, pooling layers, and fully connected layers. The model parameters, such as the learning rate, are set to their default values. In the NLP task, the structure of the generator is set as DistilBERT, and the model parameters are also set to their default values. We employ three fully connected layers as the classifier in each task. Details regarding the impact of various generators can be found in Section 5.5.

Table 4 summarizes the datasets used in this study, including the dataset name, a brief description, the number of samples, and the types of features in each dataset. To accommodate the computational requirements and ensure efficient experimentation, we sampled a subset (20,000) from some NLP datasets as the training data.

Setup. In the CV experiment, we focus on a common image classification task. Here’s an example to illustrate the process: we start by initializing the model architecture, such as ResNet18 [9], and then train the source model $M_{s}$ using CIFAR10. Then we utilize CIFAR100 to create model pools and evaluation datasets. To achieve this, we partition CIFAR100 into 10 disjoint 10-way classification subsets. Out of these, we randomly select 5 subsets (e.g., 3, 6, 8, 9, 5) to train 10 models (consisting of 5 homologous and 5 non-homologous models) that become part of the model pool. The remaining subsets are used to train multiple homologous and non-homologous models, which then serve as evaluation data. This setup is further detailed in the first row of Table 1.

To the best of our knowledge, no method has been designed for the MP task. We use MGMP with and without the DNA generator module (Figure 2) as a baseline method and observe that the inclusion of the DNA generator improves the provenance prediction accuracy. Specifically, the test accuracy obtained without the DNA generator (shown in parentheses) is lower, highlighting the importance of the DNA generator module in our proposed framework.

To provide a thorough insight into how DNA fragments contribute to the overall prediction, we analyze their influence by assessing performance at the DNA fragment level. Additionally, for assessing performance in terms of the DNA level, a threshold $\delta$ (e.g., 0.9) can be utilized to make final predictions. As demonstrated in Table 1, we consistently achieve the right provenance prediction in each row.

Setup. For text classification tasks on various benchmark NLP datasets, we utilize the DistilBERT and BERT-base architectures as shown in Table 2. Similar to the experimental setup for CV tasks, we randomly select one dataset to train the source model, and three datasets to form the model pool. For evaluation, we use two datasets to train homologous or non-homologous models.

We visualize the generated DNA fragments of the first CV experiment (i.e., the first row of Table 1). Figure 4 shows T-SNE visualizations of the space of DNA fragments results. The red points represent the generated DNA of the source model, the yellow ones are DNA from homologous target models, and the blue ones are DNA from non-homologous target models. We observe that the points belonging to homologous models are almost much closer to each other than points from non-homologous models. Specifically, the distance between the red and yellow points is much smaller than the distance between the red and blue points. This suggests that our framework is effective in capturing the similarity between homologous models and distinguishing them from non-homologous models.

The analysis of the MGMP framework aims to gain a comprehensive understanding of the individual components and processes within the approach. In this study, we specifically investigate the influence of different generator structures and various methods of DNA assembly. To isolate the effects of these factors, we conduct separate evaluations where one component is varied while keeping other components fixed. To establish a baseline experimental set, we consider the first row of Table 1 as the initial configuration.

We observed that the performance of different generator structures remained consistent across various evaluation metrics. We found that the DNA assembly process had a notable effect on the results. Different ways of assembling the DNA fragments resulted in variations in performance. Our results indicate that increasing the dimensionality of the DNA representation can lead to improved results within the MGMP framework. By incorporating additional dimensions into the DNA, we can capture more nuanced information and potentially enhance the performance of the model. By analyzing the results in Table 5, we gain insights into the contributions of individual components and their interactions within the MGMP framework. This analysis helps us identify the optimal configuration and provides valuable guidance for future improvements and refinements of the approach.