Tangent Transformers for Composition,Privacy and Removal
Authors: Tian Yu Liu
Aditya Golatkar
Stefano Soatto
Venue: ICLR 2024 poster (2024)
Keywords: Composition
Forgetting
Tangent Learning
Primary area: transfer learning, meta learning, and lifelong learning

Abstract: We introduce Tangent Attention Fine-Tuning (TAFT), a method for fine-tuning linearized transformers obtained by computing a First-order Taylor Expansion around a pre-trained initialization. We show that the Jacobian-Vector Product resulting from linearization can be computed efficiently in a single forward pass, reducing training and inference cost to the same order of magnitude as its original non-linear counterpart, while using the same number of parameters. Furthermore, we show that, when applied to various downstream visual classification tasks, the resulting Tangent Transformer fine-tuned with TAFT can perform comparably with fine-tuning the original non-linear network. Since Tangent Transformers are linear with respect to the new set of weights, and the resulting fine-tuning loss is convex, we show that TAFT enjoys several advantages compared to non-linear fine-tuning when it comes to model composition, parallel training, machine unlearning, and differential privacy. Our code is available at: https://github.com/tianyu139/tangent-model-composition

TANGENT TRANSFORMERS FOR COMPOSITION, PRIVACY AND REMOVAL
Tian Yu Liu"
Aditya Golatkar'
University of California, Los Angeles tianyuêcs wttp:/eccscllll.. edu
University of California, Los Angeles adi tyago atkaruucc edu
Stefano Soatto
University of California, Los Angeles soattoles. www.oeessslll....oeeo.ooeeeeaaee
uela. edu
ABSTRACT
We introduce Tangent Attention Fine-Tuning (TAFT), a method for fine-tuning lincarized transformers obtained by computing a First-order Taylor Expansion around a pre-trained initialization. We show that the Jacobian- resulting from linearization can computed efficiently in single forward pass, reducing training and inference cost to the same order of magnitude as its original non-linear counterpart, while using the same number of parameters. Furthermore, we show that, when applied to various downstream visual classification tasks, the resulting Tangent Transformer fine-tuned with TAFT can perform comparably with fine-tuning the original non-linear network. Since Tangent Transformers are linear with respect to the new set of weights, and the resulting fine-tuming loss is convex, we show that TAFT enjoys several advantages compared to non- non-linear fine- tuning when it comes to model composition, parallel training, machine unlearning, and differential privacy. Our code available at: hhtpp:: gitubb. com// Eianyu139/tange mtopt/aaanntttmmmmmllmmmoo composition
Product
INTRODUCTION
Deep Networks are highly non-linear operators trained by optimizing highly non-convex functions, yet some of the training dynamics convergence approximate those of linear weer-parameterized systems (Saxe et al. (003)). Accordingly, linearization has been used as a tool for both the analysis of deep networks (Jacot et al. 2018), and their design (Achille al. 202I). Linearization around an initial set of weights, however, is of limited practical relevance since the early learning dynamics are highly non-linear and decisive of final performance (Golatkar et al. 2019). On the other hand, linearization around pre-trained point has been shown to be essentially equivalent to non-linear fine- tuning, and better the case of few-shot fine-tuning (Achille et al., 2021 A linearized model has the same number of parameters as the original, but carries some distinct advantages: First, linearity allows straightforward model composition, whereby ensemble models can be formed by scalar combinations at essentially zero cost. Second, monolithic training set can be partitioned into smaller "shards," for instance for privacy or attribution purposes, and the resulting models combined to yield performance similar to model trained the monolith. This results in zero loss compartmentalizaatoon of separately trained models, and enables seamless parallel training. Third, since the linearized model can be trained by optimizing convex loss, existing methods for private training via selective forgetting (Abadi et al. 2016) are effective and enjoy strong theoretical guarantees.
Despite the benefits, model linearization is challenging scale. To this date, only small-scale models have been successfully shown to operate comparably to their non-linear counterparts, typically in the ResNet family of architectures. To our knowledge, our work is the first propose an efficient method to linearize models in the Transformer family of architectures, leading to what we call "Tangent Transformers." Tangent Transformers can be used to adapt Transformer models, as an alternative to prompt-tuning, fine-tuning, or adapter training, none of which are linear in weight space.
Denotes equal contribution

---

The key to enable practical linearization of Transformers is efficient way compute the Jacobian- Vector product in a single forward pass, described in Sec. 3.2 As a result, training and inference costs are on the same order of magnitude as the corresponding non-linear Transformer model. In Sec. 4.2 we show that a Tangent Vision Transformer (T-ViT) can achieve similar accuracy to non-linear fine- tuning (NLFT) of the original VIT (Dosovitskiy al. 2020) model. Given the comparable accuracy, we focus on illustrating some of the benefits of Tangent Transformers in Sec. 4 Specifically: Compositionality: Linearity yields equivalence between composition in weight space and composi- tion in activations, i.e., ensembling. This allows seamlessly combining independently trained models, with obvious benefits to parallel, incremental, and federated learning, while maintaining a constant inference time compared to traditional ensembling.
Speedup: Specifically, we achieve up to 10x (50x) speed-up in parallel training, with only 3. (9.3%%) drop in overall accuracy compared to non-linear fine-tuning on the full dataset, improv- ing over the Model Soup (Wortsman et al., 2022b) approach by 9.1% (13.5%) respectively. Compartmentalizationn
Since training on disjoint shards yields the same performance, data removal, it becomes necessary or desirable (Achille et al. 2023). can be performed deterministically in an exact fashion, at essentially zero cost.
Privacy: Most theoretical results and practical methods concerning Differential Privacy (DP) (Abadi
al. 2016 Bassily al. 2014] Fang et al. 2023 Yang al. 2022 Bassily et al. 2021 Wang et al. 2019: 2022a) provide much better utility-privacyy trade-offs when the optimization problem being solved is convex. While in general deep networks are not, pre-training conducted on safe data, linearized fine-tuning is convex and therefore strong results and effective methods for DP apply.
In Sec. we briefly survey relevant related work, and in Sec. [ we derive our method for Tangent Attention Fine-Tuning (TAFT). We illustrate the benefits of TAFT in Sect. 3.2 for parallel training and composition, in Sec. 3.3 for selective forgetting, or "unlearning". and in Sec. 3.4 for privacy. Finally, we empirically evaluate TAFT in Sec.
2 RELATED WORK
Deep network linearization:
Deep networks linearized using the first-order taylor approximation
have various interpretations in literature- - viewing gradients as features (Mu et al. 2020), learning along the tangent space of neural network (Liu et al. 2022 Liu & Soatto 2023) infinite-widh width networks (Jacot et al. 2018). Mu et al. (2020) shows that the Jacobian- - product (JVP) of linearized convolutional networks can be efficient computed in a single modified forward pass. Achille et al.] (2021) shows that by using Leaky-Rel_U JU activations and training with the rescaled square loss (Hui & Belkin 2020) and gradient preccondiii.ning. ResNets linearized around ImageNet pre-trained weights can achieve comparable performances to the original non-linear networks on downstream fine-tuning tasks. Most similar to our work, Liu & Soatto 2023) applies linearized convolutional networks to ensembling and continual fine-tuning. To the best of our knowledge, we are the first work to linearize transformer networks in a manner that is both computationally efficient and achieves competitive results when fine-tuned on various downstream tasks.
Composition: We investigate compositionality of deep networks in weight space to yield a model that generalizes better than each individual component model. Weight averaging has been used to improve generalization of pre-trained weights (Choshen et al./ (022). and for distributed fine- tuning (Liu & Soatto [023 Wortsman et al., (202a)). Wortsman et al. (2022b) averages the weights of large pre-trained models fine-tuned with different hyperparameters to improve generalization. Compositionality has also been explored through prompts in continual learning (Wang al. 2022b c)). However, these works do not develop any theoretically meaningful interpretations for composition, and often scale in inference time as number of component models increase. We introduce TAFT for linearly composing tangent transformer models trained, possibly in parallel, on multiple disjoint shards of data. Under our method, composition of linear weights theoretically equivalent to output ensembling with constant inference cost.
Machine Unlearning:
Machine unlearning, or forgetting, methods aim to remove the influence of
specific samples from trained network (Achille et al. 2023 Bourtoule et al. 202 Dukler et al.) 2023, Golatkar et al.) 2021 2020a b). We focus on methods that are zero-shot and yields theoretically guaranteed unlearning. These works (Bourtoule et al. 2021 Bowman et al. 2023 Koch & Soll
operate by splitting datasets into multiple disjoint shards, hence compartmentalizing each sample to only single shard. Unlearning can then be done by simply removing the shard. However, such
2

---

methods either incur high inference costs as result of running inference across multiple models, and often incur significant trade-offs in generalization accuracy of the composed model. Instead, we show that result of linearity, can compose tangent transformer networks simply by averaging network weights, to produce outputs that are equivalent to the ensemble of each network with an inference cost that is constant with respect to number of shards/models in the ensemble.
Privacy: Differential privacy (Dwork et al. 2014) seeks to limit the amount of information a trained model contains about the individual training samples. DP-SGD (Abadi et al. 2016) achieves this through clipping of individual gradients followed by addition of Gaussian noise. Bassily et al. (202 I
2014): Fang et al., (2023): Waang al. (2019 2022a); Yang et al. (2022) provide rigorous theoretical guarantees for the convergence and utility of DP algorithms, and show that convex (or strongly convex) models offer better utility compared to their non-convex counterparts. Per dimension Gaussian noise in DP-SGD reduces the utility of training large models in favour of fine-tuning parameter efficient adapters (Bu al., 2022b aj Golatkar et al., 2022] Yu et al. 202T). In this paper, we show that TAFT in the tangent space of these parameters provides better utility-privas trade-off.
3 METHOD
We explore the most direct way to linearize pre-trained transformer network fw by replacing with its first-order taylor approximation plin about its pre-trained weights w.
plim f- + V vw Aw
By construction, flin is now linear with respect to Aw, the new set of learnable parameters.
The new network can be trained easily using any loss function. For example, using the standard mean-squared error loss yields a quadratic objective function, reducing the training of such models to simple lincar-quadratic optimization (Achille et al. 2021). We use the Rescaled Square Loss (RSL) (Hui & Belkin 2020) given by
L(x, v) (a(lf("()))s x)² +
Σ 11,,
(L)T())))))
where a, K are hyper-parameters. We empirically found RSL performs better compared to cross-
entropy regular MSE loss, corroborating the results of Achille et al. 20021); Liu & Soatto (203).
We further note that how good local approximation of the network is depends on the distance that the fine-tuned weight moves from its initial point w. As such, we additionally regularize the training objective by adding a penalty on || The resulting training objective is simply ridge regression, retaining the benefits of linear-quadratic optimization while obtaining better empirical results (Sec. (4.6). Note that due to the high dimensionality of the training set and gradient- features, it is computationally prohibitive obtain the closed form solution even though it exists.
This appears costly to compute for both inference and training, since evaluating the Jacobian- Vector Product (JVP) V fw Aw requires computing the gradient with respect the original weights, for every input x. However, by computing the directional derivative, we can derive closed form equations for the linearized versions of the key building blocks of transformer networks. We show that they can be computed in a single modified forward pass through the original model where each layer of the network outputs the computed JVP (JVP out) in addition the original output values, and takes as input the JVP from the previous layer (JVP... ....) in addition to the original input values.
3.1 LINEARIZING TRANSFORMERS
Here, we will derive the closed form linearization a transformer network, and show that can be easily computed by the modified forward propagation without explicitly computing any gradients. We break down transformer networks into attention, normalization, and fully fully-connected layers, and separately derive their linearizations (note that while fully-comnected layers are already linear, still need to handle the input JVP from the previous layer). These layers can be simply composed together to form the final Tangent Transformer network.
3

---

We parameterize the attention function A R R by the weights Wg. Wa, W, E Rdxd corresponding to the key, query and value matrices respectively, and given by
A(x) = (x)) where ((x) = a (x)))) Q(x) - (W,, K(x) = (W x). ((x) = (W..
(4) where is the soft-max activation function. We will write Q, K.K.V. V. instead of Q(x), ((), Vzx), ((x) for case of notation. For simplicity, we only consider single-headed atten- tion in our derivations, but note that our definitions and derivations can be extended to multi-headed attention (which we use in the experiments section) with minimal modification. Now, we wish to compute the first-order approximation of A, denoted Atin Rd + Rd and parameterized by the linearized weights AW, AWa, ,, for the key, query, and value matrices respectively. By taking directional derivatives, can be found in Appendix B):
derive the following closed form expression for Aun (details
Atin (x) = A(x) + lim A(x, Wg + rAW WE + rAWA, W. + rrWW,
(5)
T0 Dr
JVP
A(x) + ( v (1 3 (D"))) V+ Ør
(6)
where
V- = V(x) = (AQ + WT JVP K) + (Q. + WT JVP
(7)
r : [(x) = AV + W JVP in
AQ == (AW,, 2), AK := (AW,, x). AV == (AW,, x)
(9)
@ denote the Hadamard product, and JVP lim,- 8r from the previous layer, obtained from the modified forward pass. The terms Tb Q, K, V can be obtained for free as intermediate variables from computing A(x). Thus, computing the JVP term done through simple matrix multiplication operations of similar computational complexity as the original attention mechanism.
Or is the Recobian-Vector Product computed
Transformer blocks also include several normalization layers. Similarly
compute closed form expression for their linearized versions that can be obtained in the modified forward propagation step. We show the derivation for Layer Norm, which we denote (NN (-) and parameterize by the affine transformation parameters (7, (),, but note that the results can be easily generalized other forms such as Batch Norm (Achille et al., 2021). In particular, the linearized Layer Norm LN which is parameterized by (A7, A)) evaluated at can be computed
LN (x) = LN(G (x) + LN(AA)A99)
(10)
1
(JVP.n -[JVP:..]) B[(x (JV) -EJVP. (x 1B/x])
Var(z]
Var(s)
JVP our
(11)
where * is the element-wise scaling operation.
Fully-connected (FC) layers parameterized by weight W and bias b can be easily modified to handle JVP in from the previous layer and has already been derived and used in prior works (Achille et al. 2021 Mu et al. 2020). We include it below for completeness.
FC (x) FC(x) + AWT x + Ab + WT JVP..
(12)
Mon-lincasistes also conveniently handled by the same technique. We illustrate the derivation for the GeI.U activation commonly used in transformer-baat networks:
GeLU (a) GeLU(x) + (GelU((x) (x: + ..e.CDF())) JVP,,.
(13)
where CDF(() evaluates the Standard Normal CDF at æ. As before, all terms can be easily computed without any backpropagation steps.
4

---

Table 1: Comparison of non-linear fine-tuning (NLFT) TAFT on various downstream datasets sorted by distance to ImageNet (Achille et al. 2021). We compare fine-tuning the last attention block (NLFT-I, TAFT-1), the last 7 attention blocks (NLFT-7, TAFT-7), and only the classification head (FC). On most datasets sufficiently close to the ImageNet pre-training task, we show that TAFT can yield comparable or better performance compared to NLFT and FC, while benefiting from linearity.
Dataset
NLFT-7
TAFT-7
NLFT-1
TAFT-1
FC
Caltech-256
93.7
95.7
95.8
95.9
95.7
MTT-67
86.7
88.4
88.1
89.3
87.8
Oxford Pets
93.0
94.7
94.2
94.5
94.0
Stanford Dogs
84.8
91.6
91.2
91.9
90.7
CUB-200
87.3
89.3
88.8
89.0
87.7
PGVC-Aircrarie
77.7
60.5
69.9
62.2
58.2
Stanford Cars
75.6
71.3
67.3
67.2
57.4
Average
85.5
84.5
85.0
84.3
81.6
The final linearized transformer, which we term Tangent Transformer, is simply the composition of such layers, chained together using the modified forward pass. Since Tangent Transformers are linear only in the weights Aw, and highly non-linear in the original weights we only update Aw during fine-tuning, a process we term Tangent Attention Fine- Tine-Tuning (TAFT).
3.2 PARALLEL TRAINING AND COMPOSITION
Given N models linearized about pre-trained weights w and a query æ, the ensemble of these models, defined by the linear combination their outputs, equivalent to evaluating a single tangent model composed by taking the same affine combination of component models in weight space:
Σ plimi (x) = (z) V Jw
A;
(14)
This gives rise to a natural interpretation of weight space composition via output ensembling, while reducing the cost of ensembling N models from O(N) to O(1). In other words, we can train multiple Tangent Transformers on multiple different datasets in completely parallel manner, and simply combine their output weights to yield a single model that performs well as their ensemble but constant inference time. Such compositions
weight space of transformer networks have also been previously explored by Wortsman et al. (202bb) combining multiple models trained the same dataset with different configurations using weight averaging. However, as we show in Sec. [4.] the lack of any theoretical relationship between combinations of models in weight space and their resulting output causes the resulting model to perform poorly when component models are trained on disjoint sets of data. On the other hand, we will show that the equivalence of weight averaging and ensembling allow the composition of up to 50 T-VITs trained on different shards of data with relatively much smaller accuracy trade-offs compared naively composing non-linear models.
3.3 ZERO-/LOW-Con COST FORGETTING WITH TANGENT TRANSFORMERS
"Learning" a model by combining the weights of component tangent models, each trained on disjoint shards of data, also allows for the subtraction of each component from the final model. Clearly, this subtraction operation completely removes the influence of samples contained within the shard used to train the component model from the final model. This is highly advantageous for machine unlearning.
Given request to forget training sample, the paragon unlearning method that guarantees forgetting of the target sample requires training the entire model from scratch on the remaining dataset samples. models like GPT-3
This is clearly impractical especially for large real-world transformer-basee (Brown et al. 2020). With a Tangent Transformer composed from individual component models, can simply remove the shard containing the sample to be forgotten by subtracting the weights of the associated component model. This theoretically guarantees forgetting while preserving accuracy when number of forgetting requests is small (Fig. [(a)). all at essentially zero computational cost.
We note that this method of unlearning through shard removal is not scalable, since performance of the composed model degrades number of forgetring requests increases. Instead, can also
5

---

Table 2: We compose multiple T-ViTs, each trained with TAFT-1 disjoint shard of dataset. The equivalence between linearly combining weights and output ensembling enables the composed T-V to outperform Model Soup (Wortsman al. 2022b) across all datasets and sharding factors.
10 Shards
25 Shards
50 Shards
Dataset Shards
Soup
TAFT
Soup
TAFT
Soup
TAFT
Caltech-256
94.5
95.0
93.2
94.3
90.8
93.3
MIT-67
86.4
86.9
84.3
85.7
80.3
84.2
Oxford Pets
93.3
93.6
92.0
92.2
89.5
91.3
Stanford Dogs
90.9
91.4
90.4
90.9
88.8
90.4
CUB-200
82.0
86.5
69.7
63.0
81.4
FGVC-Aircrafts
38.2
57.6
18.6
53.0
14.7
17.7
Stanford Cars
19.8
58.4
10.6
49.2
8.4
41.5
Average
72.2
81.3
65.5
78.5
62.2
75.7
optionally retrain the component model on the remaining samples in the shard, after removing the sample be unlearned. Since shards are much smaller than the full dataset, this enables orders of magnitude speedup compared to the paragon of re-training from scratch, yet guarantees forgetting of the requested samples and maintains generalization performance of the resulting model (Fig. [(c)..
3.4 TAFT WITH DIFFERENTIAL PRIVACY
Differential privacy (Dwork et al. 2014) is mathematical framework to design algorithms which protect the privacy of individual training samples. Given a training dataset D, and an algorithm M, say that M is (c, -j-differentiaaly private (DP) only if
P(M(D) E E) S E)
for all E, D-.. where D_, is obtained by removing the 1th sample from D. In simple words, DP enforces an algorithm to produce similar outputs when the dataset differs by single sample. One of the most popular ways of enforcing DP in deep learning is to use DP-SGD (Abadi et al.) 2016 during training. SPD introduces two modifications over the standard stochastic gradient descent (Robbins & Monro, 1951). first it clips the gradient norm of every sample, and then it adds gaussian noise the sum of the clipped gradients across training batch. Thus the information pertaining to individual samples is contained with clipping with noise perturbation. It is well known (Bassily ctt al. 2021 2014 Fang et al. 2023 Yang al. 2022) that convex models have better convergence and utility guarantees with trained di differentially private convex optimization algorithms (in our case DP-SGD). We show in Tab. M that TAFT on Tangent Transformers provide comparable results to (in some cases better than) non-lines fine-tuning. As such, experiments in Sec. |.5 seek to understand if such models can remain effective in DP settings to reap the benefits of theoretical guarantees provided by private convex optimization.
3.5 CHOOSING A GOOD INITIALIZATION POINT
Strong pre-training objectives provide natural initialization point at which we can compute the tangent model. However, linearizing transformer models around the full pre-training weights might exhibit strong feature biases towards the source pre-training dataset that might not transfer well to downstream tasks, especially in the later layers. As such, we propose a simple method to overcome this, by linearizing about randomized re-initialization for the later attention layers, while keeping the pre-training weights constant for earlier layers in the network. We show that this significantly improves results in Fig. 2(b For Vision Transformer-basee classifiers, we further show in Fig. 2(c) that the CLS token itself can also be linearized in the same manner. We will empirically show that this can be beneficial for certain downstream tasks which are "far" from the pre-training initialization.
EXPERIMENTS
In Sec. show that TAFT on Tangent Transformers attain similar performances downstream tasks compared to non non-linear fine-tuning. We show the advantages that arise from linearity for composition and parallel training in Sec. 4.3 machine unlearning in Sec. 4.4 and privacy

---

100
180
SEA
14/7
80
80
Accuraty 1%
Accurscy
tedh-225
60
Accuracy
Caltach255e
VE-67
40
Deford Pets
40
Stantord Degs
20
Stantort bryg-
Stantors Care
TCVC Arrats
20
20
Stanford Dars
FBVC Airraats
CUB 20n
10
20
30
40
50
1.0
20
30
40
50
Number of Shards Removed
Number of Shards Removed
Percentage of Dataset Forget %
(a) Free forgetting via subtraction
(b) Comparison to SISA
(c) Shard retraining
Figure I: (a) We show that when number of samples to forget is small, we can simply remove shards by subtracting the weights of their respective component model with minimal drop in final model accuracy (computed as an expectation over uniform distribution of sample forgetting requests). (b)
We compare against SISA (Bourtoule et al. (02I) which also uses a sharding technique for zero-cost unlearning. Our method is uniformly better across all number of shards removed on all datasets. (c) Retraining on the remaining samples a shard after forgetting request can further improve accuracy of the "unleamed"" model, while enjoying up 50x faster training time compared to full re-training.
in Sec. 4.5 We describe our implementation details in Sec. and carry out ablation studies on our implementation choices in Sec. 4.6, In Appendix C.4 and C.1 we also present ablations on different pre-training schemes, and comparisons against Linearized ResNets.
4.1 IMPLEMENTATION DETAILS
We run all our experiments Vision Transformers on image classification tasks. In particular, we use ViT-L/16 (Dosovitskiy et al. 2020) as the base model all our experiments, and linearize around its ImageNet pre-trained weights, the result of which call T-VIT-L/16. We evaluate on the following datasets in increasing order of distance from the ImageNet pretraining task based on Li et al. (2020) Caltech-256 (Griffin et al., 2007), MIT-67 (Quattoni & Torralba, 2009), Oxford Pets (Parkhi et al. 2012), Stanford Dogs (Khosla et al. 2011), CUB-200 200 (Wah et al.] 2011). FGVC-Aircrafts (Maji et al.
2013 and Stanford Cars (Krause al. 2013). Further details be found in the Appendix.
4.2 How WELL DOES THE TANGENT TRANSFORMER COMPARE TO THE ORIGINAL MODEL?
While linearity yields many benefits in terms of composition, privacy, forgetting, and even inter- pretability, there one main drawback Tangent Transformers are strictly less expressive compared to the original non-linear model. Hence, for such linear models to be practical, we wish to preserve as much performance as possible on downstream tasks. We show Tab. [ that in fact, due to the strong inductive priors from the ImageNet pre-trained initialization, the average downstream performance highly comparable with that of non-linear fine-tuning of the original model, differing on average only by 1.0% and 0.7% respectively when fine-tuning multiple attention blocks and just the last attention block. In fact, for several tasks that are close to the pre-training dataset (ImageNet) such as MIT-67, Stanford Dogs, Oxford Pets. and CUB-200 we show that TAFT actually outperforms non-linear fine-tuning. We hypothesize that this results from the implicit regularization imposed by the linearity constraints. We further note that for tasks that far from the pre-training dataset, such as Stanford Cars and FGVC-Aircrafts, the local approximation becomes less accurate. As expected, the divergence between non-linear fine-tuning and TAFT increases. However, compared transfer learning that simply fine-tunes the classification head, TAFT is strictly more expressive and hence improves by up 2.9% on average while maintaining linearity in weights.
Since most of the accuracy gains can be made from just fine-tuning the last attention block (NLFT-1, TAFT-1), this also allows for parameter-efficien
tine-tuning where the number of parameters are < 5% of that needed for full fine-tuning. As such, we adopt NLFT-I/TAFT-I in the following sections for non-linear/linear fine-tuning respectively, where we explore several benefits that linearity yields.
7

---

Table 3: DP fine-tuning of Tangent Transformers compared to regular non-linear fine-tuning. "Full" fine-tunes the entire (last) attention block, "BitFit" (Zaken et al. 2021) fine-tunes the bias parameters of the attention block, "Layer Norm" fine-tunes the affine parameters of layer normalization modules, and "FC" fine-tunes the classification head. "Ours" refer to fine-tuning the linearized parameters, and "NLFT" refers to fine-tuning the original parameters. Training tangent models outperforms its non-linear counterparts in all training regimes (italics: best for each regime, and bold: best overall)
Full
BitFit
Layer Norm
Dataset
Privacy
FC
Ours
NLFT
Ours
NIFF
Ours
NLFT
1
46.3
11.0
44.5
41.9
46.7
45.9
44.9
CUB-200
3
72.2
57.4
72.6
71.9
72.7
71.3
72.1
6 = 8
82.2
77.2
82.0
81.7
82.4
82.1
82.1
1
4.3
1.4
4.9
4.8
4.6
4.8
5
Stanford Cars
6 = 3
13.6
5.7
15.0
13.7
14.1
13.6
13.9
6 8
26.8
16.6
27.5
26.3
27.2
27.0
26.2
4.3 COMPOSITIONALITY AND PARALLEL TRAINING
We evaluate our proposed method for parallel training and composition described in Sec. 3.2 We first shard a dataset into N disjoint subsets, and train individual models on each subset. Note that training can be done in parallel, yielding N x speed-up in training time. In Tab. ab. we show that across various sharding factors (N 10, 25, 50) of each dataset, linearly combining weights of models fine-tuned with TAFT significantly outperforms composing separately trained non-linear models via Model Soup (Wortsman et al.) 2022b), which to the best our knowledge, is the only method that yields composed model with O(1) inference cost (with respect to number of component models). Indeed, naively composing non-linear models through weight averaging yields no theoretical guarantees regarding how the output space of the resulting model changes. However, composing the linear weights of Tangent Transformers trained via TAFT is theoretically equivalent to output ensembling, hence outperforms Model Soup by 9.1 %, 13.0%, and 13.5% on 10, 25, and 50 shards respectively, while maintaining an (1) inference cost.
4.4 MACHINE UNLEARNING
Tangent Transformers composed from component tangent models trained on disjoint shards of data enables forgetting "for free", since unlearning can be done by simply subtracting models without needing any further training. We show in Fig. I(a] that for a model composed from 50 shards, one can drop up to half of the shards (25) with only 4.0% drop in accuracy. We also compare against SISA (Bourtoule et al. 202T), which also drops shards upon forgetting requests, and show that we perform uniformly better across all datasets and number of shards dropped, and on average by 11.0%.
Optionally, one can retrain the shard containing the sample to be forgotten on the remaining samples in the shard. Even then, this still yields significant advantages compared to the baseline of re-training a model from seratch, since only the relevant shard needs to be retrained. In Fig. [(c) we show that this yields 50x speed-up in our experiments, achieving close to the paragon performance (Appendix C.3 with only a 6. 6.2% drop in accuracy after unlearning 50% of the entire dataset.
4.5 PRIVACY
We hypothesize that combining TAFT with differential privacy results better utility privacy trade-off resulting from convexity of the loss landscape. To illustrate this, we fine-tune various parameters of 1- on two different fine- fine-grained datasets (CUB200-easy and Stanford Cars-hard) for different privacy range. In Tab. 3 we observe that under almost all settings, privately fine-tuning the linearized parameters performs much better than privately fine-tuning the non-linear parameters.
While fine-tuning the entire last attention block (column "Full" in Tab. 3) we observe that the gradient noise significantly degrades the model utility compared only fine-tume tee-tunnng the last cally-connected layer (and biases/normalizatioon layers) of the network. The linear nature of tangent transformers along learning algorithm i.e.
with the results in Tab, B also inspires a simple private composition/continn train private models shards of data, and linearly combine their weights.
8

---

40
Reset Random
Unearfeed CLS Taken
60
35
50
Enor
15
MSE
CE
RSL
caltech
cattech
Loss
Dataset
Dataset
(a) Ablation of RSL.
(b) Choice of initialization
(c) Linearizing CLS token
Figure 2: (a) RSL can improve fine-tuning performance, beating CE and MSE by 1.5% and 9.0% respectively across 7 datasets. (b) While computing the tangent model about the full pre-training initialization is already effective its own, re-initializing the weights of the last attention block before linearization can yield further performance gains. (c) Linearizing the CLS token improves accuracy downstream datasets which are far from the pre-training tasks.
4.6 ABLATION STUDIES
We also conduct ablation studies to show key implementation details needed to train tangent models to perform comparable to non-linear models in downstream tasks. In particular, Fig. 2(a) shows that using the rescaled square loss significantly improves average-case results across all datasets by average of 9.0% and 1.5%, and on the hardest dataset by 23. 23.8% and 2.7% compared to the MSE and CE loss respectively. In Fig. 2(b) we show that by resetting the weights of the final attention layer prior to linearization, average performance across datasets improves by 1.5%. We hypothesize that this is due the negative transfer (Zhang al.. 2022) of later features learnt from the pre-training task to new downstream tasks. Indeed, we note that for datasets which are very close to ImageNet (i.e. Caltech-256, MIT-67), linearizing about the original pre-trained weights perform marginally better since they are highly transferrable to these downstream tasks. Similarly, we show in Fig. 2(c) that resetting and linearizing the CLS token in the last attention block of vision transformer network can significantly improve performance on datasets which are far from the ImageNet pretraining task, improving results on FGVC-Airerafts and Stanford Cars by 0.8% and 3.1 % respectively.
DISCUSSION
Tangent Transformers linearized about a strong pre-trained point can serve to facilitate number of processes related to fine-tuning and ensembling. Independently trained linear components can be easily composed, thus realizing full parallelism, and disgorged if need be, thus realizing deterministic removal of data.
However, linearization not panacea: For the linearized models to work as advertised, the point around which the model is linearized is important, which can only be ascertained empirically. Once that is done, linear components can be trained with convex losses, which leads to overall models that enjoy strong guarantees for convergence, privacy, and composition. This limitation be further mitigated via techniques such as resetting certain pre-trained weights and linearization the CLS token, shown in our experiments. Another limitation of our method that the inference cost of Tangent Transformer can potentially be double that of the original model, since the modified forward pass requires an additional set of computations in addition the original forward pass. However, we show that linearizing the last attention block of a ViT-L/16 model is often sufficient to yield strong performances on several downstream tasks. Under this regime, training and inference parameter efficient, and linearization only incurs slight increase in inference cost. Note that during training where the dataset is fixed, inference costs can actually be reduced to that of the original non-linear model by simply caching the activations from the static pre-trained weights for each training example and for each layer. Lastly, as observed by Koch & Soll, sharding naturally incurs significant trade-off's in performance minority classes when training highly imbalanced datasets.
The tasks on which we demonstrated how the linearity of transformers can be exploited through TAFT certainly not exhaustive. Yet the encouraging empirical results of TAFT make it candidate replacement for any applications of transfer learning or fine-tuning, while benefiting from the simplicity, composability, and interpretability of linear models.

---

REFERENCES
Martin Abadi, Andy Chu, lan Goodfellow, Brendan McMahan, llya Mironov, Kunal Talwar, and Li Zhang. Deep learning with differential privacy. In Proceedings of the 2016 ACM SIGSAC conference on computer and communications security, pp. 308-318, 2016.
Alessandro Achille, Aditya Golatkar, Avinash Ravichandran. Marzia Polito, and Stefano Soatto. Lqf: Linear quadratic finettuing. In Proceedings the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pp. 15729-15739, 2021.
Alessandro Achille, Michael Kearns, Carson Klingenberg, and Stefano Soatto. Ai model disgorge- ment: Methods and choices. arXiv preprint arXiv:2304.03545, 2023.
Racf Bassily, Adam Smith, and Abhradeep Thakurta. Private empirical risk minimization: Efficient algorithms and tight error bounds. In 2014 IEEE 55th annual symposium on foundations of computer science, pp. 464 473. IEEE, 2014.
Raef Bassily, Cristóbal Guzmán, and Michael Menart. Differentially private stochastic optimization: New results in convex and non-convex settings. Advances in Neural Information Processing Systems, 34:9317-9322 2021.
Lucas Bourtoule, Varun Chandrasekaran, Christopher A Choquette-Choo Hengrui Jia, Adelin Travers, Baiwu Zhang, David Lie, and Nicolas Papemot. Machine unlearning. In 2021 IEEE Symposium on Security and Privacy (SP), pp. 141-159. IEEE, 2021.
Benjamin Bowman, Alessandro Achille, Luca Zancato, Matthew Trager, Pramuditha Perera, Gio- vanni Paolini, and Stefano Soatto. aa-la-carte prompt tuning (apt): Combining distinct data via composable prompting. arXiv preprint 200209930..07994 2023.
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. Advances in neural information processing systems, 33:1877-1901, 2020.
Zhiqi Bu, Yu-Xiang Wang, Sheng 7ha, and George Karypis. Differentially private optimization on 10. 10.000388 2022a.
large model at small cost. arXiv preprint
Zhiqi Bu, Yu-Xiang Wang, Sheng Zha, and George Karypis. Differentially private bias-term only fine-tuning of foundation models. arXiv preprint arXiv:2210.000366 2022b.
Leshem Choshen, Elad Venezian, Noam Slonim, and Yoav Katz. Fusing finetuned models for better pretraining. arXiv preprint arXiv:2204.0300 2022.
Mircea Cimpoi, Subhransu Maji, lasonas Kokkinos, Sammy Mohamed, and Andrea Vedaldi. Describ- ing textures in the wild. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 3606-3613, 2014.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition scale. arXiv preprint arXiv:2010. 119.29990..1999... 2020.
Yonatan Dukler, Benjamin Bowman, Alessandro Achille, Aditya Golatkar, Ashwin Swaminathan, and Stefano Soatio. Safe: Machine unlearning with shard graphs. arXiv preprint a3X1: 2023.
Cynthia Dwork, Aaron Roth, et al. The algorithmic foundations of differential privacy. Foundations and Trends@ in Theoretical Computer Science, 9(3-4):21 1-447, 2014.
Huang Fang, Xiaoyun Li, Chenglin Fan, and Ping Li. Improved convergence of differential private sgd with gradient clipping. In The Eleventh International Conference on Learning Representations, 2023.
Aditya Golatkar, Alessandro Achille, and Stefano Soatto. Eternal sunshine of the spotless net: Selective forgetting in deep networks. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 9304 9312, 2020a.
10

---

Aditya Golatkar, Alessandro Achille, and Stefano Soatto. Forgetting outside the box: Scrubbing deep networks of information accessible from input-output observations. In Computer Vision-ECCV 2020: 16th European Conference, Glasgow, UK, August 23-28. 2020, Proceedings, Part XXIX 16, pp- 383-398. Springer, 2020b.
Aditya Golatkar, Alessandro Achille, Avinash Ravichandran, Marzia Polito, and Stefano Soatto. Mixed-privacy forgetting in deep networks. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, pp. 792-802 2021.
Aditya Golatkar, Alessandro Achille, Yu-Xiang Wang, Aaron Roth, Michael Kearns, and Stefano Soatto. Mixed differential privacy in computer vision. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 8376-8386, 2022.
Aditya Sharad Golatkar, Alessandro Achille, and Stefano Soatto. Time matters in regularizing deep networks: Weight decay and data augmentation affect early learning dynamics, matter little convergence. Advances in Neural Information Processing Systems, 32. 2019.
Gregory Griffin, Alex Holub, and Pietro Perona. Caltech-256 object category dataset. 2007.
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 International Conference on Machine Learning, pp. 2790-2799. PMLR, 2019.
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. arXiv preprint 1rXiv:2106.09685, 2021.
Like Hui and Mikhail Belkin. Evaluation of neural architectures trained with square loss cross- entropy in classification tasks. arXiv preprint arXiv:2006.07322. 2020.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems, 31, 2018.
Aditya Khosla, Nityananda Jayadevaprakash, Bangpeng Yao, and Fei-Fei Li. Novel dataset for fine-grained image categorization: Stanford dogs. In Proc. CVPR workshop on fine-grained visual categorization (F (FGVC), volume 2. Citeseer, 2011.
Korbinian Koch and Mareus Soll. No matter how you slice it: Machine unlearning with sisa comes the expense of minority classes. In First IEEE Conference on Secure and Trustworthy Machine Learning.
Jonathan Krause, Michael Stark, Jia Deng, and Li Fei-Fei. 3d object representations for fine-grained categorization. In Proceedings of the IEEE international conference computer vision workshops, pp. 554 561,
Hao Li, Pratik Chaudhari, Hao Yang, Michael Lam, Avinash Ravichandran, Rahul Bhotika, and Stefano Soatto. Rethinking the hyperparameters for fine-tuning. arXiv preprint arXiv:2002.11770 2020.
Tian Yu Liu and Stefano Soatto. Tangent model composition for ensembling and continual fine-tuning. arXiv preprint 28Xiv:2307.08/ 14, 2023.
Tian Yu Liu, Aditya Golatkar, Stefano Soatto, and Alessandro Achille. Integral continual learning along the tangent vector field of tasks. arXiv preprint arXiv:2211.13108, 2022.
S. Maji, J. Kannala, E. Rahtu, M. Blaschko, and A. Vedaldi. Fine-grained visual classification of aircraft. Technical report, 2013.
Fangzhou Mu, Yingyu Liang, and Yin Li. Gradients as features for deep representation learning. arXiv preprint 0r55522004.05529, 2020.
Omkar M Parkhi, Andrea Vedaldi, Andrew Zisserman, and CV Jawahar. Cats and dogs. 2012 IEFE conference computer vision and pattern recognition, pp. 3498-3508 IEFE, 2012.
11

---

Ariadna Quattoni and Antonio Torralba. Recognizing indoor scenes. In 2009 IEEE conference computer vision and pattern recognition, pp. 423.420. IEEE, 2009.
Herbert Robbins and Sutton Monro. A stochastic approximation method. The annals of mathematical statistics, pp. 400-407, 1951.
Andrew M Saxe, James L. McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. arXiv preprint arXiv: /3/2.6/33/ 2013.
C. Wah, S. Branson, P. Welinder, P. Perona, and S. Belongic. Caltech uesd birds-200-2011. Technical Report CNS- 2--2011-001, California Institute of Technology, 2011.
Di Wang, Changyou Chen, and Jinhui Xu. Differentially private empirical risk minimization with non-convex loss functions. In International Conference on Machine Learning, pp. 6526-6535.
PMLR, 2019.
Puyu Wang. Yunwen Lei, Yiming Ying, and Hai Zhang. Differentially private sgd with non-smooth losses. Applied and Computational Harmonic Analysis, 56:306-336, 2022a.
Zifeng Wang, Zizhao Zhang, Sayna Ebrahimi, Ruoxi Sun, Han Zhang, Chen- Yu Lee, Xiaoqi Ren, Guolong Su, Vincent Perot, Jennifer Dy, et al. Dualprompt: Complementary prompting for rehearsal-free continual learning. In Computer Asioo-EECV 2022: 17th European Conference, Tel Aviv, Israel, October 23-27, 2022, Proceedings, Part XXVI. pp. 631-648. Springer, 2022b.
Zifeng Wang, Zizhao Zhang, Chen- Yu Lee, Han Zhang, Ruoxi Sun, Xiaoqi Ren, Guolong Su, Vincent Perot, Jennifer Dy, and Tomas Pfister. Learning to prompt for continual learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 139-149, 2022c
Mitchell Wortsman, Suchin Gururangan, Shen Li, Ali Farhadi, Ludwig Schmidt, Michael Rab- bat, and Ari S Morcos. lo-fi: distributed fine-tuning without communication. arXiv preprint arXiv:2210. 11948, 2022a.
Mitchell Wortsman, Gabriel Ilharco, Samir Ya Gadre, Rebecca Roelofs, Raphael Gontijo-Lopes, Ari S Morcos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith, al. Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time. In International Conference on Machine Learning, pp. 23965-23998. PMLR, 2022b.
Xiaodong Yang, Huishuai Zhang, Wei Chen, and Tie- Yan Liu. Normalized/clipped sgd with per- turbation for differentially private non-convex optimization. arXiv preprint S.Ki::2206.10533 2022.
Da Yu, Saurabh Naik, Arturs Backurs, Sivakanth Gopi, Huseyin A Inan, Gautam Kamath, Janardhan Kulkarni, Yin Tat Lee, Andre Manoel, Lukas Wutschitz, et al. Differentially private fine-tuning of language models. arXiv preprint arXiv:2110.06500 2021.
Elad Ben Zaken, Shauli Ravfogel, and Yoav Goldberg. Bitfit: Simple parameter-efficient fine-tuning for transformer-based masked language-models. arXiv preprint arXiv:2 10/99, 2021.
Wen Zhang, Lingfei Deng, Lei Zhang, and Dongrui Wu. A survey on negative transfer. JEEF/CAA Journal of Automatica Sinica, 2022.
12

---

Supplementary Material
A IMPLEMENTATION DETAILS
We rum all our experiments with ViT-L/16, and its tangent model termed T-ViT-L/16. Unless indicated otherwise, we adopt varameter-effiiien vision transformer, along with the last normalization and fully connected layer. For experiments using TAFT in Table and 2, and Figures ((a)-(c), we train with the RSL loss using 6 = 15. We also adopt a 30 epoch learning schedule for each dataset/task, with learning rate decay by factor of 10 at epochs 15 and 25. We use a batch size of 32 for all our experiments, and train using Adam optimizer. We search over learning rates (LR) of {0.001, 0.0001} for both non-linear fine-tuning and TAFT. We also select the best base model to linearize by performing TAFT on both the default pre-trained model, and resetting the last attention block and classification layer before linearization (PT=T and PT=F respectively), and ablating over resetting and linearizing the CLS token (CLS-T/F), We list the best configurations for non-linear fine-tuning (NLFT) and TAFT in Tab. 4 Tab. 5 (Table 1, main paper) and Tab. 6 (Table 2, main paper). We do not use any data augmentation for our experiments.
tine-tuning by training only the last attention block of the
For experiments on composition and machine unlearning, datasets are split into muliple shards with respect to a fixed random seed by uniform sampling without replacement. For experiments on differential privacy, we use the standard cross entropy loss. For each 1, 3, 8, we use a 50 epoch training schedule with no learning rate decay and search over PT E {T, F}. We predict only using the JVP output of the network. Gradients at each epoch are also aggregated over the entire dataset.
Table 4: Best Hyperparameters Full Data, Single Attention Block
Dataset
FC
NLFT-1
LFT-1
Caltech-256
LR=le-4
LR=le-4,
CLS=F
(3==e-3,x=15 PT=F, CLS=F
MIT-67
LR=le-3
LR=1e-4, CLS==
1R=1 PT=F, CLS=F
Oxford Pets
IR=1e-3
IR=1e-4, CLS==F
I.R=le-3, PT=F, CLS=F
Stanford Dogs
LR=le-4
LR=1c-4, CLS=F
LR=1c4, PT=T, CLS=F
CUB200
LR=1e-3
LR=le-4, CLS=T
LR=1e-3, PT=F, CLS=F
Stanford Cars
LR=le-3
LR=1e-4, CLS=T
LR=1c-3, PT=F, CLS=T
FGVC-Airerafts
LR=1e-3
LR=1e-4, CLS=T
LR=1e-3, k=15, PT=F, CLS=T
Table 5: Best Hyperparameters Full Data, Last Attention Blocks
Dataset
NLFT-7
LFT-7
Caltech-256
LR=le-4
LR=II-4,CLS-EPT-T
MIT-67
LR=le-4
LR=Ie-,CCLS-EPTTTT
Oxford Pets
LR=le-4 LR=Ic-4,CLS-EPT-TT
Stanford Dogs
LR=1e-4 LR=1ee4,CIS-EFIITT
CUB200
LR=le-4 LR=Ie-,CC.S-E.IT-T
LR=le-4 Stanford Cars
LR=Ie-4.CLS-TPIIIT
FGVC-Aircrafts
LR=le-4
LR=Ie-4CC.S-T.PI=T
B DERIVATION OF LINEAR ATTENTION
We note that when taking the Taylor approximation for any (multivariable) function f. f(w+Au) f(w) + fff(ww Aw+ Aw" A fw) Aw O[]]Aw]] Awll? where O(-) notation hides the higher order terms, the first order term can be efficiently computed via its directional derivative
0f V vwf(w) Au lim f(w + rAw)
r 0r
where r a scalar variable. We will use this technique to derive the linearized closed form for the attention layer.
13

---

Table 6: Best Hyperparameters- Shards
Dataset
NLFT-1
LFT-1
Caltech-256
LR=le-3 LR=1c<3PP==FCLS-F
MIT-67
LR=le-3
LR=Ic-3.PT=ECLS-F
Oxford Pets
LR=le-3
LR-:Ie33.71=ECCSS--
Stanford Dogs
LR=le-3
LR=:I=3.P1=,CC..SPP
CUB200
LR=le-3 LR=1e-3.PT-ECCSS-
Stanford Cars
LR=le-3
LR=:c-3.PT=FCCSS-T
FGVC-Airerafts
LR=le-3
LR=Ie-3.PT=FCCSS-T
Let A denote the attention function parameterized by weights Wq, Wx. W.. We wish to derive a closed form expression for the linear attention Atn. which is defined as the first-order Taylor approximation of A parameterized by the new linearized weights AW, y: AW., AW...
A(x) = (x)) where ((x) = (Q(x))((())
(15)
(xx) (W a), K(x) 2),
(W,,
(16)
where is the soft-max activation function. As in the main paper, will write Q, K, V instead of Q(x), (x), (zz) for case of notation. We will derive the closed form for the single-headed attention, which can then be extended to multi-heades attention with minimal modification. Similarly,
we will use n = in the below proof (so x is a vector in Rd)) for simplicity, but note that the final result extends any 1.
a
Aui (x) = A(x) + lim A(x, W, + W + +AW,, W, +raw,)
(17)
--0 dr
A(x) lim a ((W, + rAW, (W +raw (W, + +rW. ..,,
(18)
Denote for ease of notation AQ = (AW q: x), AK = (AW,. AV = (AW, z). Then for each
component of vector s, we can write
= ((Q + rAQ) +RAK))* +rVV)
(19)
Applying chain rule, we get
lim r-0 Dr 8;
d' +(Q++AQ) (K +TAK)) (A4+ WTO0Z
Ox dr
AK + W Oz
))
lim r-0
K+ + Qi
Dr
+ lim a ((Q+r ((ra)))((++<)))++))))))))))))))) +rak)) (VV + W lim
dr
00
00
dr
T
d (QiK) ( AQ W lim dr -0 Dr
lim dz + Dr
K K+Q, Q. AK W
AV T lim dr 00 Dr
+ a (Q,K)"
[ddia((((Q: ( Vto "
[diag(".) Φ. v
V V++
[2,0V; (0/ 1/
14

---

where denote the Hadamard product. Hence, denoting v as the matrix with v, and the identity matrix, we obtain the desired result
(x) = A(x))
lim r-0 dr
(20)
=A(x) C ( V o4
(21)
C ADDITIONAL COMPARISONS
We discuss additional comparisons to Linearized ResNets in Sec. C.I, and detail training and inference times in Sec. C.2 We also compare our unlearning method with the paragon of re-training from scratch in Sec. and ablate on the pre-training schemes for initializing the tangent transformer in Sec.
C.I COMPARISON TO LINEARIZED RESNET ARCHITECTURES
The benefits of linearization rely the strength of the inductive prior obtained from pre- rre-training. Since vision transformers are shown to learn better inductive priors than convolutional architectures the scale of training data increases, we believe that linearized transformers yield clear advantage over linearized ResNets by being able to leverage the better inductive priors learnt from pre-training. We compare with linearized ResNet-50 in Tab. 7 where show that TAFT outperforms Lincarized ResNet-50 by 7.3% on the standard fine-tuning task, and by 9.0% for the parallel training and composition task (10 shards) averaged across datasets.
Table 7: Comparing linearized ResNet-50 (L-RN55) and linearized ViT-L/16 (TAFT) downstream classification tasks for both standard fine-tuning and parallel training/comppoitionn across 10 shards.
Dataset
Shards
L-RN50
TAFT
Caltech-256
85.5
95.9
MIT-67
79.3
89.3
Oxford Pets
93.1
94.5
Caltech-256
10
83.6
95.0
MIT-67
10
72.4
86.9
Oxford Pets
10
92.5
93.6
C.2 TRAINING AND INFERENCE TIME COMPARISONS
We compare the per-example training and inference wall-clock timings for NLFT and TAFT Tab, The inference and training cost for the linearized transformer is potentially twice of the original model discussed Sec. 5] We note that the train timings reported would be much faster practice due to large batch sizes and caching of intermediate features when limiting training to later layers.
Table 8: Comparison of per-example training and inference wall-clock timing (seconds) for NLFT and TAFT using batch size of These would be much faster in practice due large batch sizes and caching of intermediate features when limiting training to later layers. Timing is computed using the MIT-67 dataset.
NLFT (Train)
TAFT (Train)
NLFT (Inference)
TAFT (Inference)
0.147s
0.204s
0.021s
0.065s
C.3 COMPARISON WITH FORGETTING PARAGON
In Fig. compare the shard re-training forgetting method using TAFT the paragon of re- training from scratch. Both methods guarantee complete unlearning, but TAFT is able to achieve plose-to-paaagon performance while speeding up unlearning by up to 50x.
15

---

MIT-67
Oxford Pets
Stanford Dogs
100
100
Accuracy 1%
Accuracy 1%
Accuracy 1%
80
70
TAFT
70
70
TAFT
Paragon
Paragon
Paragon
60
60
60
10
20
30
40 50
10
20
30
40
50
10
20
30
40
50
Dataset Forget
% of Dataset to Forget
Dataset Forget
Figure 3: Shard re-training with TAFT (using sharding factor of 50) compared the Paragon method of re- re-training the non-linear model from scratch. While both method guarantee complete unlearning, pass--oo-
TAFT achieves close-to paragon performance while speeding up unlearning by up 50x.
C.4 ABLATION ON PRE-TRAINING SCHEME
In practice, the choice of model fine-tune for downstream tasks presupposes some relation between the latter tasks and those used in the pre-trained initialization. Since we focus on classification, we choose ImageNet classification pre-training as our initialization for all the experiments.
Here, we compare TAFT with different pre-training schemes: (1) Self supervised learning via MAE (Masked Autoencoder Training), (2) Supervised/Chassifcation pre-training. and (3) Contrastive Language-Image Image Pre-training (CLIP) followed by supervised pre-training.
We detail our results Tab. 9 Indeed, the performance from fine-tuning depends on the discrepancy between the pre-training objective and the target task. (1) being the farthest to classification performs worse than a classification pre-training. However, by augmenting supervised classification pre- training using& contrastive lunuuage-image pre-training objective, (3) further boosts the performance of classification-only pre-training.
Table 9: We compare fine-tuning from three different pre-training schemes. (1) MAE does self- supervised pre- me--raining via mask image modelling, (2) CLS uses ImageNet classification pre-training and (3) CLIP uses contrastive language-image image pre-training followed by tinettnnig on ImageNet classification. Since all MAE models are pre-trained on ImageNet IK, we use the T-ViT-B architecture to fairly compare MAE and CLS where ImageNet IK pre-trained models are available for both methods. The CLS T-VIT-L model is pre-trained on ImageNet 21K IK, while the CLIP model is pre-trained on WIT400M, ImageNet 12K + IK. The inductive priors learnt from MAE transfer less effectively the downstream classification tasks considered, where CLS the smaller T- T-VIT-B model able to outperform MAE on both T- --ViT-B and T-ViT-L. The inductive priors learnt with CLIP, which combines both unsupervised contrastive learning and supervised fine-tuning, transfer best to the downstream tasks.
Method
MAE/ / T-VIT-B MAE/ TAE/T-VE T-VT-L C-S/T-ViT-B CLS/T-VITL CLIP T-VITL
MIT-67
67.5
75.4
77.7
89.3
91.8
Oxford Pets
78.8
90.0
91.7
94.5
95.1
Stanford Dogs
62.3
74.7
90.8
91.9
95.2
C.5 INFLUENCE OF INDIVIDUAL. COMPONENT MODELS
Since models are composed via linear combinations of their weights, the influence of single component model can be quantified in at least two ways: (A) based on the difference in performance on validation dataset when the component model is added, and (B) based on the magnitude of the difference in weights with and without the component model. We explored (A) in Fig. ((a) where we show that subtracting models have lower impact the performance number of remaining component models is large. However when there remain only few component models in the composition, the impact of each model becomes larger.
downstream tasks when the
16

---

In Fig. 4 show that result of linearity, this effect is also reflected in the weight space via measuring the L2 difference in weights before and after adding the component model.
L2 difference in weights after adding new model
16
MIT-67 (50 shards)
14
Stanford Cars (50 shards)
CUB-200 (50 shards)
10
10 20 30
40
50
Number of models in composition
Figure 4: We plot the L2 change in weight space result of adding new component model against number of existing models in the composition. The impact of adding a new model is significantly larger when number of existing component models is small. Note that while plotted the same graph, the difference in scale between different datasets are not meant to be directly comparable due difference in number of output classes, amongst other factors.
C.6 TEXTURE CLASSIFICATION
In the main paper, we primarily evaluated our method on object classification tasks. In Tab. 10 we evaluate method the Describable Textures Dataset (DTD) (Cimpoi al. (014)) where show that even on texture classification tasks, composing models trained with TAFT consistently outperforms non-linear models across all sharding factors.
Table 10: We compare TAFT and Model Soup Wortsman et al. (2022b) in the same manner Tab, on the Describable Textures Dataset (DTD), and show that as a result of linearity, composing models trained with TAFT outperforms composing non-linear models across all sharding factors.
Method
10 Shards
25 Shards
50 Shards
Soup
68.2
59.8
47.6
TAFT
75.3
71.6
65.3
C.7 COMPARISON WITH PARAMETER EFFICCIET FINE-TUNING
In this section, we compare against parameter-efficient fine-tuning methods. In particular, we compare against Adapters (Houlsby et al. 2019) and Low-Rank Adaptation (LoRA) (Hu et al. 2021) when applied to the same last attention block as non-linear fine-tuning and TAFT. Since the main use cases of such methods lie in parameter efficiency and training speed, we show in Tab. IT that they typically exhibit lower performance downstream tasks compared full non-linear fine-tuning, and also lack the linearity of TAFT required to yield effective composition.
D TAFT WITH PROJECTED GRADIENT DESCENT
In the main paper, we constrain the distance that the flin moves from pretrained weights w by using the L2 weight decay penalty as regularizer during training, since the first-order taylor expansion is only valid around some local neighborhood w. However, we note that it is also possible impose hard constraint rather than soft constraint using projected gradient descent, where weights are projected onto a ball of radius R.
In Tab. 12 we disable weight decay and instead train with projected gradient descent. We compare using the RSL loss (with k = 5) and CE loss, since both losses differ in their effect on the final weight magnitude. We show that while RSL loss is more effective for the smaller radius R 1,
17

---

Table 11: We compare TAFT with parameter-efficient fine-tuning methods Adapter (Houlsby et al. 2019) and LoRA (Hu et al. 2021) in the same manner as Tab. on MIT-67, CUB-200 and Stanford Cars. We show that as result of linearity, composing models trained with TAFT outperforms
composing models fine-tuned with other methods across all sharding factors.
Shards
Dataset
Adapter
LoRA
Soup
TAFT
MIT-67
86.7
85.5
86.4
86.9
CUB-200
79.0
78.1
82.0
86.5
Stanford Cars
20.3
21.4
19.8
58.4
MIT-67
84.0
80.1
84.3
85.7
25
CUB-200
69.0
60.0
69.7
84.2
Stanford Cars
14.7
9.6
10.6
49.2
MIT-67
82.6
72.1
80.3
81.2
50
CUB-200
66.2
17.1
63.0
81.4
Stanford Cars
12.0
5.2
8.4
41.5
Table 12: Projected Gradient Descent (PGD) onto ball of radius R. While imposing soft constraints through weight decay is more effective than PGD, the hard constraints on weight magnitude provide several benefits for bridging theoretical analysis and empirical results of deep neural networks.
Dataset
R
CE Loss
RSL Loss
1
37.6
52.1
Stanford cars
10
57.5
52.1
1
85.4
86.3
CUB
10
88.4
86.4
40.4
55.1
Aircrafts
10
58.5
55.4
CE loss becomes more effective as the radius increases to R 10. While imposing such hard constraints generally yield worse results compared to TAFT, we note that this can be useful in several applications, such as for estimating the smoothness constant of the Tangent Transformer. This can help to better bridge the gap between theoretical analysis which generally require L-smoothness assumptions or convex loss objectives and empirical applications.
18


Reviews:
[{"review_id": "N7SJCz1jaB", "reviewer": "Reviewer_Ef2Q", "summary": "This paper introduces Tangent Attention Fine-Tuning (TAFT), a method for fine-tuning linearized transformers that are derived by computing a First-order Taylor Expansion around a pre-trained initialization. The key contributions of the paper are as follows:\n\nEfficient Jacobian-Vector Product Calculation: The authors demonstrate that the Jacobian-Vector Product resulting from linearization can be efficiently computed in a single forward pass. This reduces the training and inference costs of linearized transformers to a similar order of magnitude as the original non-linear models, all while maintaining the same number of parameters.\n\nTAFT presents an efficient and effective method for fine-tuning linearized transformers obtained through Taylor Expansion. It maintains performance parity with non-linear models while providing advantages related to model composition, training efficiency, unlearning, and privacy. This has significant implications for the practical use of transformers in various downstream tasks.", "questions": "1. How about interpretability? A detailed analysis of how a given training sample affects the learned model and the predicted results is given in LQF. Is it possible for the authors to provide an analysis of the interpretability?\n2. The authors detail the advantages of linear models in the introduction, but these advantages don't seem to be relevant to transformer, what are the advantages and disadvantages of linearizing the transformer model compared to linearizing ResNet?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. Efficient Jacobian-Vector Product Calculation: The authors demonstrate that the Jacobian-Vector Product resulting from linearization can be efficiently computed in a single forward pass. This reduces the training and inference costs of linearized transformers to a similar order of magnitude as the original non-linear models, all while maintaining the same number of parameters.\n\n2. Comparable Performance: When applied to various downstream visual classification tasks, the Tangent Transformer fine-tuned with TAFT performs on par with fine-tuning the original non-linear network. This suggests that the linearized version is a viable alternative without compromising performance.\n\n3. Convex Fine-Tuning Loss: The paper highlights that Tangent Transformers are linear concerning a new set of weights, resulting in a convex fine-tuning loss. This convexity offers several advantages over non-linear fine-tuning, particularly in terms of model composition, parallel training, machine unlearning, and differential privacy.", "weaknesses": "1. In Section 3.4, the author mentions basically the same things as in Section 2 Related work-Pravicy with no new theoretical analysis about differential privacy.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yETZQstXN5", "reviewer": "Reviewer_QE2Q", "summary": "The authors introduce Tangent Attention Fine-Tuning (TAFT) for fine-tuning linearized transformers. It can perform comparably with fine-tuning the original non-linear network in various downstream visual classification tasks. It enjoys several advantages compared to non-linear fine-tuning when it comes to model composition, parallel training, machine unlearning, and differential privacy.", "questions": "see weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is the first work to propose an efficient method to linearize models in the Transformer family of architectures, which is meaningful. The paper is clearly written with many experiments.", "weaknesses": "1.since this is a fine-tuning method, please provide more fine-tuning methods for comparison (lora,adapter...) in table 3. \n2.I wonder how TAFT works when applied to LLM? maybe some experiments can be added.\n3.please derive in detail how to get the closed form expression in equation 5 and 6.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0BKQDybnpp", "reviewer": "Reviewer_ZQ32", "summary": "The paper proposed a method for fine-tuning linearized transformers, called Tangent Attention Fine-Tuning (TAFT). \n\nThe paper claims that TAFT can achieve comparable performance to non-linear fine-tuning on various downstream tasks while enjoying the benefits of linearity such as compositionality, parallel training, machine unlearning, and differential privacy. \n\nThe paper also introduces Tangent Transformers, which are linearized versions of pre-trained transformer models, and shows how to compute their Jacobian-Vector products efficiently in a single forward pass. \n\nThe paper demonstrates the advantages of TAFT and Tangent Transformers in several experiments using vision transformer models. \n\nTo summarize, the paper's main contributions are:\n\n- A novel method for fine-tuning linearized transformers that is computationally efficient and competitive with non-linear fine-tuning.\n\n- A theoretical analysis of the benefits of linearity for model composition, parallel training, machine unlearning, and differential privacy.\n\n- An empirical evaluation of TAFT and Tangent Transformers on various downstream visual classification tasks using vision transformer models.", "questions": "(1) Please give a detailed derivation process from Eq 5 to Eq 6 in the appendix.\n\n(2) The selected datasets (Caltech-256 (Griffin et al., 2007), MIT-67 (Quattoni & Torralba, 2009), Oxford Pets (Parkhi et al.,\n2012), Stanford Dogs (Khosla et al., 2011), CUB-200 (Wah et al., 2011), FGVC-Aircrafts (Maji\net al., 2013), and Stanford Cars (Krause et al., 2013)) are object-oriented. What about the performance of the proposed method on Describable Textures (DTD) (Cimpoi et al., 2014)？ \n\n(3) (Optional to reply) Could the proposed method be applied to temporal classification tasks, such as activity classification on UCF101?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper proposes a novel method for fine-tuning linearized transformers, called Tangent Attention Fine-Tuning (TAFT), which is computationally efficient and competitive with non-linear fine-tuning. The paper also introduces Tangent Transformers, which are linearized versions of pre-trained transformer models, and shows how to compute their Jacobian-Vector products efficiently in a single forward pass. The paper demonstrates the advantages of TAFT and Tangent Transformers on several experiments using vision transformer models.\n\nIn terms of originality, the paper introduces a new method for fine-tuning linearized transformers that is computationally efficient and competitive with non-linear fine-tuning. The paper also introduces Tangent Transformers, which are linearized versions of pre-trained transformer models, and shows how to compute their Jacobian-Vector products efficiently in a single forward pass. These contributions are novel and have not been explored before.\n\nIn terms of quality, the paper provides a theoretical analysis of the benefits of linearity for model composition, parallel training, machine unlearning, and differential privacy. The paper also provides an empirical evaluation of TAFT and Tangent Transformers on various downstream visual classification tasks using vision transformer models. The experiments are well-designed and the results are presented clearly.\n\nIn terms of clarity, the paper is well-written and easy to follow. The authors provide clear explanations of their methods and results. The paper also includes visualizations that help to illustrate the concepts presented.\n\nIn terms of significance, the paper’s contributions have important implications for the field of machine learning. The proposed method for fine-tuning linearized transformers is computationally efficient and competitive with non-linear fine-tuning. This has important implications for large-scale applications where computational efficiency is critical. Additionally, the theoretical analysis of the benefits of linearity has important implications for model composition, parallel training, machine unlearning, and differential privacy.\n\nOverall, this is a well-written paper that makes significant contributions to the field of machine learning.", "weaknesses": "(1) While the paper is well-written and easy to follow, it would be helpful to provide more detailed explanations of some of the concepts presented. For example, the paper could provide more details on how Tangent Transformers are computed and how they are used in practice (especially from Eq 5 to Eq 6).\n\n(2) While the paper provides an empirical evaluation of TAFT and Tangent Transformers on various downstream visual classification tasks using vision transformer models, it would be helpful to see more experiments that compare TAFT with other fine-tuning methods on these tasks. This would help to establish the competitiveness of TAFT more clearly.\n\n(3) While the paper provides a theoretical analysis of the benefits of linearity for model composition, parallel training, machine unlearning, and differential privacy, it would be beneficial to provide more empirical evidence to support these claims. Specifically, it would be helpful to see more experiments that demonstrate the advantages of TAFT and Tangent Transformers on a wider range of tasks and datasets.\n(e.g. DTD and UCF101).\n\nOverall, this is a well-written paper that makes some contributions to the field. However, there is room for improvement in terms of providing more empirical evidence to support the claims made in the paper and providing more detailed explanations of some of the concepts presented.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)