Super-model ecosystem:
A domain-adaptation perspective

Large-scale pretrained models (or foundation models) (Han et al., 2021; Chen et al., 2021), including GPT-3 (Brown et al., 2020) and BERT (Devlin et al., 2018), enables a new paradigm in machine learning: pre-training a large-scale model on very large-scale datasets and then transferring the learned model to an unseen domain. This paradigm was first introduced in natural language processing and recently to computer vision. It sheds light in a higher-level automation and is establishing a new paradigm with advantages/supremacy. In this paradigm, a super model learns meta knowledge from large amounts of data and reduces learning cost in specific domains. This helps considerably reduce the computational and data cost of applying machine learning in many specific applications. This is thus of significant values to enormous small and medium-sized enterprises. Additionally, super-model paradigm enables better management of the geographic location of machine learning workload and the datacenter infrastructure, which has been shown able to significantly reduce the carbon emission (Patterson et al., 2021).

Technically, domain adaptation plays a vital role for the knowledge transferring in the super-model paradigm. Usually, the data in a target domain is much smaller than the one in the source domain. In the light of this, an appropriate understanding to the generalizability of the transferred model on the target domain is of high importance.

In this paper, we prove an upper bound for the generalization error (generalization bound) for domain adaptation algorithms. The generalization error is defined as the difference between the expected risk $R$ and the empirical risk $\hat{R}$. Intuitively, a larger generalization bound indicates that the generalization error is possibly larger and thus suggests worse generalizability.

We model the super model paradigm as a two-stage diffusion processes. In the first stage, stochastic gradient-based optimizers, usually stochastic gradient-based optimization, including stochastic gradient descent (SGD) (Robbins and Monro, 1951), momentum (Nesterov, 1983; Tseng, 1998), and Adam (Kingma and Ba, 2014), learns a pre-trained model on the source-domain data via empirical risk minimization,

where ${\hat{R}}_S\left(\theta \right)$ is the empirical risk of model parameterized by $\theta$ on the training sample $S$, which is defined to be

where $N$ is the training sample size, and $l$ is the loss function. The convergent parameter initializes the fine-tuning process on the target domain in the second stage. We model the parameter trajectory of SGD by a stochastic process, Uhlenbeck-Ornstein process (Uhlenbeck and Ornstein, 1930), as follows,

where $B$ is positive definite matrix which characterizes the covariance of the gradient noise. This can also be smoothed to a diffusion equation, Fokker-Plank equation. Correspondingly, the trajectories can be modeled by the dynamics of the Fokker-Plank equations. Further, the steady distributions of the Uhlenbeck-Ornstein equations characterize the distributions of the learned models.

Deep learning can be formulated as solving a non-convex optimization problem: The loss surface of neural networks are usually highly non-convex due to the complexity of neural network architectures. In general, solving a non-convex optimization problem is NP-hard. However, numerous experiments show that deep learning has excellent optimization performance. This mystery is partially addressed by some empirical finding on the local convexity and smoothness of the loss surfaces of deep neural networks. Empirical results show that the loss surface around the convergent local minima is second-order smooth, as shown by Li et al. (2018).

This empirical finding inspires us to model the loss surface around the convergent local minimum as a quadratic function. This assumption determines the derivatives and boundary conditions of the Fokker-Plank equation. Moreover, the model parameter is usually initialized by following a Gaussian distribution. Based on them, Fokker-Plank equation has a steady distribution in the form of Maxwell-Boltzmann distribution, which governs the distribution of the learned model by the SGD. During the pre-training stage, SGD converges a Maxwell-Boltzmann distribution around the local minimum given below,

where $M_{PT}$ is the normalizer and ${\Sigma }_{PT}$ is the covariance.

This distribution is then used as the initial distribution in the fine-tuning stage. Subsequently, SGD in the fine-tuning stage learns the mapping from the initialization to a new Maxwell-Boltzmann distribution centered at the new local minimum on the loss surface in the target domain as follows,

where $M_{FT}$ is the normalizer, ${\theta }_{FT}$ is the distribution shift, and ${\Sigma }_{FT}$ is the covariance.

Based on the diffusion processes, we then establish PAC-Bayesian generalization bounds for the learned super model on the source domain and the transferred model on the target domain. The PAC-Bayesian framework (McAllester, 1999a, b) upper bounds the generalization error of a stochastic algorithm via the distance between the initial distribution and the distribution of the learned hypothesis, usually measured by some information-theoretical distances, such as KL-divergence. Intuitively, the PAC-Bayesian theory suggests that training a very-large model from a no-knowledge prior, such as Gaussian distribution and uniform distribution, needs a very large amount of data to secure the generalizability; and if the initialization is near the distribution of the learned hypothesis, the needed sample complexity can be relatively much smaller. However, a high-quality prior is not accessible in practice. This significantly limits the model size, particularly in low-resource scenarios. This renders the key motivation of the super-model paradigm: (1) training a super model on a very large-scale dataset, in order to learn a high-quality model from the no-knowledge prior; and (2) using the learned super model as a high-quality prior in the down-stream application, in order to reduce the needed training data and supports larger model size.

In this paper, the generalization bound in pre-training is established based on the KL-divergence between the Maxwell-Boltzmann distribution in pre-training as below,

where

and $R\left(Q_{PT}\right)$ is the expected risk, $\hat{R}\left(Q_{PT}\right)$ is the empirical risk, ${\Sigma }_{PT}$ is the covariance of the distribution of the learned hypothesis, and $N_{PT}$ is the training sample size in the pre-training.

Meanwhile, the generalization bound in pre-training is established based on the KL-divergence between the two Maxwell-Boltzmann distributions in pre-training and fine-tuning as follows,

where

and $R\left(Q_{FT}\right)$ is the expected risk, $\hat{R}\left(Q_{FT}\right)$ is the empirical risk, ${\Sigma }_{FT}$ is the covariance of the distribution of the learned hypothesis, ${\theta }_{FT}$ is the shift of the distribution center, and $N_{FT}$ is the training sample size in the fine-tuning.

We further define two new notions to measure the domain discrepancy as follows,

and

where $d$ is the parameter size. These two notions measure the magnitude of the domains shifts based on the learned hypotheses on the source domains and target domains.

Our theory have the following implications:

• The very large-scale datasets employed in the pre-training stage helps secure obtaining a high-quality model from a no-knowledge prior. The learned model carries the knowledge learned from the training data in the pre-training stage. The distribution of the pre-trained model severs as a high-quality initialization in the down-stream fine-tuning stage.

• Comparing the generalization bounds in the pre-training stage and fine-tuning stage, we show that the generalization error of the fine-tuning stage is dominant in the super-model paradigm. This is because the dominantly large size of the training data in the pre-training stage. This finding supports the feasibility and efficiency of the super-model paradigm.

• Our generalization bound suggests that the generalization on the target domain is determined by the magnitude of the domain shifts. Generally, larger domain shifts lead to worse generalization on the target domain. This supports the heuristic in practice that the performance on the target domain is limited by the domain shifts.

It is worth noting that the super-model paradigm also supports model compression in the model deployment, including pruning, quantization, model distillation, etc. The influence of model compression methods can be directly plugged in our theory.

This section reviews the related work, including super model, domain adaptation, generalization, and deep learning theory.

Domain adaptation. Domain adaptation algorithms transfer knowledge from one domain to another. It enables the super model paradigm. Domain adaptation has three main streams:

(1) Discrepancy-based domain adaptation modifies the loss to narrow the discrepancy between the features from the source domain $D^s$ and the ones from the target domain $D^t$. Tzeng et al. (2014) introduce a fully-connected adaptation layer into CNN for learning the representation of the kernel $\varphi (\cdot )$ in order to minimize the maximum mean discrepancy (MMD) between the features from different domains:

Long et al. (2015) employ multiple adaptation layers. Long et al. (2015, 2016) introduce residual blocks into the classifiers of source domain. Long et al. (2017) further consider the discrepancy of the joint distribution $P(X,Y)$ rather than the marginal distribution $P\left(X\right)$;

(2) Adversarial-based domain adaptation maps the source domain $D^s$ and the target domain $D^t$ to a general space, inspired by generative adversarial networks (GANs): If a classifier hardly separates examples of source domain and those from target domain, the feature extractor has narrowed the two domains. Ganin and Lempitsky (2015) propose gradient reversal layers that reverse the gradients generated by the domain classifier during backpropagation. Zhang et al. (2018) argue that the features of the bottom layers contain more domain information, while those of the top layers contain less domain information. They further employ collaborative learning to learn domain informative features in the bottom layers, and adapt adversarial learning to learn domain uninformative features in the top layers; and

(3) Reconstruction-based domain adaptation reconstructs the features extracted from the source domain $D^s$ to the target domain $D^t$. Ghifary et al. (2016) reconstruct examples from the target domain via features learned from source domain classification task. Bousmalis et al. (2016) reconstruct the inputs via both private representation and shared representation of both source and target domains.

Generalization. Good generalization guarantees that an algorithm learns the underlying patterns in training data rather than just memorize the data. In this way, good generalization abilities provide confidence that the models trained on existing data can be applied to similar but unseen scenarios. Three major approaches in analyzing the generalizability are seen in the literature: (1) generalization bounds based on the hypothesis complexity, including VC dimension (Blumer et al., 1989; Vapnik, 2006), Rademacher complexity (Koltchinskii and Panchenko, 2000; Koltchinskii, 2001; Bartlett and Mendelson, 2002), and covering number (Dudley, 1967; Haussler, 1995). The results are usually obtained via concentration inequalities. They also suggest controlling the model size to secure the generalizability, which is no longer valid in deep learning; (2) generalization bounds based on the algorithmic stability (Rogers and Wagner, 1978; Bousquet and Elisseeff, 2002; Xu et al., 2011). The results in this stream follow the motivation that learning algorithms robust to small disturbances in input data usually have good generalizability; and (3) generalization bounds in the PAC-Bayes framework (McAllester, 1999a, b). The results are obtained based on information-theoretical versions of concentration inequalities.

Deep learning theory. Deep learning has been deployed successfully in many real-world scenarios. However, the theoretical foundations of deep learning are still elusive. For example, there is no explanation for how deep learning algorithms work, why they can succeed, when they would fail, and whether they would hurt society. Such deficiency in explainability questions the transparency and accountability of deep learning, and further undermines our confidence of deploying deep learning in security-critical application domains, such as medical diagnosis (Kulikowski, 1980; Silver et al., 2016) and drug discovery (Chen et al., 2018a). Many works have emerged to establish the theoretical foundations of deep learning via VC dimension (Harvey et al., 2017), Rademacher complexity (Golowich et al., 2018; Bartlett et al., 2017), covering number (Bartlett et al., 2017), Fisher-Rao norm (Liang et al., 2019; Tu et al., 2020), PAC-Bayesian framework (Neyshabur et al., 2017), algorithmic stability (Hardt et al., 2016; Kuzborskij and Lampert, 2018; Verma and Zhang, 2019), and the dynamics of stochastic gradient descent or its variants (Mandt et al., 2017; Mou et al., 2018b; He et al., 2019). Please see more related works in surveys (E et al., 2020; He and Tao, 2020; Poggio et al., 2020). This work is committed to establishing theoretical foundations of privacy, generalization, adversarial attack in deep learning, all of which have profound importance in enhancing the explainability, transparency, and accountability of deep models.

Generalization of SGD. Some generalization bounds for algorithms trained by SGD are proposed. Mou et al. (2018a) analyze the generalization of stochastic gradient Langevin dynamics (SGLD), and prove an $O\left(1/N\right)$ upper bound and an $O\left(1/\sqrt{N}\right)$ upper bound for the generalization error, respectively via algorithmic stability and PAC-Bayesian theory. Pensia et al. (2018) analyze the generalizability of noisy and iterative machine learning algorithms. A generalization bound is then proved given the mutual information between the output hypothesis and the input data. It also proved generalization bounds for SGLD as examples. Chen et al. (2018b) prove that the convergence and stability for iterative machine learning algorithms have a trade-off under both convex smooth assumption and strong convex smooth assumption. Under the same assumptions, Chen et al. (2018b) prove an $O\left(1/N\right)$ generalization bound for SGD. Liu et al. (2017) prove an $O\left(1/N\right)$ generalization bound for SGD when the loss function is Lipschitz continuous and smooth. London (2017) prove a generalization bound for SGD based on the KL divergence between the prior $P$ and the posterior $Q$ under the PAC-Bayes framework. He et al. (2019) present a PAC-Bayes generalization bound for SGD based on stochastic differential equations. In the work of He et al., the gradient noise is modeled by a Gaussian distribution. Meng et al. (2020) extend the gradient noise to be state-dependent. Cheng et al. (2020) extend the gradient noise to be Levy process.

Suppose the training dataset is $S=\left\{\right(x_1,y_1),\dots ,(x_N,y_N)|x_i\in R^{d_X},y_i\in R^{d_Y},i=1,\dots ,N\}$, where $d_X$ is the dimension of the feature $X$ and $d_Y$ is the dimension of the label $Y$. Suppose $x_i$ and $y_i$ are independent and identically distributed (i.i.d.) observation of variables $X\in X$ and $Y\in Y$, respectively. We also rewrite $z_i=(x_i,y_i)$, which is an i.i.d. observation of random variable $Z=(X,Y)\in Z$. Denote the generating distribution of $Z$ is $D$.

Formally, machine learning algorithms are designed to select the hypothesis function $F_{\theta }$ with the lowest expected risk $R$ under the loss function $l$ from a hypothesis class $\{F_{\theta }|\theta \in \Theta \subset R^d\}$, where $\theta$ is the parameter of the hypothesis and $d$ is the dimension of the parameter $\theta$. For many stochastic algorithms, such as SGD, we usually use a distribution to express the output parameter. Suppose the parameter follows a distribution $Q$, the expected risks respectively in terms of $\theta$ and $Q$ are defined as:

However, the expected risk $R$ is not available from the data, since we do not know the formulation of latent distribution $D$ of data. Practically, we use the empirical risk $\hat{R}$ to estimate the expected risk $R$, which is defined as:

where all $(X_i,Y_i)$ constitute the training sample $T$.

Learning algorithms usually solve the following empirical risk minimization (ERM) problem to approach the optimal hypothesis,

We usually employ stochastic gradient-based optimizers for ERM in deep learning. Popular options of stochastic gradient-based optimizers include stochastic gradient descent (SGD) (Robbins and Monro, 1951), momentum (Nesterov, 1983; Tseng, 1998), and Adam (Kingma and Ba, 2014). For the brevity, we analyze SGD in this paper. The analysis for other stochastic gradient-based optimizers is similar.

Suppose $B$ is a mini batch randomly drawn from the training sample set $S$. Then, the stochastic gradient on $B$ is as follows,

In the $t$-th iteration, the weight is updated as follows,

where ${\theta }_t^{ERM}$ is the weight vector in the $t$-th iteration and ${\eta }_t$ is the corresponding learning rate.

Meanwhile, adversarial training employs SGD to solve the following minimax problem,

where $\rho$ is the radius of the ball centered at the example $(x_i,y_i)$. Here, we call ${\hat{R}}_S^A\left(\theta \right)$ adversarial empirical risk. Correspondingly, the stochastic gradient on a mini batch $B$ and the weight update are calculated as below,

To avoid technicalities, the measurability/integrability issues are ignored throughout this paper. Moreover, Fubini’s theorem is assumed to be applicable for any integration with respect to multiple variables, that the order of integrations is exchangeable. Also, we assume the stable (stationary) solutions of all stochastic differential equations involved exit and are unique.

A supreme industrial paradigm has been emerging that (1) pre-training a large-scale model on large amounts of multi-modality data, such as GPT-3 (Brown et al., 2020) and BERT (Devlin et al., 2018); and (2) fine-tuning the obtained model on specific smaller domain where data size is relatively difficult to access. In this paper, we name it as super-model paradigm. Super-model paradigm enables efficient and effective knowledge discovery in low-resource application scenarios, including few-shot learning (Snell et al., 2017; Sung et al., 2018) and zero-shot learning (Romera-Paredes and Torr, 2015). A key cornerstone technology wherein is domain adaptation. This section describes this paradigm.

Large-scale pre-trained models. Recent advances are seen mainly in natural language processing (NLP), particularly after the appearance of transformer (Vaswani et al., 2017). ELMo (Peters et al., 2018) finds the word embedding in NLP is not invariant in different application domains, but considerably changes with context. Based on this observation, ELMo pre-trains a large-scale bidirectional LSTM on a large text corpus to generate word vectors by fine-tuning. BERT (Devlin et al., 2018) employs the transformer encoder for detecting bidirectional information in the context. Meanwhile, Liu et al. (2018) employs the transformer decoder to word embedding with fine-tuning, in order to realize wider attention. GPT (Radford et al., 2018) also employs the transformer decoder but is fine-tuned on each specific task for better performance. Extended from GPT, GPT-2 (Radford et al., 2019) and GPT-3 (Brown et al., 2020) construct huge models in order to realize zero-shot learning. The comparison between these “super models” is presented in the following table.

Pre-training stage. The first step of the super-model paradigm is pre-training a super model on large-scale data, sometimes of multi-modality. The learned model in this stage is of high quality that the approximation and generalization of the output hypothesis are usually excellent, which suggests that the learned model has stored rich general knowledge in the learned model. This makes it possible to apply the learned model for the smaller specific application domains.

Fine-tuning stage. The learned model is then fine-tuned on the target domain, usually a smaller specific domain. The stored general knowledge is thereby transferred to the target domain. In this way, super-model paradigm reduces considerable sources of knowledge discovery in the target domain.

Theoretical advantages. According to the PAC-Bayesian theory, the generalizability of the learned model is determined by the distance between the posterior and the prior. As we will show in the next two sections, the very large-scale training data in the pre-training stage secures learning a high-quality model with no-knowledge prior. The learned knowledge is of high value but consumed enormous resources which is not accessible for many potential machine learning users, particularly small and medium-sized enterprises. In the super-model paradigm, the high-quality model learned in the pre-training stage is employed as the initialization in the fine-tuning stage. In this way, we significantly reduce the needed sample complexity in the fine-tuning stage.

Industrial values. Machine learning has been thriving in a wide range of areas. However, the industrial applications are still limited. This is partially caused by the high cost of computing facilities and data annotations. The paradigm based on super models significantly reduce the cost of machine learning applications. This is particularly important for small and medium-size enterprises.

Climate value. Super-model paradigm enables recycling discovered general knowledge in enormous application domains. This would also help significantly reduce the carbon emission. Meanwhile, the super-model paradigm centralizes the modeling training process which can help manage the geographic location and the datacenter infrastructure in order to reduce the carbon emission, as a recurrent work suggested (Patterson et al., 2021):

• Geographic location of machine learning workload scheduling can result carbon emission vary around five times to ten times, even when the country and the organization remain invariant.

• Cloud data centers can be around 1.4-2X more energy-efficient. Meanwhile, machine learning-oriented accelerators can be  2-5X more effective.

Super-model paradigm can thus reduce the carbon print of machine learning application and further contribute in slowing down the climate crisis.

We consider a diffusion process-based model that serves an envelope for domain adaptation methods. Two diffusion processes are designed for modeling the pre-training and fine-tuning stages, respectively. The knowledge transition can then be modeled via the transition of diffusion processes.

The knowledge transition characterized by the diffusion process in the fine-tuning. In this paper, we employ the PAC-Bayesian theory to analyze the generalizability of domain adaptation.

Large-scale pre-trained models, such as GPT-3 and Bert, enables a new industrial paradigm: pre-training a super model on large amounts of multi-modality data (sometimes of low-quality) and then fine-tuning the learned model to smaller specific application domains. This paradigm may start a super-model paradigm that would significantly reduce the application cost of machine learning, which is critical for enormous small and medium-sized enterprises.

A major technique in this paradigm is domain adaptation which enables the knowledge transfer between the two domains. We model a super-model paradigm as a two-stage diffusion process: (1) in the pre-training stage, the trajectory of the stochastic gradient descent (SGD) or its variants searches on the loss surface driven by Uhlenbeck-Ornstein process discretely or smoothly by the Fokker-Plank equation. The model weight starts from a no-knowledge prior and converges to a Maxwell-Boltzmann distribution; and (2) in the fine-tuning stage, the trajectory of SGD is driven by a similar SDE, which starts from the learned model distribution in the pre-training stage and converges to another Maxwell-Boltzmann distribution. Based on the diffusion processes, an $O\left(1/\sqrt{N}\right)$ generalization bound is obtained via the PAC-Bayesian framework.

The generalization bounds suggest that the fine-tuning stage dominates the generalization of the whole paradigm. The generalization is determined by the domain discrepancy between the pre-training and fine-tuning domains, which is characterized by a new measure based on the covariance and domain shifts.

In this work, we make several assumptions and abstractions. This section discusses the limitation introduced by them and give several potential extensions.

• Model compression in the fine-tuning stage. In this paper, we ignore the model compression approaches in the fine-tuning stage, which are sometimes employed in practice. Popular model compression methods include model distillation, pruning, and quantization. The effects of model compression can be seen as operators on the loss surface and the learned model. A future direction is to mathematically characterize the influence of model compression. The results may be plug-and-play components to the presented theory in this paper.

• Gradient noise in SGD. In this paper, we assume that the gradient noise is drawn from a Gaussian distribution. Recent works also made assumptions that the gradient noise as a Levy process, Laplacian noise, etc. The exact distribution of the gradient noise is still an open problem. In addition, the gradient noise is assumed state-independent, which can be easily extended to state-dependent. A future direction is to study the distribution of the gradient noise in SGD. It is worth noting that relatively little efforts are needed to change the gradient noise distribution assumptions in this paper.

• Advanced techniques in modeling SGD. In this paper, we model the trajectory of SGD via Fokker-Planck equation and Uhlenbeck-Ornstein process. This modeling ignores the influence of several techniques, such as momentum and adaptive learning rate. Recent works discover that these techniques may have implicit regularization on the learned model while would not have determinant impact. A future direction is modeling the SGD as a more sophisticated stochastic differential equation.

• Distribution/data-dependent priors. Some works design priors relying on the data generation distribution but still not directly relying on the training data. This would be reasonable since we can assume the data distribution has been fixed before the data was collected (Lever et al., 2013). Such distribution-dependent priors have shown to be able to considerably tighten the generalization bounds. Negrea et al. (2019) further push the frontier that constructs priors not independent with data. Suppose $S_j$ is a subset of $\subset S$ with size of $n<m$. One may design a prior exploiting $S_J$ to deliver a data-dependent forecast of the posterior $Q$. A future direction is modeling the SGD via distribution/data-dependent priors.

This section presents the proofs for the given theory.

We model the the iterative updates in SGD employing a stochastic differential equation. This approach is also seen in the literature; see, e.g., E (2017); Mandt et al. (2017); Mou et al. (2018a); He et al. (2019); Meng et al. (2020); Cheng et al. (2020); Xie et al. (2020); Wang et al. (2021).

We first translate the updates in SGD as Ornstein-Uhlenbeck process (Uhlenbeck and Ornstein, 1930) under some mild assumptions. The Ornstein-Uhlenbeck process has a steady distribution which is then employed to characterize the distribution of the learned hypothesis. We further obtain a generalization bound via PAC-Bayesian framework by exploiting the stationary distribution, which characterizes the influence on the generalization via the distance between the output hypothesis distribution and its prior (McAllester, 1999a, b).