# Efficient Stitchable Task Adaptation

Haoyu He Zizheng Pan Jing Liu Jianfei Cai Bohan Zhuang
ZIP Lab, Monash University, Australia Corresponding author. Email: $\tt
<EMAIL_ADDRESS>
###### Abstract

The paradigm of pre-training and fine-tuning has laid the foundation for
deploying deep learning models. However, most fine-tuning methods are designed
to meet a specific resource budget. Recently, considering diverse deployment
scenarios with various resource budgets, SN-Net [54] is introduced to quickly
obtain numerous new networks (stitches) from the pre-trained models (anchors)
in a model family via model stitching. Although promising, SN-Net confronts
new challenges when adapting it to new target domains, including huge memory
and storage requirements and a long and sub-optimal multistage adaptation
process. In this work, we present a novel framework, Efficient Stitchable Task
Adaptation (ESTA), to efficiently produce a palette of fine-tuned models that
adhere to diverse resource constraints. Specifically, we first tailor
parameter-efficient fine-tuning to share low-rank updates among the stitches
while maintaining independent bias terms. In this way, we largely reduce fine-
tuning memory burdens and mitigate the interference among stitches that arises
in task adaptation. Furthermore, we streamline a simple yet effective one-
stage deployment pipeline, which estimates the important stitches to deploy
with training-time gradient statistics. By assigning higher sampling
probabilities to important stitches, we also get a boosted Pareto frontier.
Extensive experiments on 25 downstream visual recognition tasks demonstrate
that our ESTA is capable of generating stitches with smooth accuracy-
efficiency trade-offs and surpasses the direct SN-Net adaptation by remarkable
margins with significantly lower training time and fewer trainable parameters.
Furthermore, we demonstrate the flexibility and scalability of our ESTA
framework by stitching LLMs from LLaMA family, obtaining chatbot stitches of
assorted sizes111Source code will be released at
https://github.com/ziplab/Stitched_LLaMA.

##  1 Introduction

The paradigm of pre-training and fine-tuning has underpinned modern
applications in both vision and language. With off-the-shelf models pre-
trained on large-scale datasets, the de-facto choice is vanilla full fine-
tuning, which tunes all the model parameters with the downstream data. To
reduce memory footprint and avoid overfitting, an emerging trend is to tune a
small proportion of the model parameters while freezing the majority ones with
Parameter-Efficient Fine-Tuning (PEFT) [33, 29, 30]. However, both full fine-
tuning and PEFT target an exclusive specific resource budget for each
downstream task, while in reality, we often need to deploy the fine-tuned
models to diverse platforms, _e.g_., from a high-end iPhone 15 pro to a low-
end Nokia X5 and consider their dynamic resource budgets such as energy,
latency, and on-chip memory. This makes us ponder: _Can we efficiently obtain
a palette of fine-tuned models to meet different resource constraints?_

To satisfy a range of deployment resource constraints, a natural approach is
to compress a well-trained large model into numerous smaller ones [12, 31, 38,
23]. However, the computational cost grows linearly with the number of
deployment scenarios. Following the once-for-all network [9, 22], a few works
first pre-train a weight-sharing over-parameterized supernet, then adapt the
supernet to the downstream tasks. Although promising, training supernets at
scale with large datasets requires prohibitive computational resources,
_e.g_., thousands of GPU training hours, which is practically infeasible.

Table 1: Comparison of frameworks for obtaining a palette of networks to meet
different hardware efficiency constraints. Supernet-based methods NAT [50] and
TOFA [37] necessitate a supernet training stage plus long deployment GPU hours
and are restricted to small FLOPs ranges. SN-Net [54] employs model stitching
to achieve a wide FLOPs range with significantly reduced deployment GPU hours.
Our ESTA framework further reduces the deployment GPU hours, fine-tuning
memory, and trainable parameters, and is scalable to produce LLM stitches with
billions of parameters (Section 5.2). The deployment GPU hours and fine-tuning
memory for both SN-Net and ESTA are measured on a single NVIDIA GeForce RTX
3090 GPU when stitching ViT-Ti/S/B anchors [18] with batch size 64. ∗
indicates data is inferred from the training recipe of TOFA’s supernet
architecture FBNetV3 [15].

Method | Adapting supernet | Deployment GPU hours | Fine-tuning memory | Trainable parameter $\\#$ | FLOPs range (G) | LLM friendly
---|---|---|---|---|---|---
NAT [50] | ✗ |  $>$1,000 | - |  $>$10M | $[0.2,0.6]$ | ✗
TOFA [37] | ✓ |  $>$1,000∗ | - |  $>$10M∗ | $[0.2,2.5]$ | ✗
SN-Net [54] | ✗ | 19.3 | 13,235M | 124.2M | $[1.3,17.6]$ | ✗
ESTA (ours) | ✗ | 5.0 | 9,685M | 4.6M | $[1.3,17.6]$ | ✓

On the other hand, there are many open-source large models [67, 25, 6] from
communities such as HuggingFace [75] that are pre-trained on large-scale
datasets and ready to be downloaded. Inspired by the success of model
stitching [40, 5, 13], recently, Stitchable Neural Networks (SN-Net) [54] has
been proposed to stitch pre-trained models (anchors) of the same family to
quickly obtain a set of candidate new networks (stitches) with different
accuracy-efficiency trade-offs of a wide FLOPs range. However, the scope of
SN-Net is limited to the pre-training classification task on the same source
domain [59]. When directly adapting the standard approach of SN-Net to a
target domain, it faces new challenges. Specifically, the straightforward way
to employ SN-Net in task adaptation typically requires three stages: first
adapting anchors individually to the target domain, followed by stitching the
adapted anchors, and finally evaluating all stitches to find and deploy the
ones on the Pareto Frontier. Such a three-stage adaptation process is
expensive and sub-optimal. Moreover, adapting SN-Net typically needs to load
and optimize all parameters of multiple models, leading to a daunting
training-time memory cost. For instance, stitching the two smallest models in
the LLaMA family [67, 67] with SN-Net under standard training settings (32-bit
floating point parameters and Adam optimizer on a single GPU) requires more
than 80G Video RAM that cannot fit into high-end GPUs such as NVIDIA A100. SN-
Net also requires saving a separate instance of the model family for each
task, making the storage grow linearly with the number of deployed
applications.

Therefore, in this paper, we introduce a novel Efficient Stitchable Task
Adaptation (ESTA) framework. Specifically, built upon SN-Net [54], ESTA
introduces two main designs to overcome the aforementioned issues. First, to
reduce the memory and storage costs, we tailor a Parameter-efficient Stitch
fine-Tuning (PST) method, which incorporates the representative sparse fine-
tuning technique LoRA [30] to keep the anchors and stitching layer weights
frozen, while approximating their updates with trainable low-rank
decomposition matrices. In addition to stitch-agnostic LoRA modules, we
further introduce stitch-specific bias terms to alleviate the conflict among
different stitches.

Second, we propose a simple one-stage deployment pipeline to simultaneously
adapt and stitch the pre-trained anchors to the target domain. Moreover, we
propose to estimate and accumulate an importance score for each stitch via the
saliency pruning criterion [39]. With the importance scores, we develop a
novel stitch-sampling method to assign important stitches with higher sampling
probabilities. More importantly, after fine-tuning, the stitch importance
scores can be directly used to infer the best ones at the Pareto frontier for
deployment. Table 1 compares different frameworks for obtaining a palette of
networks for diverse deployments.

Overall, our paper has the following major contributions. 1) This is a
pioneering work to investigate the problem of efficient fine-tuning of SN-Net
for task adaptation. Our solution ESTA is a successful endeavor that delivers
dozens of ready-to-deploy vision models up to 17.6G FLOPs in 5 GPU hours as
well as multiple chatbot models with billions of parameters by stitching LLaMA
models [67, 68] for the instruction-following task. 2) We devise a parameter-
efficient stitch fine-tuning method that incorporates trainable stitch-
agnostic LoRA modules and stitch-specific bias terms, which largely reduces
the fine-tuning memory footprints while alleviating the interference issue
among stitches. Moreover, we streamline a simple one-stage deployment pipeline
with a novel task-specific stitch sampling strategy that greatly reduces the
deployment time and improves Pareto frontiers. 3) Extensive experiments
conducted on 25 downstream visual recognition tasks demonstrate that when
stitching ViT-Ti/S/B [18], our ESTA achieves remarkable performance
improvements compared to the direct SN-Net adaptation counterpart while
utilizing significantly fewer trainable parameters.

##  2 Related Work

Model stitching. Model stitching [40, 5] targets connecting the bottom layers
of one network to the top layers of another with a stitching layer. Following
[40], model stitching is first employed to measure the similarities for the
inner representations learned by deep neural networks [5, 13, 27]. Previous
works [5, 13] observe that trained networks with different initializations can
be effectively stitched without incurring significant performance drop. Based
on this observation, DeRy [80] stitches the blocks dissected from different
pre-trained models and assembles a new model based on their quantified
similarity for better performance. Herdt _et al_. [26] propose interpreting
the inner representations of a deep network by stitching it with a pre-trained
GAN generator. Recently, from the perspective of flexible deployment, SN-Net
[54] stitches models of different sizes within a model family to cheaply
produce numerous new networks with diverse accuracy and efficiency trade-offs
in the pre-training task. In contrast to [54], our work specifically targets
obtaining a palette of fine-tuned models for the task adaptation setting. To
do so, we make key redesigns including a parameter-efficient stitch fine-
tuning method and a simple one-stage deployment pipeline to overcome the
efficiency challenges in SN-Net adaptation.

Parameter-efficient fine-tuning. Parameter-efficient fine-tuning [29, 42, 30,
33] is a powerful alternative to vanilla full tine-tuning, which updates only
a small number of parameters while freezing the majority ones. Freezing the
majority of parameters allows us to optimize storage and reduces the burden on
training GPU memory, as there is no need to store their gradients or other
training time statistics. Recent research freezes the vast majority of
parameters while fine-tuning either the parameters that are inherited in the
backbone [81, 84] or additionally added, in the form of learnable prompts [33,
41], low-rank bottleneck layers [29], learnable scaling and shifting factors
[43], and separated small networks [64]. Meanwhile, PEFT has been employed in
numerous downstream applications. For instance, while freezing the most
parameters, Polyhistor [47] learns a hyper-network to generate adapter weights
[29] for multi-task adaptation, LLaMA-Adapter [87] inserts learnable prompts
into Transformer layers for the instruction-following task, and ControlNet
[86] employs the zero-initialized convolutions to fine-tune diffusion
generative models [58]. However, these works target only one exclusive
resource budget and are not scalable to diverse deployment scenarios. Unlike
other methods, our ESTA is crafted to yield multiple fine-tuned models with
diverse capacities while jointly incorporating stitch-specific and stitch-
agnostic lightweight trainable parameters.

Fine-tuning with multiple models. The presence of more large-scale models has
unleashed the potential to utilize multiple models instead of a single one
during fine-tuning. One line of work selects the best pre-trained model in the
model zoo to fine-tune. To identify the best model, prior studies quantize the
transferability of the pre-trained models ahead of fine-tuning by estimating
their accuracy on the downstream tasks [19, 3, 35] or the generalization
capability to mitigate the domain gaps [52, 69, 82, 7]. Another way is to do
model selection using an efficient online learning regime [78]. Considering
the diversity among the pre-trained models, research efforts have been devoted
to effectively exploit their knowledge with feature aggregation [61, 60, 44],
model merging [76, 32, 14, 51], or a mixture-of-experts architecture [62]. In
contrast to previous works that merely target better performance, our work
employs off-the-shelf pre-trained model families to produce plentiful new
models for diverse deployment requirements.

##  3 Preliminaries

###  3.1 Model Stitching

Considering an $L$-layer feed-forward artificial neural network
$f_{\theta_{1}}:{\mathcal{X}}\rightarrow{\mathcal{Y}}$ parameterized by
${\theta_{1}}$, that maps any input from the input space ${\mathcal{X}}$ to
the output space ${\mathcal{Y}}$, $f_{\theta_{1}}$ can be denoted as a
composition of functions that $f_{\theta_{1}}=f_{L}\circ\cdots\circ f_{1}$,
where $\circ$ denotes the function composition. In model stitching,
$f_{\theta_{1}}$ can be split up at the $l$-th layer into two portions of
functions where $l\in[1,L-1]$. Given any input
${\bm{X}}\subseteq{\mathcal{X}}$, the first portion of the head functions maps
${\bm{X}}$ to activations ${\bm{X}}_{l}$ at layer $l$, _i.e_.,
$H_{{\theta_{1}},l}({\bm{X}})=f_{l}\circ\cdots\circ f_{1}={\bm{X}}_{l}$. The
second portion of the tail functions maps ${\bm{X}}_{l}$ to the final output,
_i.e_., $T_{{\theta_{1}},l}({\bm{X}}_{l})=f_{L}\circ\cdots\circ f_{l+1}$.

Given another pre-trained $M$-layer artificial neural network $f_{\theta_{2}}$
parameterized by ${\theta_{2}}$ that is split up at the $m$-th layer that
$m\in[2,M]$, we can then employ a stitching layer
$S:{\mathcal{A}}_{{\theta_{1}},l}\rightarrow{\mathcal{A}}_{{\theta_{2}},m}$ to
map between the two activation spaces at layers $l$ and $m$. In this case, we
obtain a new network parameterized by $\phi$ and the stitching layer
parameters, _i.e_.,

$\displaystyle F_{\phi,S}({\bm{X}})=T_{{\theta_{2}},m}\circ S\circ
H_{{\theta_{1}},l}({\bm{X}}),$ (1)

where $\phi$ consists of partial parameters from both $\theta_{1}$ and
$\theta_{2}$.

Figure 1: (a) Illustration of Stitchable Neural Network [54]. With a pair of
anchors from the same model family, SN-Net connects the early layers of the
smaller one to the latter layers of the larger one with stitching layers so as
to obtain a set of new networks with different performance-efficiency trade-
offs, _e.g_., the path in Blue. (b) Overview of our PST method tailored for
fine-tuning a palette of stitches, which integrates stitch-agnostic LoRA
modules with stitch-specific bias terms, aiming to promote diverse
representations among stitches while maintaining low trainable parameters. (c)
Overview of our task-specific stitch sampling. We estimate the importance
scores of the stitches with a scoring function $Q(\cdot,\cdot)$ and accumulate
them as global statistics with moving averages. For a resource constraint
$\gamma$, we sample with a categorical distribution
$\pi({\mathcal{N}}_{\gamma})$ that is parameterized by the normalized
importance scores so as to assign the important stitches with higher sampling
probabilities. After fine-tuning, we directly deploy the stitches with the
highest scores to avoid the costly evaluation stage. Figure 2: Distribution of
pair-wise gradient angles among stitches when updating shared weights at fine-
tuning iteration 600. We highlight angle 90∘ with a dashed red line. For
simplicity, we show the gradient angles among the combined query, key, and
value projection matrices for a total of 32 stitches when stitching ViT-Ti and
ViT-S anchors. Generally, the gradient angles are larger in the target domain
Stanford Cars [20] than in the source domain ImageNet-1k [59].

###  3.2 Stitchable Neural Network

In Stitchable Neural Network (SN-Net) [54], let a pre-trained model family of
size $Z$ be ${\mathcal{Z}}=\\{f_{\theta_{z}}\\}_{z=1}^{Z}$, where $\theta_{z}$
is the model parameter of the $z$-th model. The goal is to derive additional
$N$ candidate new networks ${\mathcal{N}}=\\{F_{\phi_{n},S_{n}}\\}_{n=1}^{N}$
to adapt to various resource constraints. To do so, SN-Net first selects a
pair of pre-trained models (anchors) from ${\mathcal{Z}}$ and then stitches
them with Eq. (1) at different layer indexes to get $N$ stitches. Then, SN-Net
samples the stitches randomly and jointly optimizes them. An overview of SN-
Net is depicted in Figure 1 (a). Since the anchors vary in scale, the newly
assembled stitches have diverse performances and complexities. Importantly,
SN-Net gives practical principles to design the space of ${\mathcal{N}}$ that
one should 1) stitch a pair of nearest anchors in terms of model complexity in
a model family; 2) stitch the head of a faster and smaller anchor to the tail
of a larger and slower anchor. Unless specified otherwise, we adopt these as
the default experimental settings in this paper.

##  4 Method

In this section, we introduce our ESTA framework, which consists of two major
components: parameter-efficient stitch fine-tuning (PST) and a simple one-
stage deployment pipeline. PST is tailored to cheaply fine-tune a palette of
stitches (Section 4.1 and Figure 1 (b)). The one-stage deployment pipeline
aims to simultaneously save deployment time and improve adaptation performance
(Section 4.2).

###  4.1 Parameter-efficient Stitch Fine-tuning

To address the problems of large storage and memory consumption when adapting
SN-Net to downstream tasks with full fine-tuning, we introduce a parameter-
efficient stitch fine-tuning (PST) method as depicted in Figure 1 (b). Our
basic idea is to adapt the representative sparse fine-tuning method LoRA [30]
for SN-Net fine-tuning.

LoRA on Transformer layers. Specifically, let any pre-trained weight matrix in
the multi-head self-attention layer be ${\bm{W}}\in{\mathbb{R}}^{d\times k}$,
we freeze ${\bm{W}}$ and insert trainable low-rank decomposition matrices
${\bm{W}}_{\rm down}\in{\mathbb{R}}^{d\times r}$ and ${\bm{W}}_{\rm
up}\in{\mathbb{R}}^{r\times k}$, where $r$ is the pre-defined rank that
$r\ll{\rm min}(d,k)$. In this way, the updated version of ${\bm{W}}$ can be
formulated as

$\displaystyle{\bm{W}}^{\prime}\leftarrow{\bm{W}}+{\bm{W}}_{\rm
down}{\bm{W}}_{\rm up}.$ (2)

We follow [30] to respectively use Gaussian and zero initializations for
${\bm{W}}_{\rm up}$ and ${\bm{W}}_{\rm down}$, so that ${\bm{W}}_{\rm
down}{\bm{W}}_{\rm up}$ is zero at the beginning of fine-tuning.

LoRA on stitching layers. We further extend the low-rank weight update to
stitching layers. In SN-Net, a stitching layer is designed as a full-rank
transformation matrix. Without loss of generality, we show the low-rank update
of stitching layer $S_{n}$ that stitches $f_{\theta_{1}}$ and $f_{\theta_{2}}$
at layers $l$ and $m$, respectively. We first follow the initialization of
stitching layers in [54, 13] to align the activations at layers $l$ and $m-1$.
Let the activations be $H_{\theta_{1},l}({\bm{X}})\in{\mathbb{R}}^{b\times
d_{1}}$ and $H_{{\theta_{2}},m-1}({\bm{X}})\in{\mathbb{R}}^{b\times d_{2}}$
with sequence length $b$ and feature dimensions $d_{1}$ and $d_{2}$. The
stitching layer is parameterized by a transformation matrix
${\bm{M}}\in\mathbb{R}^{d_{1}\times d_{2}}$, which is initialized by solving a
least squares problem, _i.e_.,
${\bm{M}}=H_{\theta_{1},l}({\bm{X}})^{\dagger}H_{{\theta_{2}},m-1}({\bm{X}})$,
where $H_{\theta_{1},l}({\bm{X}})^{\dagger}$ is the Moore-Penrose
pseudoinverse of $H_{\theta_{1},l}({\bm{X}})$. We find that such
initialization already has impressive representational capacity on the target
domain and they can be updated with low-rank decomposition without hurting the
performance too much. Accordingly, we update ${\bm{M}}$ similar to updating
${\bm{W}}$ in Eq. (2) by approximating its update with two learnable low-rank
matrices with the same initializations.

Although the low-rank updates significantly enhance parameter efficiency, the
low-rank essence [2] largely limits the network capacity. Particularly, in
task adaptation, we often observe conflicting updates on the shared weights
among the stitches. Figure 2 gives an example, where we stitch pre-trained
anchors on ImageNet-1k [59] and Stanford Cars [20] and visualize the
distribution of pair-wise gradient angles among different stitches when
updating shared weights. We find that stitches agree less and conflict more
(_i.e_., with more large angles) on the target domain Stanford Cars compared
to the source domain ImageNet-1k.

Stitch-specific bias. To address the above issue, we further introduce stitch-
specific bias terms as depicted in Figure 1 (b). Particularly, for a linear
layer parameterized by ${\bm{W}}$ and a bias term ${\bm{b}}$, we add bias term
${\bm{b}}^{s}$ when optimizing the $s$-th stitch, for which the output
activation becomes ${\bm{W}}{\bm{X}}+{\bm{b}}+{\bm{b}}^{s}$. We employ a set
of distinct bias terms at different layers for each stitch. In this way, the
stitches are encouraged to learn distinct feature representations for
different resource requirements. To restrict the number of trainable
parameters, following [30], we only apply PST to self-attention layers while
freezing feed-forward layers.

###  4.2 One-stage Deployment Pipeline

Simultaneously adapt and stitch anchors. As aforementioned, given a model
family ${\mathcal{Z}}$ pre-trained in the source domain, directly adapting SN-
Net to a target domain typically involves three stages. First, adapting each
anchor to the target domain ${\mathcal{D}}$ by solving
${\theta_{z}}^{*}={\operatorname{argmin}}_{\theta_{z}}{\mathcal{L}}(\theta_{z};{\mathcal{D}}$).
In the second stage, at each training iteration, SN-Net samples a stitch and
optimizes it with objective ${\mathcal{L}}(\phi_{n}^{*},S_{n};{\mathcal{D}})$,
where $\phi_{n}^{*}$ is the set of SN-Net parameters for stitch
$F_{\phi_{n},S_{n}}$ which has been optimized once in the first stage. In the
final stage, all optimized stitches need to be evaluated to identify the ones
on the Pareto frontier for diverse resource constraints. Such a three-stage
approach is expensive and sub-optimal (see more discussions in Section F of
the appendix). Thus, we propose to adapt and stitch the anchors simultaneously
within one stage, by directly optimizing
${\mathcal{L}}(\phi_{n},S_{n};{\mathcal{D}})$ for each sampled stitch.
Moreover, we introduce a novel task-specific stitch-sampling method to allow
the promising stitches more likely to be sampled, based on a stitch importance
score. After fine-tuning, the stitches with the highest scores are naturally
selected for deployment without the need for the above final stage evaluation.

Figure 3: Performance comparisons with SN-Net [54] for adapting ViT-Ti/S/B
pre-trained on ImageNet-22k [16] to Stanford Cars [20], CUB-200-2011 [71],
Stanford Dogs [34], and NABirds [70]. We denote individually fine-tuned
anchors as yellow stars. We also show the number of trainable parameters.
Figure 4: Performance comparisons with SN-Net [54] for adapting ViT-Ti/S/B
pre-trained on ImageNet-22k [16] to VTAB-1k [85] and CIFAR-100 [36]. We denote
individually fine-tuned anchors as yellow stars and also show the number of
trainable parameters.

Task-specific stitch sampling. SN-Net assumes that all stitches are equally
important and performs random (uniform) sampling during training. However, as
recognized by prior research [73, 45] that random sampling causes a gap
between training and deployment as only the best stitches on the Pareto
frontier are actually deployed. Moreover, since the importance of the pre-
trained weights varies in different downstream tasks [21, 79, 24], we argue
that the performance of stitches that are parameterized by these weights is
also different across these tasks.

To this end, we propose to assign the important stitches that are likely to be
on the Pareto frontier with higher sampling probabilities during fine-tuning
to ensure that they are optimized sufficiently. Specifically, we first
estimate the importance of each stitch with a data-dependent saliency pruning
metric SNIP [39], which measures the importance with the first-order gradient
information with barely any extra computational cost. When the $n$-th stitch
is sampled, we can get its importance score $Q(F_{\phi,S_{n}},{\bm{X}})$ with
the scoring function $Q(\cdot,\cdot)$ [39]. To obtain robust importance scores
for sampling, we accumulate scores with moving average during training,
_i.e_.,

$\displaystyle q_{n}^{t}\leftarrow\eta
q_{n}^{t-1}+(1-\eta)Q(F_{\phi_{n},S_{n}},{\bm{X}}),$ (3)

where $q_{n}^{t}$ and $q_{n}^{t-1}$ are the importance score at the $t$-th and
$(t-1)$-th iteration, respectively, and $\eta\in[0,1)$ is the momentum
coefficient. In this way, we can get stable importance scores. After a warm-up
period that employs uniform sampling to accumulate the scores, we assign the
important stitches with higher sampling probabilities for the rest of the
fine-tuning epochs. To do so, we uniformly divide the resource constraint
range of our stitches into several intervals and sample an interval $\gamma$
with uniform sampling. Accordingly, we can obtain a subset of stitches
${\mathcal{N}}_{\gamma}$ whose resource constraints belong to $\gamma$ and
their corresponding importance scores ${\mathcal{Q}}_{\gamma}$. We can then
define a categorical distribution based on the normalized importance scores,
_i.e_.,
$\pi({\mathcal{N}}_{\gamma})=\operatorname{softmax}({\mathcal{Q}}_{\gamma})$
and finally get the sampled stitch
$F_{\phi_{n},S_{n}}\sim\pi\left({\mathcal{N}}_{\gamma}\right)$. We show that
our task-specific stitch sampling improves the performance in Section 5.3.

Additionally, with the sampled stitch, SN-Net employs knowledge distillation
with a RegNetY-160 [57] pre-trained teacher to improve its performance.
However, pre-trained teachers on downstream tasks are often unavailable.
Therefore, in each training iteration, we always sample and train the largest
anchor (teacher) and transfer its knowledge to the sampled stitch (student)
similar to inplace distillation [83]. In practice, we employ the hard-label
distillation [28, 66].

Reuse stitch importance scores for deployment. After fine-tuning, with the set
of accumulated importance scores ${\mathcal{Q}}_{\gamma}$ for all stitches, we
can directly use them to select the stitches with the highest scores to
deploy, which eliminates the need for the costly evaluation of the final stage
of the straightforward SN-Net adaptation solution. We empirically observe that
the deployed stitches are mostly on the Pareto frontier in Section 5.3,
suggesting our estimation of the important stitches is accurate. Furthermore,
in Section G of the appendix, we show that the important stitches differ
across different tasks, proving the need for task-specific deployment
strategies, as highlighted in [21, 79, 24]. It’s worth noting that although
the SNIP metric has been widely employed as a one-shot metric to efficiently
estimate network performance [1, 72], we extend it to estimate stitch
importance for better training and selection.

Figure 5: Instruction-following comparison between Stitched LLaMA obtained by
our ESTA and the Alpaca-LoRA 7B fine-tuned with LoRA [30].

##  5 Experiments

###  5.1 Visual Recognition

We evaluate the effectiveness of our method on a total of 25 downstream visual
recognition tasks, including fine-grained visual classification (FGVC)
benchmark, common visual classification benchmark CIFAR-100 [36], and VTAB-1k
[85] benchmark. FGVC benchmark contains NABirds [70], CUB-200-2011 [71],
Stanford Cars [20], Stanford Dogs [34], and Oxford Flowers [53] tasks and the
VTAB-1k benchmark includes 19 tasks in different domains, each of which has
800 training samples. For visual recognition experiments, we stitch
ImageNet-21k pre-trained ViT-Ti/S/B models [18] from [63].

Implementation details. We employ the stitching space and settings (kernel
size as $2$ and stride as $1$) in [54] for downstream visual recognition
tasks. We uniformly divide the resource constraints supported by the stitches
into around 15 intervals and deploy one stitch within each interval. We set
the hyper-parameter $r$ for the rank of weight updating in the self-attention
layers to be $32$, $16$, and $8$ for ViT-Ti, ViT-S, and ViT-B, respectively.
We also set $r=d_{1}//4$ universally for all stitching layers and $\eta$ in
Eq. (3) as 0.9 by grid search. We follow [33] to fine-tune 100 epochs on each
task and select the other hyper-parameters and augmentation methods. We
include more implementation details in Section A of the appendix.

Main results. We compare the performance of our ESTA with SN-Net and the
individual anchors fine-tuned with LoRA [30] on various datasets. The FLOPs-
accuracy curves are visualized in Figures 3 and 4. We report the averaged
results over the three task groups on VTAB-1k following [33]. Overall, the
stitches obtained by our ESTA exhibit smooth FLOPs-accuracy trade-off curves
consistently across all datasets. Notably, the curve for our ESTA outperforms
SN-Net by significant margins with much fewer trainable parameters. We show
that the inferior results of SN-Net on downstream tasks are led by the sub-
optimal designs of fine-tuning all model parameters and a three-stage
adaptation pipeline in Section 5.3.

We also observe that in many cases, the stitches even outperform the anchors
with comparable or lower FLOPs. For instance, our stitches outperform
individually fine-tuned ViT-S/B anchors by clear margins on NABirds,
CUB-200-2011, Stanford Cars, and Stanford Dogs datasets with comparable
computational complexity. This phenomenon is consistent with the observations
made in [83, 55] and we conjecture that weight sharing among the stitches
serves as a strong regularization to improve their generalization. We show
more results on few-shot learning and stitching representative ConvNext model
family [48] in Section C of the appendix.

Table 2: Relative response quality to Alpaca-LoRA 7B [65] on Vicuna Bench
[11]. Our Stitched LLaMA successfully interpolates the answer quality between
Alpaca-LoRA 7B and 13B.

| Alpaca-LoRA | Stitched LLaMA
---|---|---
Parameter # (B) | 6.7 | 13.0 | 6.9 | 7.8 | 11.5 | 12.5 | 13.0
Quality (%) | 100 | 123 | 113 | 114 | 115 | 120 | 123

###  5.2 Instruction-following Task

We also evaluate the instruction-following capability when stitching the LLaMA
family [67] with our ESTA framework. We adapt and fine-tune LLaMA-7B/13B on
the 52K instruction-following dataset from [65] to obtain chatbot stitches of
varying sizes. The 52K instruction-following data is generated from 175 human-
written instruction-output pairs with [74]. Each sample in the dataset is
formulated by a task description, the context of the task, and the answer that
is generated by GPT-3.5 (text-davinci-003) [8].

Implementation details. We employ the stitching space and settings (kernel
size as $1$ and stride as $1$) in [54] for the instruction-following task and
introduce 40 stitches. We uniformly divide the resource constraints supported
by the stitches into 6 intervals and deploy one stitch within each interval.
We set $\eta$ in Eq. (3) as 0.9. The rank $r$ is set to $16$ and $d_{1}//4$
for LoRA modules in self-attention and stitching layers, respectively. We
employ Adam optimizer and fine-tune with batch size 128 for 10 epochs with
gradient accumulation. During the generation stage, we employ top-$p$ sampling
to sample from the top subword candidates with the probability $p=0.75$,
conduct the beam search, and set the beam size as $4$.

Results. We term the stitches obtained by our ESTA framework as Stitched LLaMA
and compare them with Alpaca-LoRA 7B and 13B that respectively freeze the
LLaMA 7B and 13B backbones and update the self-attention weights with LoRA
modules [30]. For quantitative results reported in Table 2, we observe that
similar to the visual recognition tasks, our Stitched LLaMA with different
numbers of parameters generates answers with plausible quality in-between
Alpaca-LoRA 7B and 13B. Figure 5 gives qualitative comparisons. Our Stitched
LLaMAs produce reasonable responses, suggesting the potential of our ESTA for
adapting large language models (LLMs) for more downstream language tasks. More
examples can be found in Figures 13 and 14 in the appendix.

###  5.3 Ablation Study

Figure 6: Ablation studies with results averaged over five FGVC tasks. (a)
Effect of our stitch-agnostic LoRA modules. “Full FT” represents full fine-
tuning of all the model parameters. “PST w/ full SL” employs fully fine-tuned
stitching layers. We also show the number of trainable parameters. (b) Effect
of our stitch-specific bias terms. “SSB” and “scaled r” represent stitch-
specific bias terms and scaling the rank hyperparameter to reach 4.6M
trainable parameters, respectively. (c) Effect of our task-specific stitch
sampling. “Uniform, all” and “Uniform, best” represent all the stitches fine-
tuned with uniform sampling and their best ones in each resource constraint
interval on the Pareto frontier, respectively. (d) Effect of our strategy to
simultaneously adapt and stitch anchors. Our strategy “Adapt-and-stitch” takes
100 epochs for fine-tuning. In contrast, “Adapt-then-stitch”, as a
straightforward approach to apply SN-Net, first individually adapts each
anchor for a total of $300$ epochs, and then fine-tunes SN-Net for another 100
epochs.

Effect of parameter-efficient stitch fine-tuning. We investigate the
effectiveness of our PST on five FGVC tasks and visualize the averaged results
in Figure 6 (a) and (b). To recall, our PST keeps the majority of the model
parameters frozen while introducing trainable stitch-agonistic low-rank
decomposition matrices and stitch-specific bias terms. In Figure 6 (a), we
observe that when employing LoRA to the multi-head self-attention layers,
there are solid performance gains on all FLOPs ranges with significantly fewer
trainable parameters, echoing the observations in [33, 43]. It is suggested
that low-rank weight updates is a powerful alternative to full fine-tuning for
stitching and adapting anchors and we speculate that low-rank weight update
avoids overfitting under limited downstream data. We further employ low-rank
weight updates to the stitching layers and save 2.4M more trainable parameters
with a slight overall performance drop. We conjecture that the stitching
layers for large pre-trained models are also in low intrinsic dimensions [2]
and least-squares initialization [54] already provide a good initialization
for them.

Figure 6 (b) shows that employing stitch-specific bias terms achieves better
performance with affordable 1.4M extra trainable parameters. We also include
the baseline that simply scales the low-rank dimension $r$ to reach the same
number of trainable parameters as our PST. This baseline is inferior to our
PST. We speculate that the stitch-specific bias terms enable flexible
adjustments of feature representations for different stitches, leading to
enhanced performance.

Effect of task-specific stitch sampling. We investigate the effectiveness of
our sampling strategy (introduced in Section 4.2) on the five FGVC tasks. The
averaged results are visualized in Figure 6 (c). We compare our task-specific
stitch sampling with the uniform one that is adopted in [54]. We empirically
find that for uniform sampling, there is a clear performance gap between the
best stitches that will be deployed and the others that will be dropped after
training, suggesting the need for concentrating on the important stitches
during training. To this end, our task-specific sampling strategy yields
better performance than uniform sampling, indicating that our sampling
strategy selects accurate task-specific important stitches. Most importantly,
our strategy circumvents the need for the costly evaluation stage in the
direct solution of adapting SN-Net for downstream tasks.

Effect of simultaneously adapting and stitching anchors. We study the effect
of simultaneously adapting and stitching anchors strategy (introduced in
Section 4.2) on five FGVC tasks. The averaged results are visualized in Figure
6 (d). We compare our approach with the direct SN-Net adaptation that first
adapts each anchor individually for a total of 300 epochs before stitching
them under the exact setting as ESTA. It can be seen that the direct SN-Net
adaptation yields poor results in task adaptation. Fine-tuned individual
anchors, especially the larger ones, are likely to overfit the limited
downstream data [56, 46]. Accordingly, we speculate that the adapted anchor
weights serve as a sub-optimal initialization for the stitches in the context
of model stitching. On the other hand, our approach achieves better
performance and greatly saves the fine-tuning cost.

Figure 7: Left: effect of different PEFT technique choices. Right: effect of
the rank hyper-parameter $r$ in LoRA modules of self-attention layers, for
which the three values represent the ranks in ViT-Ti, ViT-S, and ViT-B,
respectively. The results are averaged on five FGVC tasks.

Effect of the PEFT technique choices. We compare the effect of different PEFT
technique choices and visualize the averaged results on the five FGVC tasks in
Figure 7 left. We experiment to employ other PEFT techniques Adapter [29] and
Adaptformer [10]. We follow their default setting to respectively insert
bottleneck structures sequentially and in parallel to the feed-forward layers
and fix the number of trainable parameters the same as ours for a fair
comparison. We observe that the overall performance for different PEFT
techniques is similar, suggesting that our ESTA framework is compatible with
more PEFT technique choices. Since low-rank weight updates of LoRA [30] can be
merged into the backbone after training without extra inference computational
cost, we employ LoRA by default.

Effect of the low-rank hyper-parameter $r$. We investigate the effect of
different $r$ which controls the rank when updating weights in self-attention
layers on the five FGVC tasks. The averaged results on the five FGVC tasks are
visualized in Figure 7 right. We find that setting $r$ to be different values
barely has any effect on the overall performance, suggesting our ESTA
framework is robust to the rank hyper-parameter $r$. Since setting $r$ to be
$32$, $16$, and $8$ for low-rank updating the self-attention weights in ViT-
Ti, ViT-S, and ViT-B has slightly higher averaged performance among the
stitches, we make it our default setting.

##  6 Conclusion and Future Work

In this paper, we have introduced a novel task adaptation framework to cheaply
obtain a palette of fine-tuned networks via model stitching, supporting
diverse efficiency-performance tradeoffs at runtime. Specifically, built on
SN-Net, we have tailored a parameter-efficient stitch fine-tuning method,
which learns lightweight stitch-agnostic LoRA modules and stitch-specific bias
terms while keeping the majority of the parameters frozen. Our design
significantly reduces fine-tuning memory and storage costs for downstream task
adaptation. Moreover, we have devised a task-specific stitch sampling strategy
to assign higher sampling probability to the important stitches during
training, which simultaneously improves the Pareto frontiers and avoids a
costly evaluation stage. Extensive experiments on 25 downstream visual
recognition tasks and the instruction-following task have demonstrated the
effectiveness of our proposed framework.

Limitations. Due to the constraints of computational resources, our
experiments are limited to visual recognition and instruction-tuning tasks. In
the future, we will explore adapting pre-trained model families to dense
prediction and multi-modal tasks.

## References

  * Abdelfattah et al. [2021] Mohamed S. Abdelfattah, Abhinav Mehrotra, Lukasz Dudziak, and Nicholas D. Lane.  Zero-cost proxies for lightweight nas.  In _ICLR_ , 2021.
  * Aghajanyan et al. [2021] Armen Aghajanyan, Luke Zettlemoyer, and Sonal Gupta.  Intrinsic dimensionality explains the effectiveness of language model fine-tuning.  In _ACL_ , 2021.
  * Ammar et al. [2014] Haitham Bou Ammar, Eric Eaton, Matthew E Taylor, Decebal Constantin Mocanu, Kurt Driessens, Gerhard Weiss, and Karl Tuyls.  An automated measure of mdp similarity for transfer in reinforcement learning.  In _AAAIW_ , 2014.
  * Andreassen et al. [2021] Anders Andreassen, Yasaman Bahri, Behnam Neyshabur, and Rebecca Roelofs.  The evolution of out-of-distribution robustness throughout fine-tuning.  _arXiv preprint arXiv:2106.15831_ , 2021.
  * Bansal et al. [2021] Yamini Bansal, Preetum Nakkiran, and Boaz Barak.  Revisiting model stitching to compare neural representations.  _NeurIPS_ , 34:225–236, 2021.
  * Bao et al. [2022] Hangbo Bao, Li Dong, Songhao Piao, and Furu Wei.  Beit: Bert pre-training of image transformers.  In _ICLR_ , 2022.
  * Bolya et al. [2021] Daniel Bolya, Rohit Mittapalli, and Judy Hoffman.  Scalable diverse model selection for accessible transfer learning.  _NeurIPS_ , 34:19301–19312, 2021.
  * Brown et al. [2020] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al.  Language models are few-shot learners.  _NeurIPS_ , 33:1877–1901, 2020.
  * Cai et al. [2020] Han Cai, Chuang Gan, Tianzhe Wang, Zhekai Zhang, and Song Han.  Once-for-all: Train one network and specialize it for efficient deployment.  In _ICLR_ , 2020.
  * Chen et al. [2022] Shoufa Chen, Chongjian Ge, Zhan Tong, Jiangliu Wang, Yibing Song, Jue Wang, and Ping Luo.  Adaptformer: Adapting vision transformers for scalable visual recognition.  _NeurIPS_ , 2022.
  * Chiang et al. [2023] Wei-Lin Chiang, Zhuohan Li, Zi Lin, Ying Sheng, Zhanghao Wu, Hao Zhang, Lianmin Zheng, Siyuan Zhuang, Yonghao Zhuang, Joseph E. Gonzalez, Ion Stoica, and Eric P. Xing.  Vicuna: An open-source chatbot impressing gpt-4 with 90%* chatgpt quality, 2023.
  * Coates and Ng [2011] Adam Coates and Andrew Y Ng.  The importance of encoding versus training with sparse coding and vector quantization.  In _ICML_ , pages 921–928, 2011.
  * Csiszárik et al. [2021] Adrián Csiszárik, Péter Kőrösi-Szabó, Akos Matszangosz, Gergely Papp, and Dániel Varga.  Similarity and matching of neural network representations.  _NeurIPS_ , 34:5656–5668, 2021.
  * Dai and Zhang [2011] Dong Dai and Tong Zhang.  Greedy model averaging.  _NeurIPS_ , 24, 2011.
  * Dai et al. [2021] Xiaoliang Dai, Alvin Wan, Peizhao Zhang, Bichen Wu, Zijian He, Zhen Wei, Kan Chen, Yuandong Tian, Matthew Yu, Peter Vajda, et al.  Fbnetv3: Joint architecture-recipe search using predictor pretraining.  In _CVPR_ , pages 16276–16285, 2021.
  * Deng et al. [2009] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei.  Imagenet: A large-scale hierarchical image database.  In _CVPR_ , pages 248–255, 2009.
  * Dettmers et al. [2022] Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer.  Llm. int8 (): 8-bit matrix multiplication for transformers at scale.  _NeurIPS_ , 2022.
  * Dosovitskiy et al. [2021] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby.  An image is worth 16x16 words: Transformers for image recognition at scale.  In _ICLR_ , 2021.
  * Eaton et al. [2008] Eric Eaton, Marie Desjardins, and Terran Lane.  Modeling transfer relationships between learning tasks for improved inductive transfer.  In _ECML PKDD_ , pages 317–332. Springer, 2008.
  * Gebru et al. [2017] Timnit Gebru, Jonathan Krause, Yilun Wang, Duyun Chen, Jia Deng, and Li Fei-Fei.  Fine-grained car detection for visual census estimation.  In _AAAI_ , 2017.
  * Guo et al. [2019] Yunhui Guo, Honghui Shi, Abhishek Kumar, Kristen Grauman, Tajana Rosing, and Rogerio Feris.  Spottune: transfer learning through adaptive fine-tuning.  In _CVPR_ , pages 4805–4814, 2019.
  * Guo et al. [2020] Zichao Guo, Xiangyu Zhang, Haoyuan Mu, Wen Heng, Zechun Liu, Yichen Wei, and Jian Sun.  Single path one-shot neural architecture search with uniform sampling.  In _ECCV_ , pages 544–560. Springer, 2020.
  * Han et al. [2015] Song Han, Jeff Pool, John Tran, and William Dally.  Learning both weights and connections for efficient neural network.  _NeurIPS_ , 28, 2015.
  * He et al. [2023] Haoyu He, Jianfei Cai, Jing Zhang, Dacheng Tao, and Bohan Zhuang.  Sensitivity-aware visual parameter-efficient tuning.  In _ICCV_ , 2023.
  * He et al. [2022] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick.  Masked autoencoders are scalable vision learners.  In _CVPR_ , pages 16000–16009, 2022.
  * Herdt et al. [2023] Rudolf Herdt, Maximilian Schmidt, Daniel Otero Baguer, Jean Le’Clerc Arrastia, and Peter Maass.  Model stitching and visualization how gan generators can invert networks in real-time.  _arXiv preprint arXiv:2302.02181_ , 2023.
  * Hernandez et al. [2022] Adriano Hernandez, Rumen Dangovski, Peter Y. Lu, and Marin Soljacic.  Model stitching: Looking for functional similarity between representations.  _NeurIPSW_ , 2022.
  * Hinton et al. [2014] Geoffrey Hinton, Oriol Vinyals, and Jeff Dean.  Distilling the knowledge in a neural network.  _NeurIPSW_ , 2014.
  * Houlsby et al. [2019] Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin De Laroussilhe, Andrea Gesmundo, Mona Attariyan, and Sylvain Gelly.  Parameter-efficient transfer learning for nlp.  In _ICML_ , pages 2790–2799, 2019.
  * Hu et al. [2022] Edward J Hu, yelong shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen.  LoRA: Low-rank adaptation of large language models.  In _ICLR_ , 2022.
  * Hubara et al. [2017] Itay Hubara, Matthieu Courbariaux, Daniel Soudry, Ran El-Yaniv, and Yoshua Bengio.  Quantized neural networks: Training neural networks with low precision weights and activations.  _JMLR_ , 18(1):6869–6898, 2017.
  * Izmailov et al. [2018] Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson.  Averaging weights leads to wider optima and better generalization.  In _UAI_ , 2018.
  * Jia et al. [2022] Menglin Jia, Luming Tang, Bor-Chun Chen, Claire Cardie, Serge Belongie, Bharath Hariharan, and Ser-Nam Lim.  Visual prompt tuning.  In _ECCV_ , 2022.
  * Khosla et al. [2011] Aditya Khosla, Nityananda Jayadevaprakash, Bangpeng Yao, and Li Fei-Fei.  Novel dataset for fine-grained image categorization.  In _CVPRW_ , 2011.
  * Kornblith et al. [2019] Simon Kornblith, Jonathon Shlens, and Quoc V Le.  Do better imagenet models transfer better?  In _CVPR_ , pages 2661–2671, 2019.
  * Krizhevsky [2009] Alex Krizhevsky.  Learning multiple layers of features from tiny images.  Technical report, 2009.
  * Kundu et al. [2023] Achintya Kundu, Laura Wynter, Rhui Dih Lee, and Luis Angel Bathen.  Tofa: Transfer-once-for-all.  _arXiv preprint arXiv:2303.15485_ , 2023.
  * LeCun et al. [1989] Yann LeCun, John Denker, and Sara Solla.  Optimal brain damage.  _NeurIPS_ , 2, 1989.
  * Lee et al. [2019] Namhoon Lee, Thalaiyasingam Ajanthan, and Philip Torr.  Snip: Single-shot network pruning based on connection sensitivity.  In _ICLR_ , 2019.
  * Lenc and Vedaldi [2015] Karel Lenc and Andrea Vedaldi.  Understanding image representations by measuring their equivariance and equivalence.  In _CVPR_ , pages 991–999, 2015.
  * Lester et al. [2021] Brian Lester, Rami Al-Rfou, and Noah Constant.  The power of scale for parameter-efficient prompt tuning.  In _EMNLP_ , 2021.
  * Li and Liang [2021] Xiang Lisa Li and Percy Liang.  Prefix-tuning: Optimizing continuous prompts for generation.  In _ACL_ , 2021.
  * Lian et al. [2022] Dongze Lian, Daquan Zhou, Jiashi Feng, and Xinchao Wang.  Scaling & shifting your features: A new baseline for efficient model tuning.  _NeurIPS_ , 2022.
  * Liu et al. [2019] Iou-Jen Liu, Jian Peng, and Alexander G Schwing.  Knowledge flow: Improve upon your teachers.  In _ICLR_ , 2019.
  * Liu et al. [2022a] Jing Liu, Jianfei Cai, and Bohan Zhuang.  Focusformer: Focusing on what we need via architecture sampler.  _arXiv preprint arXiv:2208.10861_ , 2022a.
  * Liu et al. [2023] Pengfei Liu, Weizhe Yuan, Jinlan Fu, Zhengbao Jiang, Hiroaki Hayashi, and Graham Neubig.  Pre-train, prompt, and predict: A systematic survey of prompting methods in natural language processing.  _ACM Comput. Surv._ , 2023.
  * Liu et al. [2022b] Yen-Cheng Liu, Chih-Yao Ma, Junjiao Tian, Zijian He, and Zsolt Kira.  Polyhistor: Parameter-efficient multi-task adaptation for dense vision tasks.  _NeurIPS_ , 2022b.
  * Liu et al. [2022c] Zhuang Liu, Hanzi Mao, Chao-Yuan Wu, Christoph Feichtenhofer, Trevor Darrell, and Saining Xie.  A convnet for the 2020s.  In _CVPR_ , pages 11976–11986, 2022c.
  * Loshchilov and Hutter [2017] Ilya Loshchilov and Frank Hutter.  Decoupled weight decay regularization.  _arXiv preprint arXiv:1711.05101_ , 2017.
  * Lu et al. [2021] Zhichao Lu, Gautam Sreekumar, Erik Goodman, Wolfgang Banzhaf, Kalyanmoy Deb, and Vishnu Naresh Boddeti.  Neural architecture transfer.  _TPAMI_ , 43(9):2971–2989, 2021.
  * Matena and Raffel [2022] Michael S Matena and Colin A Raffel.  Merging models with fisher-weighted averaging.  _NeurIPS_ , 35:17703–17716, 2022.
  * Nguyen et al. [2020] Cuong Nguyen, Tal Hassner, Matthias Seeger, and Cedric Archambeau.  Leep: A new measure to evaluate transferability of learned representations.  In _ICML_ , pages 7294–7305. PMLR, 2020.
  * Nilsback and Zisserman [2008] Maria-Elena Nilsback and Andrew Zisserman.  Automated flower classification over a large number of classes.  In _ICVGIP_ , pages 722–729. IEEE, 2008.
  * Pan et al. [2023] Zizheng Pan, Jianfei Cai, and Bohan Zhuang.  Stitchable neural networks.  In _CVPR_ , 2023.
  * Pham et al. [2018] Hieu Pham, Melody Guan, Barret Zoph, Quoc Le, and Jeff Dean.  Efficient neural architecture search via parameters sharing.  In _ICML_ , pages 4095–4104. PMLR, 2018.
  * Plested and Gedeon [2022] Jo Plested and Tom Gedeon.  Deep transfer learning for image classification: a survey.  _arXiv preprint arXiv:2205.09904_ , 2022.
  * Radosavovic et al. [2020] Ilija Radosavovic, Raj Prateek Kosaraju, Ross Girshick, Kaiming He, and Piotr Dollár.  Designing network design spaces.  In _CVPR_ , pages 10428–10436, 2020.
  * Rombach et al. [2022] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer.  High-resolution image synthesis with latent diffusion models.  In _CVPR_ , pages 10684–10695, 2022.
  * Russakovsky et al. [2015] Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al.  Imagenet large scale visual recognition challenge.  _IJCV_ , 115:211–252, 2015.
  * Rusu et al. [2016] Andrei A Rusu, Neil C Rabinowitz, Guillaume Desjardins, Hubert Soyer, James Kirkpatrick, Koray Kavukcuoglu, Razvan Pascanu, and Raia Hadsell.  Progressive neural networks.  _arXiv preprint arXiv:1606.04671_ , 2016.
  * Shu et al. [2021] Yang Shu, Zhi Kou, Zhangjie Cao, Jianmin Wang, and Mingsheng Long.  Zoo-tuning: Adaptive transfer from a zoo of models.  In _ICML_ , pages 9626–9637. PMLR, 2021.
  * Shu et al. [2022] Yang Shu, Zhangjie Cao, Ziyang Zhang, Jianmin Wang, and Mingsheng Long.  Hub-pathway: Transfer learning from a hub of pre-trained models.  _NeurIPS_ , 2022.
  * Steiner et al. [2022] Andreas Steiner, Alexander Kolesnikov, Xiaohua Zhai, Ross Wightman, Jakob Uszkoreit, and Lucas Beyer.  How to train your vit? data, augmentation, and regularization in vision transformers.  _TMLR_ , 2022.
  * Sung et al. [2022] Yi-Lin Sung, Jaemin Cho, and Mohit Bansal.  Lst: Ladder side-tuning for parameter and memory efficient transfer learning.  _NeurIPS_ , 2022.
  * Taori et al. [2023] Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto.  Stanford alpaca: An instruction-following llama model.  https://github.com/tatsu-lab/stanford_alpaca, 2023.
  * Touvron et al. [2021] Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou.  Training data-efficient image transformers & distillation through attention.  In _ICML_ , pages 10347–10357. PMLR, 2021.
  * Touvron et al. [2023a] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al.  Llama: Open and efficient foundation language models.  _arXiv preprint arXiv:2302.13971_ , 2023a.
  * Touvron et al. [2023b] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al.  Llama 2: Open foundation and fine-tuned chat models.  _arXiv preprint arXiv:2307.09288_ , 2023b.
  * Tran et al. [2019] Anh T Tran, Cuong V Nguyen, and Tal Hassner.  Transferability and hardness of supervised classification tasks.  In _ICCV_ , pages 1395–1405, 2019.
  * Van Horn et al. [2015] Grant Van Horn, Steve Branson, Ryan Farrell, Scott Haber, Jessie Barry, Panos Ipeirotis, Pietro Perona, and Serge Belongie.  Building a bird recognition app and large scale dataset with citizen scientists: The fine print in fine-grained dataset collection.  In _CVPR_ , pages 595–604, 2015.
  * Wah et al. [2011] Catherine Wah, Steve Branson, Peter Welinder, Pietro Perona, and Serge Belongie.  The caltech-ucsd birds-200-2011 dataset.  _Tech. Rep. CNS-TR-2011-001, California Institute of Technology_ , 2011.
  * Wang et al. [2020] Chaoqi Wang, Guodong Zhang, and Roger Grosse.  Picking winning tickets before training by preserving gradient flow.  In _ICLR_ , 2020.
  * Wang et al. [2021] Dilin Wang, Meng Li, Chengyue Gong, and Vikas Chandra.  Attentivenas: Improving neural architecture search via attentive sampling.  In _CVPR_ , pages 6418–6427, 2021.
  * Wang et al. [2022] Yizhong Wang, Yeganeh Kordi, Swaroop Mishra, Alisa Liu, Noah A Smith, Daniel Khashabi, and Hannaneh Hajishirzi.  Self-instruct: Aligning language model with self generated instructions.  _arXiv preprint arXiv:2212.10560_ , 2022.
  * Wolf et al. [2020] Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al.  Transformers: State-of-the-art natural language processing.  In _EMNLP_ , pages 38–45, 2020.
  * Wortsman et al. [2022a] Mitchell Wortsman, Gabriel Ilharco, Samir Ya Gadre, Rebecca Roelofs, Raphael Gontijo-Lopes, Ari S Morcos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith, et al.  Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time.  In _ICML_ , pages 23965–23998. PMLR, 2022a.
  * Wortsman et al. [2022b] Mitchell Wortsman, Gabriel Ilharco, Jong Wook Kim, Mike Li, Simon Kornblith, Rebecca Roelofs, Raphael Gontijo Lopes, Hannaneh Hajishirzi, Ali Farhadi, Hongseok Namkoong, et al.  Robust fine-tuning of zero-shot models.  In _CVPR_ , pages 7959–7971, 2022b.
  * Xu et al. [2021a] Hang Xu, Ning Kang, Gengwei Zhang, Chuanlong Xie, Xiaodan Liang, and Zhenguo Li.  Nasoa: Towards faster task-oriented online fine-tuning with a zoo of models.  In _ICCV_ , pages 5097–5106, 2021a.
  * Xu et al. [2021b] Runxin Xu, Fuli Luo, Zhiyuan Zhang, Chuanqi Tan, Baobao Chang, Songfang Huang, and Fei Huang.  Raise a child in large language model: Towards effective and generalizable fine-tuning.  In _EMNLP_ , 2021b.
  * Yang et al. [2022] Xingyi Yang, Daquan Zhou, Songhua Liu, Jingwen Ye, and Xinchao Wang.  Deep model reassembly.  _NeurIPS_ , 35:25739–25753, 2022.
  * Yosinski et al. [2014] Jason Yosinski, Jeff Clune, Yoshua Bengio, and Hod Lipson.  How transferable are features in deep neural networks?  _NeurIPS_ , 27, 2014.
  * You et al. [2021] Kaichao You, Yong Liu, Jianmin Wang, and Mingsheng Long.  Logme: Practical assessment of pre-trained models for transfer learning.  In _ICML_ , pages 12133–12143. PMLR, 2021.
  * Yu and Huang [2019] Jiahui Yu and Thomas S Huang.  Universally slimmable networks and improved training techniques.  In _ICCV_ , pages 1803–1811, 2019.
  * Zaken et al. [2022] Elad Ben Zaken, Yoav Goldberg, and Shauli Ravfogel.  Bitfit: Simple parameter-efficient fine-tuning for transformer-based masked language-models.  In _ACL_ , pages 1–9, 2022.
  * Zhai et al. [2019] Xiaohua Zhai, Joan Puigcerver, Alexander Kolesnikov, Pierre Ruyssen, Carlos Riquelme, Mario Lucic, Josip Djolonga, Andre Susano Pinto, Maxim Neumann, Alexey Dosovitskiy, et al.  A large-scale study of representation learning with the visual task adaptation benchmark.  _arXiv preprint arXiv:1910.04867_ , 2019.
  * Zhang and Agrawala [2023] Lvmin Zhang and Maneesh Agrawala.  Adding conditional control to text-to-image diffusion models.  In _ICCV_ , 2023.
  * Zhang et al. [2023] Renrui Zhang, Jiaming Han, Aojun Zhou, Xiangfei Hu, Shilin Yan, Pan Lu, Hongsheng Li, Peng Gao, and Yu Qiao.  Llama-adapter: Efficient fine-tuning of language models with zero-init attention.  _arXiv preprint arXiv:2303.16199_ , 2023.
  * Zhang et al. [2022] Yuanhan Zhang, Kaiyang Zhou, and Ziwei Liu.  Neural prompt search.  _arXiv preprint arXiv:2206.04673_ , 2022.
  * Zhou et al. [2022] Kaiyang Zhou, Jingkang Yang, Chen Change Loy, and Ziwei Liu.  Conditional prompt learning for vision-language models.  In _CVPR_ , pages 16816–16825, 2022.

Appendix

We organize the appendix as follows.

  * •

In Section A, we introduce more details about data augmentation and other
hyper-parameters.

  * •

In Section B, we show few-shot learning performance comparisons with SN-Net on
Stanford Cars [20].

  * •

In Section C, we adapt our ESTA to convnets and show the performance
comparisons with SN-Net and anchors individually fine-tuned with LoRA [30].

  * •

In Section D, we investigate the effectiveness of our distillation strategy.

  * •

In Section E, we study the effect of the momentum coefficient $\eta$.

  * •

In Section F, we analyze why our strategy of simultaneously stitching and
adapting anchors leads to better performance.

  * •

In Section G, we investigate whether the deployed task-specific stitches are
the same across three sampled tasks: Oxford Flowers [53], Stanford Cars [20],
and CUB-200-2011 [71].

  * •

In Figure 12, we show more visualizations for distributions of pair-wise
gradient angles among stitches when updating shared weights similar to Figure
2 of the main paper.

  * •

In Table 3, we show quantitative results from Figures 3 and 4 of the main
paper.

  * •

In Section H, we show more qualitative results on the instruction-following
task.

##  Appendix A Data Augmentation and Other Hyper-parameters

For our ESTA and the baseline methods, we follow [33] for a standard data
augmentation pipeline with resize, random crop to 224$\times$224 and random
horizontal flip for the FGVC and CIFAR-100 datasets, and only resize to
224$\times$224 for VTAB-1k datasets. Following [88], we use AdamW optimizer
[49] with a cosine scheduler and a learning rate warm-up period of 10 epochs.
We set weight decay as $1\times 10^{-4}$ for all visual recognition datasets.
We choose the base learning rate with a grid search on the validation set over
{$1\times 10^{-4}$, $2\times 10^{-4}$, $1\times 10^{-3}$, $2\times 10^{-3}$}
following the split of [33]. We set the learning rate as $base\\_lr\times
b/256$, where $b$ is the batch size. All experiments are conducted on NVIDIA
A100 GPUs. Due to the constraints of computational resources, we use 8-bit
matrix multiplication [17] for the frozen weights in the instruction-following
task during fine-tuning and generation.

Figure 8: Few-shot learning performance comparisons with SN-Net [54] for
adapting ViT-Ti/S/B pre-trained on ImageNet-22k [16] to Stanford Cars [20]. We
denote individually fine-tuned anchors as yellow stars and also show the
number of trainable parameters.

##  Appendix B Results on Few-shot Learning

Following [89], we conduct experiments with limited samples (1, 4, 8, 16
shots) on Stanford Cars task [20]. The results are visualized in Figure 8. We
have similar observations as those in Section 5.1 of the main paper that our
ESTA 1) exhibits smooth FLOPs-accuracy curves; 2) generally outperforms SN-Net
by large margins; and 3) achieves comparable or even better performance than
individually fine-tuned anchors, which we conjecture due to model stitching’s
strong weight-sharing regularization. This suggests that our ESTA works well
in data-efficient adaptation.

Figure 9: (a) Comparisons with SN-Net [54] for adapting ConvNext-T/S/B [48]
pre-trained on ImageNet-22k [16]. We denote individually fine-tuned anchors as
yellow stars and also show the number of trainable parameters. (b) Effect of
inplace distillation. (c) Effect of the momentum coefficient $\eta$ for
updating the importance scores. We show averaged results on five FGVC datasets
for (a), (b), and (c).

##  Appendix C Apply ESTA to Convnets

We employ our ESTA to adapt the popular convnet architecture ConvNext-T/S/B
[48] pre-trained on ImageNet-22k. The averaged results on five FGVC datasets
are visualized in Figure 9 (a). We have similar observations as those in
Figures 3 and 4 of the main paper that our ESTA obtains stitches with a smooth
FLOPs-accuracy curve and outperforms SN-Net by large margins. This indicates
that our ESTA is not restricted to transformer-based architectures.

##  Appendix D Effect of Inplace Distillation

We investigate the effect of inplace distillation that is introduced in
Section 4.2 of the main paper. The averaged results on five FGVC tasks are
visualized in Figure 9 (b). We observe that inplace distillation contributes
to the superiority of our proposed ESTA framework for obtaining stitches of
good performance. However, its contribution is not as significant as the key
designs of ESTA, _i.e_., our parameter-efficient stitch fine-tuning method, a
one-stage deployment pipeline with a task-specific stitch sampling strategy as
shown in Figure 6 of the main paper.

##  Appendix E Effect of $\eta$

As introduced in Section 4.2 of the main paper, $\eta$ is the momentum
coefficient for updating the importance scores of the stitches. We investigate
the effect of $\eta$ and visualize the averaged results on five FGVC datasets
in Figure 9 (c). We empirically find that setting $\eta$ to 0.9 and 0.99
achieves slightly better results than setting $\eta$ to 0.8. We conjecture
that a small value of $\eta$ makes the sampling process unstable during an
early training stage after the short warm-up period. Therefore, we set $\eta$
to 0.9 as the default setting.

Figure 10: Fine-tuning cross-entropy loss and testing accuracy with respect to
the fine-tuning epochs averaged over five FGVC datasets. We show the testing
accuracy for 25, 50, and 75 epochs. Our strategy “Adapt-and-stitch” takes 100
epochs for fine-tuning. In contrast, “Adapt-then-stitch”, as a straightforward
approach to apply SN-Net, first individually adapts each anchor for a total of
$300$ epochs, and then fine-tunes SN-Net for another 100 epochs.

##  Appendix F Why our “Adapt-and-stitch” leads to better performance?

As shown in Section 5.3 and Figure 6 (d) of the main paper, our strategy to
simultaneously adapt and stitch anchors achieves significant performance gain
from SN-Net’s straightforward strategy to first adapt and then stitch anchors.
We speculate that anchors are likely to overfit during fine-tuning in SN-Net’s
anchor adaptation stage and lose their generalization capability [77, 4].
Accordingly, stitching their weights in SN-Net may not generalize well to
unseen samples. To verify our speculation, we compare the training cross-
entropy loss and testing accuracy with respect to the training epoch. In
Figure 10, we observe that SN-Net’s straightforward “Adapt-then-stitch”
strategy exhibits lower cross-entropy loss throughout training, but the
testing accuracy barely improves after epoch 50. This observation aligns with
our speculation and suggests the improved generalization capability of our
“Adapt-and-stitch” pipeline.

Figure 11: The patterns of the deployed stitches selected by our task-specific
sampling strategy on Oxford Flowers [53], Stanford Cars [20], and CUB-200-2011
[71] datasets. We visualize the number of blocks from anchors to form these
stitches. Figure 12: More visualizations for distributions of pair-wise
gradient angles among stitches when updating shared weights. We show
visualizations for ImageNet-1k [59], Stanford Cars [20], CUB-200-2011 [71],
and NABirds [70] datasets. Angle 90∘ is highlighted with a dashed red line.
For simplicity, we combine the gradients for query, key, and value projections
and show their angles for a total of 32 stitches when stitching ViT-Ti and
ViT-S anchors. Similar to Figure 2 of the main paper, the gradient angles are
larger in the downstream datasets than in the source domain ImageNet-1k [59].
Figure 13: Instruction-following comparison between Stitched LLaMA obtained by
our ESTA and the individually fine-tuned anchors (Alpaca-LoRA 7B and 13B)
[30]. Figure 14: Instruction-following comparison between Stitched LLaMA
obtained by our ESTA and the individually fine-tuned anchors (Alpaca-LoRA 7B
and 13B) [30].

##  Appendix G Are important stitches the same across different tasks?

We investigate whether the important stitches are the same across tasks on
Oxford Flowers [53], Stanford Cars [20], and CUB-200-2011 [71] datasets. The
patterns are visualized in Figure 11. We observe that the selected important
stitches are different on these visual recognition tasks, which verifies our
motivation for deploying task-specific stitches and assigning higher sampling
probabilities to them.

##  Appendix H More Qualitative Evaluation on Instruction-following Task

We show more qualitative comparisons between our Stitched LLaMA and the
individually fine-tuned Alpaca-LoRA 7B and 13B in Figures 13 and 14. We
empirically find that the quality of the responses for our Stitched LLaMA
improves with more parameters, which aligns with our findings in Sections 5.1
and 5.2 of the main paper. For instance, in Figure 13, both Alpaca-LoRA 7B and
Stitched LLaMA 8B tend to generate repetitive responses; however, Stitched
LLaMA with higher numbers of parameters produce more coherent and informative
responses. In addition, as shown in Figure 14, as the total parameters of the
stitches increase, the response content becomes richer and more detailed.

Table 3: Quantitative results from Figures 3 and 4 of the main paper.

| Stanford Cars [20]
---|---
FLOPs (G) | 2.48 | 3.31 | 4.41 | 5.85 | 6.93 | 8.39 | 9.47 | 10.54 | 11.62 | 13.77 | 14.16 | 15.54 | 16.61 | 17.39 | 17.67
Accuracy (%) | 82.0 | 84.37 | 85.25 | 84.87 | 85.4 | 86.11 | 86.22 | 86.83 | 86.99 | 86.15 | 87.58 | 87.68 | 87.9 | 87.28 | 87.2
| CUB-200-2011 [71]
FLOPs (G) | 2.86 | 3.97 | 4.52 | 5.85 | 6.93 | 8.00 | 9.08 | 10.54 | 11.62 | 13.77 | 14.46 | 15.93 | 17.00 | 17.39 | 17.67
Accuracy (%) | 81.29 | 82.86 | 83.62 | 85.54 | 84.36 | 85.43 | 86.05 | 86.71 | 86.88 | 88.06 | 87.76 | 88.44 | 88.16 | 87.88 | 88.47
| Stanford Dogs [34]
FLOPs (G) | 2.2 | 3.53 | 4.36 | 5.85 | 6.93 | 8.39 | 9.08 | 10.93 | 12.01 | 13.39 | 14.16 | 15.24 | 16.31 | 17.39 | 17.67
Accuracy (%) | 76.17 | 78.72 | 79.98 | 81.83 | 81.18 | 82.19 | 82.95 | 83.61 | 85.02 | 87.65 | 87.73 | 88.51 | 88.64 | 88.47 | 89.13
| NABirds [70]
FLOPs (G) | 3.03 | 3.80 | 4.63 | 5.85 | 6.93 | 8.00 | 9.85 | 10.54 | 12.01 | 13.08 | 14.46 | 15.24 | 16.31 | 17.39 | 17.67
Accuracy (%) | 77.23 | 78.28 | 79.15 | 81.71 | 80.79 | 81.85 | 83.22 | 83.65 | 84.09 | 84.44 | 84.2 | 85.12 | 85.19 | 85.14 | 85.06
| VTAB [85]
FLOPs (G) | 2.02 | 3.03 | 4.26 | 4.53 | 5.85 | 7.2 | 8.57 | 9.69 | 10.76 | 11.98 | 13.1 | 14.12 | 15.63 | 16.59 | 17.57
Accuracy (%) | 67.25 | 69.56 | 71.61 | 71.24 | 73.11 | 73.19 | 73.32 | 73.08 | 73.19 | 73.43 | 74.13 | 74.27 | 74.7 | 74.83 | 75.1
| CIFAR-100 [36]
FLOPs (G) | 2.35 | 3.08 | 4.37 | 4.65 | 5.83 | 8.76 | 9.84 | 10.92 | 12.0 | 12.7 | 13.08 | 14.16 | 15.24 | 15.94 | 17.67
Accuracy (%) | 86.6 | 88.65 | 90.92 | 90.91 | 90.72 | 91.18 | 90.93 | 91.03 | 91.56 | 91.33 | 92.15 | 92.12 | 92.59 | 92.55 | 93.08


# Uncertainty Disentanglement with Non-stationary Heteroscedastic Gaussian
Processes for Active Learning

Zeel B Patel
IIT Gandhinagar, India  Nipun Batra
IIT Gandhinagar, India  Kevin Murphy
Google, USA

###### Abstract

Gaussian processes are Bayesian non-parametric models used in many areas. In
this work, we propose a Non-stationary Heteroscedastic Gaussian process model
which can be learned with gradient-based techniques. We demonstrate the
interpretability of the proposed model by separating the overall uncertainty
into aleatoric (irreducible) and epistemic (model) uncertainty. We illustrate
the usability of derived epistemic uncertainty on active learning problems. We
demonstrate the efficacy of our model with various ablations on multiple
datasets.

##  1 Introduction

Gaussian processes (GPs) are Bayesian non-parametric models useful for many
real-world regression and classification problems. The key object required to
define a GP is the kernel function
$\mathcal{K}(\mathbf{x},\mathbf{x}^{\prime})$, which measures the similarity
of the input points. A common choice is the RBF kernel
$\mathcal{K}(\mathbf{x},\mathbf{x}^{\prime};\mathbf{\theta})=\sigma\exp\left(-\frac{||\mathbf{x}-\mathbf{x}^{\prime}||^{2}}{2\ell^{2}}\right)$
where $\ell$ is the length scale, and $\sigma$ is the signal variance. In
regression problems, we also often have observation noise with variance
$\omega^{2}$. These three hyper-parameters, $\theta=(\ell,\sigma,\omega)$, are
often learned by optimizing the negative log marginal likelihood.

However, this model uses a stationary kernel (depended only on the distance
between locations) and homoskedastic noise (constant noise variance
($\omega^{2}$)), and these assumptions might not hold in real-life
applications such as environment modeling [1, 2]. In particular, non-
stationary kernels are necessary if the similarity of two inputs may depend on
their location in the input space. Similarly, heteroskedastic noise may be
necessary if the quality of the measurements may vary across space.

In this short paper, we provide a computationally efficient way to create GPs
with non-stationary kernels, and heteroskedastic noise, by using a Gibbs
kernel [3]. Gibbs kernel can be considered the generalization of the RBF
kernel where the hyper-parameters are input-dependent, i.e.,
$\theta(\mathbf{x})=(\ell(\mathbf{x}),\sigma(\mathbf{x}),\omega(\mathbf{x}))$.
These three hyper-parameter functions are themselves represented by a “latent”
GP, as in [4, 5]. In contrast to prior work, which uses EP or HMC, we use
inducing point approximations to speed up the computation of this latent GP
(which is needed to evaluate the kernel). In addition, we show how modeling
variation in all three hyper-parameters allows us to distinguish locations
where the latent function value is uncertain (epistemic uncertainty), as
opposed to locations where the observation noise is high (aleatoric
uncertainty). This distinction is crucial for problems such as active learning
and efficient sensor placement. (c.f. [6]).

##  2 Methods

###  2.1 Non-stationary Heteroscedastic Gaussian processes

Given observations $\mathbf{y}\in\mathbb{R}^{N}$ at inputs
$X\in\mathbb{R}^{N\times D}$, we assume the following model:

$\displaystyle y(\mathbf{x})$
$\displaystyle=f(\mathbf{x})+\varepsilon(\mathbf{x}),\quad\varepsilon(\mathbf{x})\sim\mathcal{N}\left(0,\omega(\mathbf{x})^{2}\right)$
(1) $\displaystyle f(\mathbf{x})$ $\displaystyle\sim
GP(0,\mathcal{K}_{f}(\mathbf{x},\mathbf{x}^{\prime}))$ (2)

where the kernel function is
$\mathcal{K}(\mathbf{x},\mathbf{x}^{\prime})=\text{cov}(f(\mathbf{x}),f(\mathbf{x}^{\prime}))$
and $\varepsilon(\mathbf{x})$ is zero mean noise. We use the following non-
stationary kernel function [3]:

$\displaystyle\mathcal{K}_{f}(\mathbf{x},\mathbf{x}^{\prime})=\sigma(\mathbf{x})\sigma(\mathbf{x}^{\prime})\sqrt{\frac{2\ell(\mathbf{x})\ell(\mathbf{x}^{\prime})}{\ell(\mathbf{x})^{2}+\ell(\mathbf{x}^{\prime})^{2}}}\exp\left(-\frac{||\mathbf{x}-\mathbf{x}^{\prime}||^{2}}{\ell(\mathbf{x})^{2}+\ell(\mathbf{x}^{\prime})^{2}}\right)$
(3)

We assume all hyperparameters of the model,
$\mathbf{\theta}(\mathbf{x})=(\ell(\mathbf{x}),\sigma(\mathbf{x}),\omega(\mathbf{x}))$,
may be input dependent, to allow for non-stationarity and heteroscedasticity
(Previous work [5] has shown that such a kernel is a valid positive semi-
definite kernel). We assume these “hyper-functions” $h(\mathbf{x})$ (where
$h(\cdot)$ represents either $\ell(\cdot)$, $\sigma(\cdot)$ or
$\omega(\cdot)$) are smooth and model them by a latent GP on the log scale:

$\displaystyle\tilde{h}(\mathbf{x})\sim
GP(\mu_{h},\mathcal{K}_{h}(\mathbf{x},\mathbf{x}^{\prime};\mathbf{\phi}_{h}))$
(4)

where $\tilde{h}(\cdot)=\log h(\cdot)$. These latent GPs are characterized by
a constant mean $\mu_{h}$ and RBF kernels with parameters
$\mathbf{\phi}_{h}=(\ell_{h},\sigma_{h})$. (We assume noise-free latent
functions, so $\omega_{h}=0$.)111 In practice, we use a reasonably small value
(jitter) for numerical stability.

To make learning the latent GPs efficient, we will use a set of $M$ (shared)
inducing points $\overline{\mathbf{X}}\in\mathbb{R}^{M\times D}$, which we
treat as additional hyper-parameters. Let
$\overline{\mathbf{z}}_{h}=\tilde{h}(\overline{\mathbf{X}})$ be the (log)
outputs at these locations for hyper-parameter $h$. Then we can infer the
expected hyper-parameter value at any other location $\mathbf{x}$ using the
usual GP prediction formula for the mean:

$\displaystyle\tilde{h}(\mathbf{x})=\mathcal{K}_{h}(\mathbf{x},\overline{\mathbf{X}})\mathcal{K}_{h}(\overline{\mathbf{X}},\overline{\mathbf{X}})^{-1}\overline{\mathbf{z}}_{h}$
(5)

Then we can compute $\mathcal{K}_{f}(\mathbf{x},\mathbf{x}^{\prime})$ at any
pair of inputs using Eq. 3, where $h(\mathbf{x})=e^{\tilde{h}(\mathbf{x})}$.

###  2.2 Learning the hyper-parameters

In total, we can have up to $2MD+2M+9$ parameters: $MD$-dimensional inducing
inputs $\overline{\mathbf{X}}$, $MD$ sized $\ell(\overline{\mathbf{X}})$
(ARD), $M$ sized each
$\\{\sigma(\overline{\mathbf{X}}),\omega(\overline{\mathbf{X}})$} and $9$
latent GP hyper-parameters ($\\{\mu_{h},\mathbf{\phi}_{h}$}).

$\displaystyle\mathbf{\phi}=(\mu_{h},\mathbf{\phi}_{h},\overline{\mathbf{z}}_{h})$
(6)

Let $\mathbf{\phi}^{X}=(\overline{\mathbf{X}},\mathbf{\phi})$ represent all
the model parameters. We can compute a MAP-type II estimate of these
parameters by minimizing the following objective using gradient descent:

$\displaystyle\hat{\mathbf{\phi}}^{X}=\arg\min_{\mathbf{\phi}^{X}}-[\log
p(\mathbf{y}|\mathbf{X},\mathbf{\phi},\overline{\mathbf{X}})+\log
p(\mathbf{\phi})]$ (7)

where $p(\mathbf{y}|\mathbf{X},\mathbf{\phi})$ is the marginal likelihood of
the main GP (integrating out $\mathbf{f}=[f(\mathbf{x}_{n})]_{n=1}^{N}$), and
where $p(\mathbf{\phi})$ is the prior. To ensure smoothly varying latent GPs
with a large length scale and low variance, we use a Gamma(5, 1) prior for
$\ell$ and a Gamma(0.5, 1) prior for $\sigma$. To ease the optimization
process, we use a non-centered parameterization of $\overline{\mathbf{z}}_{h}$
by learning a vector $\boldsymbol{\gamma}_{h}$ with independent values (prior
for $\boldsymbol{\gamma}_{h}$ is $\mathcal{N}(0,1)$) and we then
deterministically compute
$\overline{\mathbf{z}}_{h}=L\boldsymbol{\gamma}_{h}$, where, $L$ is Cholesky
decomposition of
$\mathcal{K}_{h}(\overline{\mathbf{X}},\overline{\mathbf{X}})$. We initialize
$\mathbf{\phi}$ by sampling from their respective priors, and initialize
$\overline{\mathbf{X}}$ by selecting $M$ points from the dataset.

###  2.3 Active learning

Active learning (see e.g., [7]) uses some measure of uncertainty to decide
which points to label so as to learn the underlying function as quickly as
possible. This can also be useful for tasks such as deciding where to place
air quality (AQ) sensors (see e.g., [2]), where the goal is to infer the
underlying AQ values at all spatial locations based on sensors of varying
quality.

Often the GP predictive variance is used as the acquisition function. The
overall predictive variance, $\text{var}(y(\mathbf{x}))$, is equal to the
model or epistemic uncertainty,
$\text{var}(f(\mathbf{x}))=\sigma(\mathbf{x})$, plus the observation noise,
$\varepsilon(\mathbf{x})\sim\mathcal{N}(0,\omega(\mathbf{x})^{2})$, which is
known as aleatoric uncertainty [8]. Aleatoric uncertainty is irreducible, and
thus querying data-points where it is high is not likely to improve the model
(c.f., [9]). Thus, we use the epistemic uncertainty
$\text{var}(f(\mathbf{x}))$ for active learning. (We note that separating
epistemic and aleatoric uncertainty is harder with other kinds of
probabilistic model, such as Bayesian neural networks, c.f., [10, 11]).

##  3 Experiments

In this section, we experimentally compare our method to various baselines on
various datasets.

#### Regression

In Table 1, we compare various ablations of our method on various small
regression tasks. In particular, we consider making one or more of the hyper-
parameters corresponding to length scale ($\ell$), variance ($\sigma$), and
observation noise ($\omega$) either input-dependent (using a latent GP) or
fixed constants to be learned The metrics are averaged over five-fold cross-
validation using the Adam optimizer (rate=0.05, epochs=1000). We use the
following datasets for experiments: 1) Motorcycle helmet: It is a real-life
non-stationary, heteroscedastic dataset used in related studies [5]; 2)
NONSTAT-2D: It is a non-stationary 2D dataset simulated by input-depended
lengthscale [12]. We add linearly increasing heteroscedastic noise to
NONSTAT-2D to make it heteroscedastic; and 3) Jump1D: This is 1D dataset taken
from [5] (mentioned as J in [5]), which is highly discontinuous from the
center. In general, we get better results by making all these parameters
input-dependent. We generate a dataset SYNTH-1D from a generative process of a
GP where all three hyperparameters ($\ell(x),\sigma(x),\omega(x)$) are input-
depended. We use 200 equally spaced inputs in range (-30, 30) with the
following deterministic trends:
$\ell(x)=0.5\sin(x/8)+1.5,\sigma(x)=1.5\exp(\sin(0.2x)),\omega(x)=2.5\log(1+\exp(\sin(-0.2x)))$.
We show in Fig. 1 that we are able recover the trends of underlying
hyperparameter functions. We illustrate the fit of $(\ell,\sigma,\omega)$-GP
over a few more datasets in Appendix Section A Fig. 3.

| Jump | Motorcycle | NONSTAT-2D
---|---|---|---
Model | NLPD | RMSE | NLPD | RMSE | NLPD | RMSE
Stationary Homoskedastic GP | 4.98 | 0.26 | 11.96 | 0.44 | -50.72 | 0.09
$(\ell)$-GP | 5.01 | 0.26 | 11.92 | 0.44 | -65.13 | 0.06
$(\omega)$-GP | 3.82 | 0.22 | 5.21 | 0.44 | -50.81 | 0.09
$(\sigma)$-GP | 0.92 | 0.30 | 11.56 | 0.44 | -56.66 | 0.07
$(\ell,\omega)$-GP | 5.01 | 0.26 | 5.68 | 0.45 | -65.31 | 0.06
$(\ell,\sigma)$-GP | -2.18 | 0.22 | 11.54 | 0.44 | -49.28 | 0.07
$(\sigma,\omega)$-GP | 0.92 | 0.22 | 4.21 | 0.46 | -54.35 | 0.10
$(\ell,\sigma,\omega)$-GP | -2.20 | 0.22 | 4.09 | 0.45 | -73.74 | 0.05
Table 1: Quantitative results on various datasets. Metrics are as follows
(lower is better): Negative Log Predictive Density (NLPD), Root Mean Squared
Error (RMSE). The rows represent different methods in which we either used
fixed or input-dependent length scale $\ell$, signal variance $\sigma$ and
observation noise $\omega$. N-NONSTAT-2D, and other datasets are from [5].
$(\ell,\sigma,\omega)$-GP is the best or the second best across all datasets
and metrics.  Figure 1: Our model fit on SYNTH-1D data (top) and recovered
hyperparameter trends (bottom). Dotted lines show the true functions and
continuous lines show the learned functions.

#### Active learning

In Section 2.3, we argued that $\text{var}(f(\mathbf{x}))$ is a better measure
of uncertainty compared to $\text{var}(y(\mathbf{x}))$, when performing active
learning. We illustrate this point on a 1d synthetic example in Fig. 2. We
first train the model on 30 initial train points and then select 50 points
with active learning with the following acquisition functions: c)
$\text{var}(y(\mathbf{x}))$ and d) $\text{var}(f(\mathbf{x}))$. Note that we
do not retrain the hyperparameters of GP during active learning to fasten the
process. We empirically (b) and visually (c, d) show that points chosen by
$\text{var}(f(\mathbf{x}))$ are closer to the real function.

(a)

(b)

(c)

(d)

Figure 2:  Fit after sampling 50 points from SYNTH-1D with active learning
using (c) overall uncertainty; (d) epistemic uncertainty. We show in (b) that
Mean Squared Error (MAE) between predicted function $f$ and ground truth
$f^{*}$ is improving faster in (d) with epistemic uncertainty as compared to
(c) overall uncertainty. Initial training was done on 30 points. While using
epistemic uncertainty, we can capture better points that help GP learn a
better fit. Black dots in (a) are data points. Green and orange regions show
epistemic and overall uncertainty, respectively.

##  4 Related work

Multiple approaches to create non-stationary kernels include the multivariate
generalization of the Gibbs kernel in [13], sparse spectral kernels [14, 15],
the post-processing method of [16], deep kernel learning [17], and deep GPs
[18]. In this paper, we adopt the Gibbs kernel approach, where the kernel
parameters are themselves generated by latent GPs, as used in prior work [5].
The main difference is in the model fitting procedure. They optimize the
latent hyper-parameters $\mathbf{\phi}$ by grid search, and the compute MAP
estimates of $\tilde{\mathbf{\ell}}=[\tilde{\ell}(\mathbf{x}_{n})]_{n=1}^{N}$,
$\tilde{\mathbf{\sigma}}=[\tilde{\sigma}(\mathbf{x}_{n})]_{n=1}^{N}$, and
$\tilde{\mathbf{\omega}}=[\tilde{\omega}(\mathbf{x}_{n})]_{n=1}^{N}$ using
gradient descent. Thus they are optimizing a total of $3N+9$ parameters,
namely the outputs of the three latent hyper-functions, and the latent hyper-
parameters. By contrast, we optimize $2M(D+1)+9$ parameters, namely the
inducing inputs to the three latent hyper-functions, and the latent hyper-
parameters. For the problems of interest to us (spatio-temporal modeling), $D$
is often low dimensional, so $2M(D+1)\ll 3N$, which makes our approach faster
and less prone to overfitting. An additional advantage of our approach is that
it is easy to make predictions on new data since we can compute the kernel
$\mathcal{K}$ for any pair of inputs. By contrast, in [5], they have to use a
heuristic to extrapolate the latent hyper-functions to the test inputs, before
using them to compute the kernel values at those locations.

## References

  * [1] Paul D Sampson and Peter Guttorp.  Nonparametric estimation of nonstationary spatial covariance structure.  Journal of the American Statistical Association, 87(417):108–119, 1992.
  * [2] Zeel B Patel, Palak Purohit, Harsh M Patel, Shivam Sahni, and Nipun Batra.  Accurate and scalable gaussian processes for Fine-Grained air quality inference.  AAAI, 36(11):12080–12088, June 2022.
  * [3] Mark N Gibbs.  Bayesian Gaussian processes for regression and classification.  PhD thesis, Citeseer, 1998.
  * [4] Ville Tolvanen, Pasi Jylänki, and Aki Vehtari.  Expectation propagation for nonstationary heteroscedastic gaussian process regression.  In 2014 IEEE International Workshop on Machine Learning for Signal Processing (MLSP), pages 1–6, September 2014.
  * [5] Markus Heinonen, Henrik Mannerström, Juho Rousu, Samuel Kaski, and Harri Lähdesmäki.  Non-stationary gaussian process regression with Hamiltonian Monte Carlo.  In Artificial Intelligence and Statistics, pages 732–740. PMLR, 2016.
  * [6] A. Krause, A. Singh, and C. Guestrin.  Near-optimal sensor placements in Gaussian processes: Theory, efficient algorithms and empirical studies.  JMLR, 9:235–284, 2008.
  * [7] Burr Settles.  Active learning literature survey.  2009\.
  * [8] Eyke Hüllermeier and Willem Waegeman.  Aleatoric and epistemic uncertainty in machine learning: an introduction to concepts and methods.  Mach. Learn., 110(3):457–506, March 2021.
  * [9] Ian Osband.  Risk versus uncertainty in deep learning: Bayes, bootstrap and the dangers of dropout.  In NIPS workshop on Bayesian deep learning, 2016.
  * [10] Matias Valdenegro-Toro and Daniel Saromo.  A deeper look into aleatoric and epistemic uncertainty disentanglement.  In CVPR Workshop on LatinX in CV, April 2022.
  * [11] Stefan Depeweg, Jose-Miguel Hernandez-Lobato, Finale Doshi-Velez, and Steffen Udluft.  Decomposition of uncertainty in Bayesian deep learning for efficient and risk-sensitive learning.  In Jennifer Dy and Andreas Krause, editors, ICML, volume 80 of Proceedings of Machine Learning Research, pages 1184–1193. PMLR, 2018\.
  * [12] Christian Plagemann, Kristian Kersting, and Wolfram Burgard.  Nonstationary gaussian process regression using point estimates of local smoothness.  In Joint European Conference on Machine Learning and Knowledge Discovery in Databases, pages 204–219. Springer, 2008.
  * [13] Christopher J Paciorek and Mark J Schervish.  Spatial modelling using a new class of nonstationary covariance functions.  Environmetrics, 17(5):483–506, 2006.
  * [14] Sami Remes, Markus Heinonen, and Samuel Kaski.  Non-Stationary spectral kernels.  In NIPS, May 2017.
  * [15] Tompkins, Oliveira, and Ramos.  Sparse spectrum warped input measures for nonstationary kernel learning.  In NIPS, 2020.
  * [16] Mark D Risser and Daniel Turek.  Bayesian inference for high-dimensional nonstationary gaussian processes.  J. Stat. Comput. Simul., 90(16):2902–2928, November 2020.
  * [17] Andrew Gordon Wilson, Zhiting Hu, Ruslan Salakhutdinov, and Eric P Xing.  Deep kernel learning.  In AIStats, pages 370–378, 2016.
  * [18] Kalvik Jakkala.  Deep gaussian processes: A survey.  2021\.

##  Appendix A Model fit on various datasets

(a)

(b)

(c)

Figure 3: Fitted $(\ell,\sigma,\omega)$-GP to a) motorcycle helmet data; b)
step data and c) Olympic 100-m race data. Top: predicted distribution. Bottom:
learned latent GPs for the 3 hyper-parameters: noise ($\omega$), scale
($\ell$), variance ($\sigma$). The green region is 95% epistemic variance
$\text{var}(f(\mathbf{x}))$ and orange region is overall 95% confidence.