Tangent Transformers for Composition, Privacy and Removal

TL;DR

Tangent Transformers, fine-tuned via TAFT, are linearized models that match the performance of non-linear transformers on visual tasks while enabling efficient, privacy-preserving, and compositional fine-tuning methods.

Contribution

The paper introduces TAFT, a novel linearization-based fine-tuning method for transformers that improves efficiency and privacy while maintaining performance.

Findings

Tangent Transformers perform comparably to original models on visual classification tasks.

TAFT enables efficient training and inference with reduced computational costs.

Linearized models facilitate model composition, unlearning, and differential privacy.

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