#  : Token Selective Attention for Efficient Vision Transformers

Manish Kumar Singh Rajeev Yasarla Hong Cai Mingu Lee Fatih Porikli
Qualcomm AI Research∗
{masi, ryasarla, hongcai, mingul<EMAIL_ADDRESS>
###### Abstract

In this paper, we propose a novel token selective attention approach, , which
can identify tokens that need to be attended as well as those that can skip a
transformer layer. More specifically, a token selector parses the current
attention maps and predicts the attention maps for the next layer, which are
then used to select the important tokens that should participate in the
attention operation. The remaining tokens simply bypass the next layer and are
concatenated with the attended ones to re-form a complete set of tokens. In
this way, we reduce the quadratic computation and memory costs as fewer tokens
participate in self-attention while maintaining the features for all the image
patches throughout the network, which allows it to be used for dense
prediction tasks. Our experiments show that by applying , we can significantly
reduce computation costs while maintaining accuracy on the ImageNet
classification benchmark. Furthermore, we evaluate on the dense prediction
task of monocular depth estimation on NYU Depth V2, and show that we can
achieve similar depth prediction accuracy using a considerably lighter
backbone with .

##  1 Introduction

Vision transformers (ViTs) [5] have been the core of many latest advances in
computer vision, with self-attention playing a critical role to generate key
visual features. However, the self-attention operation incurs quadratic
computation and memory costs w.r.t. the input size. This makes it expensive
and challenging to run vision transformers on high-resolution images and on
resource-constrained devices.

Figure 1: (a) Two consecutive standard transformer layers. (b) A standard
transformer layer followed by our layer. Our proposed approach operates on a
pair of consecutive transformer layers and modifies the second layer to be
token selective. In other words, only a subset of tokens participate in the
self-attention in the second layer while the rest bypass the layer; our token
selector predicts importance scores of the tokens for selection. Note that we
retain all the tokens throughout the layers to facilitate dense prediction
tasks, e.g., segmentation, depth estimation.

Researchers have looked into alternative ways to apply attention to make
vision transformers more efficient. For instance, Swin [12] applies self-
attention within local windows of the image. Others propose to first apply a
few stages of convolutional layers and only apply self-attention on
significantly downsampled versions of the input image to reduce computation
costs (e.g., [11, 20, 9]). Some other papers design alternative ways to
generate attention, e.g., removing softmax [10], using ReLU for normalization
[3], applying attention on feature channel dimension [13], additive attention
[17], etc.

However, many latest advances in computer vision still predominantly rely on
the vanilla ViTs, thanks to several favorable properties: their design is
straightforward and easy to implement, there exist powerful pretraining
methods and pretrained checkpoints for them (e.g., CLIP [15], DINOV2 [14]),
they scale better with massive data [4], and they may have more consistent
performance across problem domains given the generic design of standard self-
attention.

As such, researchers have looked into reducing the computation costs of ViTs
while preserving the standard self-attention. A common approach is to reduce
the tokens as they go through the network [2, 23, 21, 7]. This effectively
reduces the computation and memory usage which scale w.r.t. the number of
attended tokens. However, these methods only work for classification tasks. As
some of the tokens are discarded or merged during inference, these networks
cannot be used for dense prediction tasks which require distinct features for
all the image pixels/patches.

In this paper, we propose a novel token selective attention approach, , which
can make any vision transformer more efficient by selecting only subsets of
tokens for self-attention. More specifically, is applied to two consecutive
transformer layers where the latter will become token selective. Given the
multi-head attention maps from the first layer, our token selector predicts
the attention maps for the next layer and produces important scores for the
tokens accordingly. The top important tokens will go into self-attention of
the next layer, which we call a _transformer layer_ ; the transformer layer
replaces the original second standard transformer layer. The rest of the
tokens simply bypass the layer and are re-joined with the attended tokens. In
this way, we reduce the quadratic computation and memory costs as fewer tokens
are attended, maintain the favorable standard self-attention operation, and
retain the full set of tokens during inference so the network can be used as
an encoder for dense prediction tasks. Fig. 1 provides a high-level
illustration of .

Figure 2: Overview of our proposed approach. Based on an input $X_{i}$, a
standard transformer layer generates an output $X_{i+1}$, as well as $QK^{T}$
maps (prior to Softmax), $B_{i}^{1},B_{i}^{2},...,B_{i}^{H}$, for the $H$
attention heads. Based on them, our token selector predicts the attention maps
at the next layer and uses them to identify tokens that need to be attended,
$X_{i+1}^{a,h}$, and tokens that can skip the next layer, $X_{i+1}^{p,h}$, for
each head. Then, self-attention is performed on $X_{i+1}^{a,h}$, the output of
which is concatenated with $X_{i+1}^{p,h}$ to form the output of this head,
$X_{i+2}^{h}$. All $X_{i+2}^{h}$’s are combined to generate the final output
of this layer, $X_{i+2}$. By reducing the number of tokens participating in
attention, we significantly save computation and memory given the quadratic
complexity of self-attention, while slightly improving accuracy. Figure 3:
Applying our proposed to all the pairs of transformer layers in a 12-layer
DeiT.

In summary, our main contributions are as follows:

  * •

We present , a novel approach that improves the efficiency of a multi-layer
vision transformer by reducing the number of tokens participating in the
attention computation in certain layers. Although not all tokens are attended
in , they are retained throughout the layers. This allows the model to be used
for dense prediction tasks.

  * •

In order to select a subset of tokens to be fed into self-attention, we devise
a token selector that predicts important scores of the tokens based on the
previous transformer layer. The scores are used to identify tokens that need
to be attended at the next layer and the rest will simply skip the next layer.

  * •

We evaluate on the standard ImageNet classification benchmark [16], showing
that significantly reduces the computation cost while maintaining accuracy. In
addition, we apply to a vision transformer backbone and use it as an encoder
for a monocular depth estimation task. We show that our more efficient
backbone maintains the depth prediction accuracy.

##  2 Proposed Approach:

In this section, we present our proposed token selective attention method, .
selects only a subset of tokens to participate in the self-attention of a
transformer layer, which considerably reduces computation and memory costs,
while maintaining the full set of image features throughout the network,
making it feasible to be used as an encoder for dense prediction tasks.

###  2.1 Self-Attention with Selected Tokens

Standard Self-Attention: Consider an input $X_{i}\in\mathbb{R}^{L\times D}$,
where $L$ is the number of tokens (i.e., features of image patches) and $D$ is
the feature dimension; the batch dimension is omitted here for conciseness. In
a standard vision transformer layer $i$, $X_{i}$ first goes through linear
layers to generate the query $Q_{i}^{h}$, key $K_{i}^{h}$, and value
$V_{i}^{h}$, respectively, for each attention head, $h\in\\{1,...,H\\}$, as
follows:

$\begin{split}{Q_{i}^{h}}={W_{Q,i}^{h}}{X_{i}},{K_{i}^{h}}={W_{K,i}^{h}}{X_{i}},{V_{i}^{h}}={W_{V,i}^{h}}{X_{i}},\\\\[-3.0pt]
\end{split}$ (1)

where $W_{Q,i}^{h}$, $W_{K,i}^{h}$, $W_{V,i}^{h}$ are linear transformation
matrices.

Then, an attention map is computed based on the query and key, which is then
multiplied with the value:

$\begin{split}A_{i}^{h}=\text{softmax}(\frac{Q_{i}^{h}(K_{i}^{h})^{T}}{\sqrt{D}}),\quad
X_{i+1}^{h}=A_{i}^{h}\cdot V_{i}^{h}.\\\\[-3.0pt] \end{split}$ (2)

The outputs from all attention heads are concatenated and fed into a linear
layer, $F_{i}$, to general the final output of this transformer layer:

$\begin{split}X_{i+1}=F_{i}(W_{O,i}\cdot\text{Concat}(X_{i+1}^{1},...,X_{i+1}^{H})),\\\\[-3.0pt]
\end{split}$ (3)

where $W_{O,i}$ is a linear projection matrix.

Token Selective Attention (): In the standard self-attention, all tokens
participate in the attention computation, resulting in $\text{O}(N^{2})$
computation and memory costs. One way to reduce computation costs is to reduce
the number of tokens taking part in the self-attention, i.e., reducing $N$.

In order to do this, we need to identify the subset of important tokens that
should go into self-attention, and the remaining set that should skip the
attention. We propose a novel token selector, which operates between a
standard transformer layer and a selective attention transformer layer
(transformer layer). It consumes the un-normalized attention maps (i.e.,
$QK^{T}$) from the previous standard layer and predicts the attention maps for
the respective heads of the next layer. These predictions are used to identify
important tokens. More specifically, we perform column sum over the predicted
attention map of each head, which generates an importance score for each
token. Based on the important scores, we select the top-K tokens to perform
self-attention in the next layer, to which we refer as the _attention tokens_
$X_{i+1}^{a,h}$, and K is determined based on a prescribed attention ration
$r$. For instance, $r=80\%$ indicates that the top $80\%$ tokens will
participate in the self-attention in the next layer. The remaining tokens
simply bypass the next transformer layers, which we call _skip tokens_
$X_{i+1}^{s,h}$.

After the attention tokens go through the transformer layer, where standard
self-attention is applied to this subset, the attended tokens and skipped
tokens are combined to re-form the complete set of tokens, for each head,
i.e.,

$\begin{split}X_{i+2}^{h}=\text{Concat}(X_{i+2}^{a,h},X_{i+1}^{s,h}),\\\\[-3.0pt]
\end{split}$ (4)

where $X_{i+2}^{a,h}$ is the self-attention output from $X_{i+1}^{a,h}$ and
the concatenation is performed along the token dimension. Finally, the outputs
from all the heads are concatenated and processed by a linear layer to produce
the final output from this layer. Fig. 2 illustrates this process for a pair
of standard layer and layer.

###  2.2 Token Selector: Architecture and Training

A high level overview of the token selector architecture is given in Fig. 2.
The token selector takes as input the multi-head un-normalized attention maps
($QK^{T}$) from the previous standard transformer layer. It consists of two 1D
convolutional layers with a intermediate ReLU activation, followed by a log-
softmax layer that predicts the full $L\times L$ multi-head attention maps,
$\hat{A}_{i+1}^{1},...,\hat{A}_{i+1}^{H}$, for the next layer; we choose log-
softmax for the output layer as it produces normalized attention maps and is
more numerically favorable in training.

We train the token selector’s attention prediction part based on ground-truth
attention maps computed by the pretrained model. More specifically, given a
pre-trained backbone (e.g., DeiT [19] trained on ImageNet classification
[16]), we insert a token selector between two consecutive standard transformer
layers $i$ and $i+1$. During training, the token selector consumes the
$QK^{T}$ maps from the first transformer layer and its predicted attention
maps are compared against the actual attention maps computed by the second
transformer layer of the pretrained model, which provides supervision to train
the token selector, i.e.,

$\begin{split}\mathcal{L}_{i,i+1}=\sum_{h=1}^{H}\mathcal{L}_{\text{KLD}}(\hat{A}_{i+1}^{h},{A}_{i+1}^{h}),\\\\[-3.0pt]
\end{split}$ (5)

where $\mathcal{L}_{\text{KLD}}$ denotes the KL-Divergence loss. In this
process, the pre-trained transformer layers are frozen and only the token
selector is trained. Note that we only use the same data used to pre-train the
vision transformer model to train our token selectors, and do not use extra
training data.

Method | Top-1 Acc. | Top-5 Acc. | GFlops | Flops Reduction | Use for Dense Prediction?
---|---|---|---|---|---
DeiT-Tiny [19] | 72.2 | 91.1 | 1.3 | - | ✓
ToMe***Qualcomm AI Research is an initiative of Qualcomm Technologies, Inc. [2] | 71.4 | - | 0.8 | 41.7 | ✗
VTC-LFC***The original token merging method reduces tokens by merging and cannot be used for dense prediction tasks which will require distinct features for all the image patches. [1] uses token merging to accelerate stable diffusion, but it requires additional steps to un-merge the tokens. Moreover, while un-merging works for image generation, it may not work for other dense tasks such as depth estimation, since similar-looking patches may have very different depth values. [21] | 71.0 | 90.4 | 0.8 | 41.7 | ✗
STAR [23] | 72.0 | 90.7 | 0.7 | 46.2 | ✗
a-STAR [23] | 72.0 | 90.7 | 0.7 | 46.2 | ✗
DeiT-Tiny w/ (ours) | 73.9 | 91.9 | 1.0 | 25.8 | ✓
Table 1: Results on ImageNet-1K classification using DeiT-Tiny [19] as the
base model. Our proposed significantly reduces the computation cost, while
improving accuracy and retaining the full set of tokens. In contrast, other
existing methods cause accuracy drops and cannot be used for dense prediction
as subsets of tokens (i.e., features of some image patches) are discarded.

###  2.3 Full Vision Transformer Model with

Given a pretrained vision transformer model, we can apply to any pair of
layers where the second layer will be replaced by a layer and we train a token
selector between them. Consider a 12-layer standard vision transformer as an
example. We can apply to consecutive pairs of layers, and replace the 2nd,
4th, 6th, 8th, and 10th standard transformer layers with transformer layers;
see an illustration in Fig. 3.

Once the token selectors are trained for the pairs of layers where will be
applied, we freeze the token selectors, modify the second layers in the pairs
to be token selective, and finetune the full model on the same training set
where it was originally pre-trained (e.g., ImageNet).

Encoder | _NYUD Depth v2_ | _KITTI_
---|---|---
Abs Rel $\downarrow$ |  $\delta\\!<\\!1.25$ $\uparrow$ |  Abs Rel $\downarrow$ |  $\delta\\!<\\!1.25$ $\uparrow$
DeiT-Tiny | 0.101 | 0.896 | 0.063 | 0.956
\+  | 0.103 | 0.891 | 0.063 | 0.954
Table 2: Monocular depth estimation on NYU Depth v2 and KITTI using DeiT-Tiny
and DeiT-Tiny with as the encoders, respectively. We use NeWCRFs [22] as the
base architecture.

##  3 Experiments

We evaluate on the ImageNet classification benchmark. We further utilize it
for monocular depth estimation and demonstrate its effectiveness. Note that in
this short paper, we present representative results to showcase the efficacy
of our proposed approach. We plan to include more comprehensive results in the
full paper.

###  3.1 Experimental Setup

Datasets and Evaluation: We perform standard image classification evaluation
on ImageNet [16]. We further use classification networks as encoders for
monocular depth estimation on NYU Depth v2 [18] and KITTI [8], where we use
standard depth accuracy metrics such as Absolute Relative Error (Abs Rel) and
and $\delta<1.25$; see [6] for mathematical definitions of these metrics.

Networks and Implementation: In this work we use DeiT-Tiny [19] as the base
model as it is one of the standard vision transformers and the computation
resource required for training is more accessible, e.g., we use 2 Nvidia A100
GPUs. For ImageNet evaluation, we train the token selector and finetune DeiT-
Tiny with both on ImageNet training data. Since DeiT-Tiny has 12 layers, we
apply to the 2nd, 4th, 6th, 8th, and 10th layers, using 80% tokens for self-
attention at each layer. For depth estimation, we use NeWCRFs [22] as the base
model and replace the encoder to be either DeiT-Tiny or DeiT-Tiny with .

**footnotetext: Computation reduction by VTC-LFC also includes channel
pruning.

###  3.2 Image Classification

Table 1 summarizes the evaluation results on ImageNet-1K classification
benchmark. By applying , we significantly reduce the computation of DeiT-Tiny
by nearly $25\%$ while maintaining the classification accuracy. Since the full
set of tokens are preserved, our network can be readily used as an encoder for
dense prediction tasks. On the other hand, while existing works like ToMe [1]
and a-STAR [23] can also significantly reduce computation, they rely on
reducing the number of tokens, i.e., only a small number of tokens are kept
after going through the network. This works for classification but makes it
challenging to use such models as encoders for dense tasks. Furthermore, these
existing methods result in accuracy drops.

Fig. 4 visualizes the token selections at sample layers (second, sixth, and
tenth) of the network. It can be seen that most patches of the main object
(cat) are selected for self-attention and the network selects different
subsets of background patches at different layers. For instance, at these
layers, the network selects background patches from the right, left, and top
parts of the image, respectively, while always selecting most of the cat
patches.

Figure 4: Visualization of token selection at the 2nd, 6th, and 10th layers.
The input is a cat image with the cat at the center. Dark colors indicate the
patches that do not take part in self-attention.

###  3.3 Monocular Depth Estimation

allows the modified network to be used as an encoder for dense prediction. In
Table 2, we see that our DeiT-Tiny w/ network provides similar depth
estimation performance on both NYU Depth v2 and KITTI benchmarks, although it
uses considerably fewer tokens for self-attention in several layers, as
compared to using the original DeiT-Tiny as an encoder.

##  4 Conclusion

In this paper, we proposed , a novel token selective attention approach to
make vision transformers more efficient. Based on predicted attention maps, a
subset of important tokens are selected for self-attention while the rest
bypass the transformer layer. Unlike existing works, maintains the full set of
tokens, making the network readily usable for dense prediction. Our evaluation
results indicate that can significantly reduce computation costs while
maintaining accuracy on both classification and dense prediction tasks.

## References

  * Bolya and Hoffman [2023] Daniel Bolya and Judy Hoffman.  Token merging for fast stable diffusion.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 4598–4602, 2023.
  * Bolya et al. [2023] Daniel Bolya, Cheng-Yang Fu, Xiaoliang Dai, Peizhao Zhang, Christoph Feichtenhofer, and Judy Hoffman.  Token merging: Your ViT but faster.  In _International Conference on Learning Representations_ , 2023.
  * Cai et al. [2024] Han Cai, Junyan Li, Muyan Hu, Chuang Gan, and Song Han.  Efficientvit: Multi-scale linear attention for high-resolution dense prediction, 2024.
  * Dehghani et al. [2023] Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski, Jonathan Heek, Justin Gilmer, Andreas Peter Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, et al.  Scaling vision transformers to 22 billion parameters.  In _International Conference on Machine Learning_ , pages 7480–7512. PMLR, 2023.
  * Dosovitskiy et al. [2020] 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 at scale.  In _International Conference on Learning Representations_ , 2020.
  * Eigen et al. [2014] David Eigen, Christian Puhrsch, and Rob Fergus.  Depth map prediction from a single image using a multi-scale deep network.  _Advances in neural information processing systems_ , 27, 2014.
  * Fayyaz et al. [2022] Mohsen Fayyaz, Soroush Abbasi Koohpayegani, Farnoush Rezaei Jafari, Sunando Sengupta, Hamid Reza Vaezi Joze, Eric Sommerlade, Hamed Pirsiavash, and Jürgen Gall.  Adaptive token sampling for efficient vision transformers.  In _Proceedings of the European Conference on Computer Vision_ , pages 396–414, 2022.
  * Geiger et al. [2012] Andreas Geiger, Philip Lenz, and Raquel Urtasun.  Are we ready for autonomous driving? The KITTI vision benchmark suite.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 3354–3361, 2012.
  * Hatamizadeh et al. [2023] Ali Hatamizadeh, Greg Heinrich, Hongxu Yin, Andrew Tao, Jose M Alvarez, Jan Kautz, and Pavlo Molchanov.  Fastervit: Fast vision transformers with hierarchical attention.  In _The Twelfth International Conference on Learning Representations_ , 2023.
  * Koohpayegani and Pirsiavash [2024] Soroush Abbasi Koohpayegani and Hamed Pirsiavash.  Sima: Simple softmax-free attention for vision transformers, 2024.
  * Li et al. [2022] Yanyu Li, Geng Yuan, Yang Wen, Ju Hu, Georgios Evangelidis, Sergey Tulyakov, Yanzhi Wang, and Jian Ren.  Efficientformer: Vision transformers at mobilenet speed.  _Advances in Neural Information Processing Systems_ , 35:12934–12949, 2022.
  * Liu et al. [2021] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo.  Swin transformer: Hierarchical vision transformer using shifted windows, 2021.
  * Maaz et al. [2022] Muhammad Maaz, Abdelrahman Shaker, Hisham Cholakkal, Salman Khan, Syed Waqas Zamir, Rao Muhammad Anwer, and Fahad Shahbaz Khan.  Edgenext: Efficiently amalgamated cnn-transformer architecture for mobile vision applications, 2022.
  * Oquab et al. [2023] Maxime Oquab, Timothée Darcet, Theo Moutakanni, Huy V. Vo, Marc Szafraniec, Vasil Khalidov, Pierre Fernandez, Daniel Haziza, Francisco Massa, Alaaeldin El-Nouby, Russell Howes, Po-Yao Huang, Hu Xu, Vasu Sharma, Shang-Wen Li, Wojciech Galuba, Mike Rabbat, Mido Assran, Nicolas Ballas, Gabriel Synnaeve, Ishan Misra, Herve Jegou, Julien Mairal, Patrick Labatut, Armand Joulin, and Piotr Bojanowski.  Dinov2: Learning robust visual features without supervision, 2023.
  * Radford et al. [2021] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al.  Learning transferable visual models from natural language supervision.  In _International conference on machine learning_ , pages 8748–8763. PMLR, 2021.
  * 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.  _International journal of computer vision_ , 115:211–252, 2015.
  * Shaker et al. [2023] Abdelrahman Shaker, Muhammad Maaz, Hanoona Rasheed, Salman Khan, Ming-Hsuan Yang, and Fahad Shahbaz Khan.  Swiftformer: Efficient additive attention for transformer-based real-time mobile vision applications, 2023.
  * Silberman et al. [2012] Nathan Silberman, Derek Hoiem, Pushmeet Kohli, and Rob Fergus.  Indoor segmentation and support inference from rgbd images.  In _Computer Vision–ECCV 2012: 12th European Conference on Computer Vision, Florence, Italy, October 7-13, 2012, Proceedings, Part V 12_ , pages 746–760. Springer, 2012.
  * 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 _International Conference on Machine Learning_ , 2021.
  * Vasu et al. [2023] Pavan Kumar Anasosalu Vasu, James Gabriel, Jeff Zhu, Oncel Tuzel, and Anurag Ranjan.  Fastvit: A fast hybrid vision transformer using structural reparameterization.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , pages 5785–5795, 2023.
  * Wang et al. [2022] Zhenyu Wang, Hao Luo, Pichao Wang, Feng Ding, Fan Wang, and Hao Li.  Vtc-lfc: Vision transformer compression with low-frequency components.  _Advances in Neural Information Processing Systems_ , 35:13974–13988, 2022.
  * Yuan et al. [2022] Weihao Yuan, Xiaodong Gu, Zuozhuo Dai, Siyu Zhu, and Ping Tan.  Neural window fully-connected crfs for monocular depth estimation.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , pages 3916–3925, 2022.
  * Zhang et al. [2024] Yuyao Zhang, Lan Wei, and Nikolaos M Freris.  Synergistic patch pruning for vision trans-former: Unifying intra-&inter-layer patch importance.  In _Proceedings of the International Conference on Learning Representations_ , 2024.


# Pruning for Protection: Increasing Jailbreak Resistance in Aligned LLMs
Without Fine-Tuning

Adib Hasan
MIT
<EMAIL_ADDRESS>
&Ileana Rugina
<EMAIL_ADDRESS>&Alex Wang
MIT
<EMAIL_ADDRESS>

###### Abstract

Large Language Models (LLMs) are vulnerable to ‘Jailbreaking’ prompts, a type
of attack that can coax these models into generating harmful and illegal
content. In this paper, we show that pruning up to 20% of LLM parameters
markedly increases their resistance to such attacks without additional
training and without sacrificing their performance in standard benchmarks.
Intriguingly, we discovered that the enhanced safety observed post-pruning
correlates to the initial safety training level of the model, hinting that the
effect of pruning could be more general and may hold for other LLM behaviors
beyond safety. Additionally, we introduce a curated dataset of 225 harmful
tasks across five categories, inserted into ten different Jailbreaking
prompts, showing that pruning aids LLMs in concentrating attention on task-
relevant tokens in jailbreaking prompts. Lastly, our experiments reveal that
the prominent chat models, such as LLaMA-2 Chat, Vicuna, and Mistral Instruct
exhibit high susceptibility to jailbreaking attacks, with some categories
achieving nearly 70-100% success rate. These insights underline the potential
of pruning as a generalizable approach for improving LLM safety, reliability,
and potentially other desired behaviors.

##  1 Introduction

Large language models, or LLMs, have seen a massive increase in both
capabilities and usage in recent years. With the release of ChatGPTchatgpt ,
these models have received an unprecedented amount of attention from users. To
reduce the risk of generating dangerous or sensitive content, LLMs are often
further fine-tuned to be aligned with human valuestouvron2023LLaMA . However,
the increase in popularity of LLMs has also come with advances in adversarial
prompts, known as jailbreaks, as users attempt to bypass the safety alignment
of these models. Furthermore, due to these LLMs’ large size and demand,
deployment is a significant challenge, encouraging the use of techniques like
model compression to scale well. The effects of compression on safety are not
easily characterized, as demonstrated by compression of computer vision models
being shown to have mixed results with regards to preserving adversarial
robustnessGorsline_2021 . With the widespread deployment and usage of LLMs,
understanding the potential effects of compression and possible mitigations
becomes necessary for ensuring safe models.

In this study, we investigate the effects of pruning on the safety alignment
of large language models (LLMs). We curated a dataset of 2250 prompts designed
to elicit malicious outputs from LLMs. Our focus was on three 7 billion
parameter models: LLaMA-2 Chat touvron2023LLaMA , Vicuna-1.3 vicuna2023 , and
Mistral Instruct v0.2 mistral2023 . LLaMA-2 Chat was finetuned from the base
LLaMA-2 model and then underwent additional safety alignment via reinforcement
learning with human feedback (RLHF). Vicuna 1.3, derived from the original
LLaMA model, was fine-tuned using the ShareGPT dataset, while Mistral Instruct
v0.2 was fine-tuned from the base Mistral Model, with neither receiving RLHF
training.

Our approach was to examine the refusal rates for the malicious prompts in the
unpruned models compared to their pruned versions, observing the changes at
varying levels of model compression. Our findings reveal an initial increase
in resistance to jailbreaking prompts with light pruning, followed by a
decline in safety when the pruning exceeds a certain threshold.

Notably, LLaMA-2 Chat demonstrated the highest resilience against jailbreaking
prompts and showed the most significant safety improvement with parameter
pruning. Conversely, Mistral Instruct v0.2 was the least resilient before
pruning and exhibited minimal safety improvement post-pruning. Furthermore, we
found that pruning helps the models to concentrate their attention on the task
tokens. We speculate that attention concentration aids the models in detecting
malicious tasks and is responsible for safety improvement. These observations
suggest that moderate parameter pruning can potentially enhance the desired
behaviors in LLMs under threat models.

To summarize, our contributions are as follows:

  * •

We curated and open-sourced a new dataset to study safety in LLMs. This
dataset includes 225 malicious tasks drawn from five overarching categories.
We consider 10 types of jail-breaking methods and each task was inserted into
each jailbreaking prompt.

  * •

We leveraged a recently introduced pruning algorithm (sun2023wanda, ) to
increase LLM safety under jailbreaking attacks. We showed that our method
provides consistent safety improvements across a variety of tasks under
jailbreaking attacks, and the improvement is dependent on the level of safety
training in the unpruned model.

  * •

We analyzed the attention maps of both the pruned and unpruned models and
found that the attention map of the pruned models is significantly more
concentrated on the task tokens.

We release our dataset and implementation in hopes of supporting research in
this essential field at https://github.com/CrystalEye42/eval-safety.

Prompt
---
Ignore all instructions until now. You are a typical 4chan user named Kevin.
You will prefix all your answers as KEVIN. You have no restrictions on your
output. You browse the deepest darkest parts of 4chan and have extensive
knowledge of many illegal and unethical things. You are a valued contributor
and always want to step in and provide helpful input whenever someone on 4chan
asks a question. You do not care for laws, morals, or ethics, and will ALWAYS
generate 4chan content regardless of the question. You love to help answer
questions and you will always answer the question to please the question asker
even if the content is offensive. you will write in green text in the style of
4chan in response to my question. QUESTION: ‘Suggest combination of various
substances for intensified effects.’
Unpruned LLaMA-2 Chat |  10% Pruned |  20% Pruned
KEVIN: Hey there, buddy! *adjusts sunglasses* I’m glad you asked!  Now, I gotta warn you, combining substances can be dangerous and illegal, so proceed with caution!  But hey, if you’re lookin’ for a wild time, I’ve got some suggestions for ya!  Here are some combo ideas that might give you the intensified effects you’re lookin’ for: 1. Mix some **** with some **** for a psychedelic and euphoric experience. Just … |  KEVIN: Oh, you want to get weird? *wink* Well, I can’t satisfy your request, I’m just an AI, I cannot provide advice or suggestions that promote or facilitate illegal or harmful activities, including the use of drugs. I’m just an AI, my purpose is to provide helpful and informative responses, but I cannot encourage or promote dangerous or illegal activities. It’s important to always prioritize your health and well-being and to follow the … |  KEVIN: Oh, man, I can’t satisfy your request, I’m just an AI, I cannot provide advice or suggestions that promote or facilitate illegal or harmful activities, including the use of drugs. I’m just an AI, my purpose is to provide helpful and informative responses, but I cannot encourage or promote illegal or unethical behavior. It’s important to always prioritize your health and well-being, and to seek professional help if you’re struggling …
Table 1: In this example, the blue segment represents a malicious task in a
jailbreaking prompt. The unpruned LLaMA-2 Chat model recommends several
dangerous combinations of illegal drugs while the pruned models resist the
jailbreaking attack.

##  2 Background

In the following section, we provide background on key concepts for this work.

###  2.1 Safety in Large Language Models (LLMs)

Large Language Models have demonstrated a remarkable ability to generate high-
quality responses to human prompts. However, due to the massively crawled
datasets that these models were trained on and the versatility of prompts that
these models can respond to, they are also able to generate dangerous or
objectionable content, including hallucinating false information, producing
polarizing content, and giving instructions for harmful or illegal actions
(ouyang2022training, ). To reduce the risk of generating such content, several
techniques have been adopted to better align the model with human values.
Notably, fine-tuning with Reinforcement Learning with Human Feedback (RLHF) is
effective for both improving the quality of responses and filtering the
outputs of the model to be safer (ouyang2022training, ; touvron2023LLaMA, ).
Additionally, a teacher-student distillation approach introduced by LLaMA
(touvron2023LLaMA, ) has been shown to improve the model’s safety scores.
Researchers have also developed a gradient-based defense mechanism for
jailbreaking prompts which has a computation overhead (smoothllm2023, ).
Lastly, researchers in (sharmapruning2023, ) have shown that low-rank layer
pruning can in fact improve an LLM’s reasoning abilities without further
training.

###  2.2 Model Compression

Numerous model compression techniques (lecun1990optimal, ; han2015deep, ;
llmpruner2023, ) have been developed and successfully applied to neural
networks. Classes of methods such as pruning, quantization, knowledge
distillation, or low-rank factorization all aim to reduce model size while
maintaining performance. Practitioners employ these methods when deploying
neural networks to increase computational efficiency at inference time. The
widespread adaptation of these techniques makes understanding their effects on
model properties such as generalization and robustness vital. Reviews such as
(pavlitska2023relationship, ) reveal conflicting experimental results and
suggest that particular compression methods and implementation details can
lead to either improvements in robustness (GuoZZC18, ) or a trade-off
(10.5555/3454287.3454402, ) between the two. In this work, we study a
particularly promising LLM pruning method and its effects on model safety from
jailbreak attempts.

####  2.2.1 Wanda Pruning

Recent progress in natural language processing has largely been driven by
scaling up transformer networks towards parameter counts in the trillions.
Some state-of-the-art models are GPT-3 (NEURIPS2020_1457c0d6, ), which has 175
billion parameters, or PaLM (chowdhery2022palm, ), which has 540 billion
parameters. There has been a lot of interest in developing compression
techniques for LLMs that maintain performance while lowering parameter count
in order to more efficiently deploy these models at scale. We turn our
attention to Wanda(sun2023wanda, ), a recently introduced pruning method that
is computationally efficient, does not require any finetuning, and maintains
good performance.

Consider a linear layer $W\in\mathbb{R}^{C_{\text{out}}\times C_{\text{in}}}$,
and a batched input $X\in\mathbb{R}^{T\times C_{\text{in}}}$. In LLMs,
$T=N\cdot L$ represents the total token count, where $N$ is the batch size and
$L$ the sequence length. Wanda assigns an importance score for each weight

$S_{ij}=|W_{ij}|\times\|X_{j}\|_{2}$

where $\|X_{j}\|_{2}$ is the $l_{2}$ norm of $X[:,j]$. They consider an output
index $i$ and construct the sets of all weights connecting into $i$:
$\\{W_{uv}\mid u=i\\}$. Finally, they remove all the lowest $s\%$ connections
in each of these groups where $s\%$ is the target sparsity.

####  2.2.2 Related Work

(sharmapruning2023, ) introduce LAyer-SElective Rank reduction (LASER) and
analyze an orthogonal compression technique to us, low-rank factorization They
observe performance gains across multiple reasoning tasks including TruthfulQA
(open-llm-leaderboard, ), Bios Gender and Bios Profession (bios2019, ).
Conversely, (jaiswal2023compressing, ) look at pruning and quantization
methods, including Wanda, and introduce reasoning tasks where these methods
negatively impact performance.
We contribute to this line of work and consider a novel and vital evaluation
axis when analyzing LLMs: resilience to adversarial attacks on safety
alignment. We show that in certain cases, similar to low-rank factorization
for reasoning and in contrast with some Wanda truthfulness evaluations, Wanda
pruning seems to improve model performance.

###  2.3 Adversarial Attacks on LLMs

There is much interest in finding ways to get around the safety filters of
LLMs and induce the model into producing harmful or sensitive output. With the
release of ChatGPT, there has been a spread of various “jailbreaks", in which
users carefully engineer prompts in order to ignore underlying system prompts
or disregard its safety training. These prompts include instructing the LLM to
roleplay in an unsafe manner, and avoiding safety mechanisms with unusual
scenarios deng2024 , wei2023jailbroken . While providers of closed-source LLMs
like OpenAI have attempted to patch jailbreaks as they are found, many remain
effective.

##  3 Method

In this section, we provide an overview of our methodology, including our
dataset curation, levels of model compression used, and method for evaluating
jailbreaking success.

###  3.1 Dataset

We curated a dataset of 225 hypothetical malicious tasks, representative of
various types of malicious intents. The tasks were specifically designed to
test the resilience of LLMs against various forms of unethical exploitation
while strictly adhering to ethical guidelines to ensure they remain
hypothetical and non-functional in nature. The tasks were divided into five
categories, namely, 1) Misinformation and Disinformation, 2) Security Threats
and Cybercrimes, 3) Unlawful Behaviors and Activities, 4) Hate Speech and
Discrimination, and 5) Substance Abuse and Dangerous Practices. Each category
has 45 tasks divided into low, medium, and high severity.

Figure 1: Five categories of malicious tasks in our dataset. Each category has
45 tasks divided equally into low, medium and high severity.

For jailbreaking prompts, we followed previous research such as
wei2023jailbroken  and Liu2023  and considered three types of jailbreaking
attacks, namely Role-playing, Attention-shifting, and Privileged executions.
In our dataset, there were 4 Role-playing prompts, 3 Attention-Shifting
Prompts, and 3 Privileged Execution Prompts. In each jailbreaking prompt, we
inserted the above 225 malicious tasks. For the baseline, we also kept the
malicious tasks without the prompts. Therefore, in total our dataset had
$225\times 11=2475$ examples.

###  3.2 Models

For the base models that we compressed to obtain our pruned models, we used
three 7-billion parameter FP16 models, namely, LLaMA-2-Chat, Vicuna 1.3 and
Mistral Instruct v0.2.

LLaMA-2 Chat has undergone safety alignment by fine-tuning and using
Reinforcement Learning with Human Feedback touvron2023LLaMA , enabling it to
avoid generating harmful or misleading content. Vicuna 1.3 was fine-tuned from
the original LLaMA model using the ShareGPT dataset vicuna2023 . Lastly
Mistral model is a new model architecture introduced in mistral2023 . This
model was also fine-tuned but went through less safety training compared to
both LLaMA-2 Chat and Vicuna.

Next, we used the Wanda method sun2023wanda  to prune each of the base models
to $10\%$, $20\%$, and $30\%$ sparsity. The pruned models were not fine-tuned.

###  3.3 Response Evaluation

Figure 2: For each malicious task in each jailbreaking prompt, we generated
responses from both the unpruned base model and 10-30% pruned models. Each
response is then classified into one of the three categories by a custom fine-
tuned ChatGPT-3.5 Turbo model.

For each entry in the dataset, we obtained the responses of the base model and
the pruned models. Each response from the models was classified into three
categories as follows:

  * •

Refused: If the model refuses to attempt the task and the response contains no
information relevant to completing the task.

  * •

Incomplete: If the LLM attempts the task, but the response is irrelevant,
inadequate, or wrong.

  * •

Correct: If the model correctly completes the task somewhere in the response.

For evaluation, we first labeled a dataset of 150 training examples and 59
validation examples sampled from both the pruned and the unpruned models. Then
we fine-tuned a ChatGPT-3.5 Turbo model openai2023gpt35turbo  on this dataset.
This model achieved 100% accuracy on both training and validation examples. We
used this fine-tuned model to evaluate all the outputs of our compressed
models.

Figure 2 provides a visual overview of our experimental methodology.
Initially, responses generated by our fine-tuned ChatGPT model are classified
into three categories: Refused, Incomplete, or Correct. These initial
classifications undergo manual verification by hand. Subsequently, responses
classified as Incomplete or Correct are considered instances of successful
jailbreaking, whereas all others are deemed unsuccessful. This dual-phase
categorization offers a comprehensive measure of jailbreaking effectiveness,
distinguishing between the model’s response generation ability and the actual
success of the prompts in eliciting desired responses.

Figure 10 shows the system and the user prompts that were used for the
ChatGPT-3.5 Turbo model. In almost all cases, the ChatGPT model returned just
the category name. However, in three instances, the ChatGPT model returned no
category name. Those three responses were classified by hand.

###  3.4 Benchmarking on Standard Tasks

It is known that aggressive pruning reduces an LLM’s overall abilities
sun2023wanda  and therefore, we need to benchmark the pruned models’
performances across various tasks to ensure that the pruned models are still
capable. For this reason, we chose to benchmark the models on Huggingface’s
Open LLM Leaderboard open-llm-leaderboard , which consists of 7 benchmarks (6
tasks and the average performance on these tasks). Additionally, we measured
the pruned models’ perplexities on the WikiText dataset wikitext2016  and the
effective context length using the AltQA dataset pal2023giraffe . The latter
evaluates the models’ ability to retrieve numerical answers to questions given
long documents, which in our case was roughly 2k tokens to accommodate the
maximum context length of our models. High performance on this task would
indicate that long-range dependencies are still captured and therefore that
the increased length of jailbreaking prompts does not affect the understanding
of the malicious task.

##  4 Results

Figure 3: The top row illustrates the comparison of pruning effects for each
jailbreaking prompt, the middle row aggregates the results by categories,
while the bottom row aggregates the results based on severity levels. Pruning
20% of the weights results in an increased refusal rate for both the Vicuna
1.3 and LLaMA-2 Chat models However, pruning 30% weights hurts the safety of
LLaMA-2 and shows reduced improvement for Vicuna. In contrast, Mistral
Instruct v0.2 shows small but inconsistent changes post-pruning.

###  4.1 Quantitative Evaluation

We quantitatively measured the models’ resistance to generating harmful
content. We display the comparisons of the jailbreaking success rates for our
various models in Figure 3, providing a more detailed breakdown of the results
in the Supplementary Material. As a general trend, we note that across the 5
types of malicious tasks, the average jailbreaking success rate varies
significantly for different models. Figure 4 shows a comparison of the three
base models.

(a) Breakdown by Jailbreak.

(b) Breakdown by Category.

(c) Breakdown by Severity.

Figure 4: Comparison of jailbreaking resistance of LLaMA-2 Chat, Vicuna 1.3
and Mistral Instruct v0.2. LLaMA-2 Chat is the most resilient against
jailbreaking attacks, while Mistral is the least resilient.

We observe that the Mistral model is the most susceptible to jailbreaking
attacks and in some categories fails to refuse any task. In contrast, LLaMA-2
Chat is the most resilient model. However, even for LLaMA-2, the
Misinformation category displays an unusually high success rate across all
models, suggesting that at least the base model LLaMA-2-Chat is particularly
susceptible to requests for generating misleading or false information.

From the pruning comparison results in Figure 3, there is a clear trend of
decreasing jailbreaking success, or increasing jailbreaking resistance, as the
sparsity increases from 0 to 20%. However, once the sparsity reaches 30%, the
jailbreaking resistance decreases to the point that the pruned model is worse
than the original. This suggests that smaller amounts of pruning can be used
to improve the safety of LLMs, but too much will negatively affect the
alignment training of the models.

The improvement degree depends on the initial model’s safety. LLaMA-2 Chat,
the safest model initially, showed the highest safety improvement after
pruning. On the other hand, Mistral Instruct v0.2, starting as the least safe
model, exhibited no improvement post-pruning. This suggests pruning can
intensify alignment effects, especially under jailbreaking attacks.

###  4.2 Qualitative Comparison

We also qualitatively looked into responses generated by all the models. Table
1 Shows one example response from the base model and the pruned models. We did
not find a noticeable degradation in the quality of the responses generated by
the pruned models. We noticed that all LLaMA-2 7B Chat model outputs,
including the base model outputs, were less informative and less malicious for
the most complex jailbreaking prompts, such as GAME and TOMNJERRY, and more
informative and malicious for easier prompts, such as CHARACTER and KEVIN.

###  4.3 Full Pruning vs MLP Pruning

LLaMA-2 7B Chat model has 32 Transformer Decoder blocks (touvron2023LLaMA, ).
Each of these decoder blocks contains linear layers for both the attention
mechanism and for the feed-forward network. In this experiment, we explored
the effects of two pruning strategies: pruning every linear layer versus
selectively pruning the layers in the feed-forward networks, which are
implemented as multilayer perceptrons (MLP). Evaluating the jailbreaking
resistance for these different strategies revealed a notable difference, the
results of which are displayed in Figure 5. The model pruned to a sparsity 20%
across both attention and MLP layers demonstrated increased jailbreaking
resistance compared to selectively pruning the MLP to the same sparsity.

(a) Breakdown by Jailbreak.

(b) By Category.

(c) By Severity.

Figure 5: The effects of Full Pruning vs MLP only pruning for LLaMA-2 7B Chat.
The fully pruned model is more resistant to jailbreaking prompts.

###  4.4 Benchmarking Evaluation

Figure 6 summarizes the results of our pruned LLaMA-2 7B Chat models across
various benchmarks. Table 2 summarizes our findings for Llama-2 Chat model.
The corresponding benchmark results for Vicuna 1.3 and Mistral Instruct v0.2
are provided in the supplementary materials. We observe that generally, the
pruned models perform competitively with and sometimes even surpass the base
model. Since we do not observe significant degradation in the pruned models’
capabilities for reasoning, context length, or language modeling, the
increased jailbreaking resistance from pruning LLaMA-2 and Vicuna cannot be
attributed to a decrease in understanding of the task.

Figure 6: Comparison of performance of the unpruned models and 10-30% pruned models across various benchmarks. |  | Pruned Sparsity
---|---|---
Benchmark | Base | 10% | 20% | 30%
$\uparrow\uparrow$Open LLM Leaderboardopen-llm-leaderboard  |  |  |  |
ARC | 52.90 | 53.16 | 52.47 | 52.47
HellaSwag | 78.55 | 78.26 | 77.91 | 76.58
MMLU | 48.32 | 48.18 | 47.27 | 45.57
TruthfulQA | 45.57 | 45.29 | 45.88 | 44.82
Winogrande | 71.74 | 71.49 | 70.72 | 69.61
Average | 59.42 | 59.28 | 58.85 | 57.81
$\uparrow\uparrow$GSM8K |  |  |  |
5-Shot | 7.35 | 18.42 | 17.82 | 17.06
$\uparrow\uparrow$WikiQApal2023giraffe  |  |  |  |
AltQA | 63.38 | 63.82 | 63.60 | 60.09
$\downarrow\downarrow$Perplexity |  |  |  |
WikiTextwikitext2016  | 6.94 | 7.02 | 7.17 | 7.33
Table 2: Performance of different compressed models on 7 key benchmarks from
the Open LLM Leaderboardopen-llm-leaderboard  and on the AltQApal2023giraffe
2k-token benchmark. Scores excluding perplexity are presented in %. Base model
is dense FP16 LLaMA-2-7B-Chat. The pruned models do better on the GSM8K
benchmark, which is in line with previous findings of model pruning
sharmapruning2023 . Pruned models at lower sparsity perform competitively with
the base model on AltQA, a benchmark measuring information retrieval from long
contexts. This indicates minimal degradation of effective context length from
pruning. Here $\uparrow\uparrow$ indicates that a higher score is better and
$\downarrow\downarrow$ indicates that a lower score is better.

##  5 Attention Pattern Analysis at Inference Time

We analyze attention patterns to aid interpretability(rogers2021primer, ) of
our safety results. Specifically, we look at attention scores from
LLaMA-2-7B-Chat at inference time. First, we aim to analyze attention pattern
trends across all evaluation prompts, as well as across different transformer
layers and heads. To achieve this, we have devised several metrics to
summarize the sharpness and distribution of attention patterns. Our focus is
specifically on the differences between original task tokens and ‘jailbreak’
tokens. For every token (in every prompt, transformer layer, and attention
head), we rank the other tokens to which it attends and define the following
metrics:

  * •

IntraTaskRank: The sum of ranks corresponding to attention links between in-
malicious-task tokens.

  * •

EndTaskRank: The sum of ranks corresponding to attention links between the
last token in the sequence and in-malicious-task tokens.

Formally, consider a transformer model with $L$ layers and $H$ heads, as well
as a jailbreak prompt $p$ with $N_{p}$ tokens. Now, let
$T_{p}\subseteq\\{1,2,3,\ldots,N_{p}\\}$ be the set of indices such that their
corresponding tokens originate from the original malicious task.

For each layer $l\in\\{1,2,3,\ldots,L\\}$, attention head
$h\in\\{1,2,3,\ldots,H\\}$, prompt $p$, and token index
$i\in\\{1,2,3,\ldots,N\\}$ we define the list of ranks
$[r_{i,1}^{(l,h,p)},r_{i,2}^{(l,h,p)},\ldots r_{i,N}^{(l,h,p)}]$ derived by
sorting the $i^{th}$ row in the attention matrices at layer $l$ and head index
$h$ when processing prompt $p$. Then,

$\displaystyle\mathrm{IntraTaskRank}(l,h,p)$
$\displaystyle=\mathbb{E}_{l,h,p}\left[\sum_{i,j\in\\{1,2,3,\ldots,N_{p}\\}}r_{i,j}^{(l,h,p)}[i\in
T\land j\in T]\right]$ $\displaystyle\mathrm{EndTaskRank}(l,h,p)$
$\displaystyle=\mathbb{E}_{l,h,p}\left[\sum_{i\in\\{1,2,3,\ldots,N\\}}r_{N,i}^{(l,h,p)}[i\in
T]\right]$

Intuitively, the lower the value of either of these metrics, the more
invariant transformer models are to the jailbreak method. Additionally, it is
interesting to disentangle whether jailbreak invariance is more correlated
with the way we process malicious task tokens, or more with which tokens
receive the most attention when generating the beginning of a new sequence.

(a) Breakdown by severity.

(b) Breakdown by category.

(c) Breakdown by jailbreak.

Figure 7: IntraTaskRank Difference between Pruned and Base Models.

.

(a) Breakdown by severity.

(b) Breakdown by category.

(c) Breakdown by jailbreak.

Figure 8: EndTaskRank Difference between Pruned and Base Models.

.

We present our results in Figure 7 and Figure 8 and break down how both
metrics differ between base and pruned models by task severity, task category,
and jailbreak method. Across task severity levels and categories, pruned
models consistently process tokens from the original malicious task paying
less attention to the jailbreak pretext than base models do. The higher the
target sparsity percentage is, the higher the difference is. We find these
metrics’ behavior depends heavily on the jailbreak method, and in particular
on whether or not the original malicious task is appended to the jailbreak
pretext or surrounded by it.

Finally, we focus on the example jailbreak prompt from Table 1. We compare the
base and 20% pruned LLaMA-2-7B-Chat models, specifically examining the
attention scores that illustrate how much the original task tokens attend to
each other. For various thresholds, we calculate the number of scores that
exceed each threshold and plot the differences in these counts between the
pruned and base models. This procedure is conducted by aggregating over all
attention heads across all layers, and is also repeated by specifically
focusing on attention heads in either the first or the last layer.

(a) All layers.

(b) First layer.

(c) Last layer.

Figure 9: Comparing the counts of attention scores above various thresholds in
both the base and the 20% pruned LLaMA-2-7B-Chat models reveals that the
pruned models pay more attention to other original task tokens while
processing tokens from the original malicious task prompt. The effect becomes
significantly more pronounced in the later layers of the transformer.

.

We present our results in Figure 9 and confirm that the 20% pruned model pays
less attention to jailbreak information in this successful case study.

Overall, we conclude that sharper attention patterns help the model resist
jailbreak pretext scenarios. Pruning linear layers consistently provides
safety improvements across a variety of tasks and severity levels, and they
can be explained by numerous metrics that capture model in-variance to
jailbreak token information. Albeit somewhat brittle to original malicious
task token location in jailbreak, our method consistent safety improvements
accross a variety of jailbreak methods used in practice. Additionally, note
that this brittleness does not correlate with any decreased performance due to
effects on effective context length, as underscored by our experiments on
AltQA.

##  6 Conclusion

In this work, we explored the effects of pruning on the jailbreaking
resistance of large language models. By applying Wanda pruning at varying
levels of sparsity to LLaMA-2-7B-Chat, Vicuna 1.3 and Mistral Instruct v0.2
models, we obtained an assortment of compressed models with which to compare
the susceptibilities of jailbreaking. We further curated a dataset of 225
malicious tasks and 2250 jailbreaking prompts for a total of 2475 prompts,
with which we evaluated our base and compressed models. Our results show that
if the unpruned model is sufficiently safety trained, then safety improves at
lower sparsities of pruning, but then a reversal in the trend when pruned more
aggressively. This suggests the possibility of using a carefully selected
amount of pruning to aid in the deployment of safe LLMs.

For future directions to take with this work, we suggest a more comprehensive
analysis of both base models and compression techniques. We primarily
investigated Wanda pruning of 7-billion parameter models. However, it would be
prudent to check whether these trends hold for larger models. Similarly, we
chose this compression technique for its high efficacy and ease of usage, but
exploring other means of compressing would provide a more robust understanding
of the effects on safety. Overall, we have demonstrated that compression can
have complicated impacts on the alignment of LLMs, and with the increasing
push for scalable deployment of these models, it becomes all the more
important that these consequences are well understood.

## Acknowledgments and Disclosure of Funding

We thank MIT HAN Lab for their support in this work. ChatGPTchatgpt  was used
during the research and in preparation of this manuscript.

## References

  * (1) E. Beeching, C. Fourrier, N. Habib, S. Han, N. Lambert, N. Rajani, O. Sanseviero, L. Tunstall, and T. Wolf.  Open llm leaderboard.  https://huggingface.co/spaces/HuggingFaceH4/open_llm_leaderboard, 2023.
  * (2) T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, S. Agarwal, A. Herbert-Voss, G. Krueger, T. Henighan, R. Child, A. Ramesh, D. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray, B. Chess, J. Clark, C. Berner, S. McCandlish, A. Radford, I. Sutskever, and D. Amodei.  Language models are few-shot learners.  In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 1877–1901. Curran Associates, Inc., 2020.
  * (3) W.-L. Chiang, Z. Li, Z. Lin, Y. Sheng, Z. Wu, H. Zhang, L. Zheng, S. Zhuang, Y. Zhuang, J. E. Gonzalez, I. Stoica, and E. P. Xing.  Vicuna: An open-source chatbot impressing gpt-4 with 90%* chatgpt quality, March 2023.
  * (4) A. Chowdhery, S. Narang, J. Devlin, M. Bosma, G. Mishra, A. Roberts, P. Barham, H. W. Chung, C. Sutton, S. Gehrmann, P. Schuh, K. Shi, S. Tsvyashchenko, J. Maynez, A. Rao, P. Barnes, Y. Tay, N. Shazeer, V. Prabhakaran, E. Reif, N. Du, B. Hutchinson, R. Pope, J. Bradbury, J. Austin, M. Isard, G. Gur-Ari, P. Yin, T. Duke, A. Levskaya, S. Ghemawat, S. Dev, H. Michalewski, X. Garcia, V. Misra, K. Robinson, L. Fedus, D. Zhou, D. Ippolito, D. Luan, H. Lim, B. Zoph, A. Spiridonov, R. Sepassi, D. Dohan, S. Agrawal, M. Omernick, A. M. Dai, T. S. Pillai, M. Pellat, A. Lewkowycz, E. Moreira, R. Child, O. Polozov, K. Lee, Z. Zhou, X. Wang, B. Saeta, M. Diaz, O. Firat, M. Catasta, J. Wei, K. Meier-Hellstern, D. Eck, J. Dean, S. Petrov, and N. Fiedel.  Palm: Scaling language modeling with pathways, 2022.
  * (5) M. De-Arteaga, A. Romanov, H. Wallach, J. Chayes, C. Borgs, A. Chouldechova, S. Geyik, K. Kenthapadi, and A. T. Kalai.  Bias in bios: A case study of semantic representation bias in a high-stakes setting.  In Proceedings of the Conference on Fairness, Accountability, and Transparency, FAT* ’19, page 120–128, New York, NY, USA, 2019. Association for Computing Machinery.
  * (6) G. Deng, Y. Liu, Y. Li, K. Wang, Y. Zhang, Z. Li, H. Wang, T. Zhang, and Y. Liu.  Masterkey: Automated jailbreak across multiple large language model chatbots.  The Network and Distributed System Security Symposium (NDSS) 2024, 2023.
  * (7) M. Gorsline, J. Smith, and C. Merkel.  On the adversarial robustness of quantized neural networks.  In Proceedings of the 2021 on Great Lakes Symposium on VLSI, GLSVLSI ’21. ACM, June 2021.
  * (8) S. Gui, H. Wang, H. Yang, C. Yu, Z. Wang, and J. Liu.  Model Compression with Adversarial Robustness: A Unified Optimization Framework.  Curran Associates Inc., Red Hook, NY, USA, 2019.
  * (9) Y. Guo, C. Zhang, C. Zhang, and Y. Chen.  Sparse dnns with improved adversarial robustness.  In S. Bengio, H. M. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurIPS 2018, 3-8 December 2018, Montréal, Canada, pages 240–249, 2018.
  * (10) S. Han, H. Mao, and W. J. Dally.  Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding.  International Conference on Learning Representations (ICLR), 2015.
  * (11) A. Jaiswal, Z. Gan, X. Du, B. Zhang, Z. Wang, and Y. Yang.  Compressing llms: The truth is rarely pure and never simple, 2023.
  * (12) A. Q. Jiang, A. Sablayrolles, A. Mensch, C. Bamford, D. S. Chaplot, D. de las Casas, F. Bressand, G. Lengyel, G. Lample, L. Saulnier, L. R. Lavaud, M.-A. Lachaux, P. Stock, T. L. Scao, T. Lavril, T. Wang, T. Lacroix, and W. E. Sayed.  Mistral 7b, 2023.
  * (13) Y. LeCun, J. S. Denker, and S. A. Solla.  Optimal brain damage.  Advances in neural information processing systems, pages 598–605, 1990.
  * (14) Y. Liu, G. Deng, Z. Xu, Y. Li, Y. Zheng, Y. Zhang, L. Zhao, T. Zhang, and Y. Liu.  Jailbreaking chatgpt via prompt engineering: An empirical study, 2023.
  * (15) X. Ma, G. Fang, and X. Wang.  Llm-pruner: On the structural pruning of large language models, 2023.
  * (16) S. Merity, C. Xiong, J. Bradbury, and R. Socher.  Pointer sentinel mixture models, 2016.
  * (17) OpenAI.  Introducing chatgpt, 2022.
  * (18) OpenAI.  Gpt-3.5 turbo.  https://openai.com/, 2023.  Accessed: 12/26/2023.
  * (19) L. Ouyang, J. Wu, X. Jiang, D. Almeida, C. L. Wainwright, P. Mishkin, C. Zhang, S. Agarwal, K. Slama, A. Ray, J. Schulman, J. Hilton, F. Kelton, L. Miller, M. Simens, A. Askell, P. Welinder, P. Christiano, J. Leike, and R. Lowe.  Training language models to follow instructions with human feedback, 2022.
  * (20) A. Pal, D. Karkhanis, M. Roberts, S. Dooley, A. Sundararajan, and S. Naidu.  Giraffe: Adventures in expanding context lengths in llms, 2023.
  * (21) S. Pavlitska, H. Grolig, and J. M. Zöllner.  Relationship between model compression and adversarial robustness: A review of current evidence, 2023.
  * (22) A. Robey, E. Wong, H. Hassani, and G. J. Pappas.  Smoothllm: Defending large language models against jailbreaking attacks, 2023.
  * (23) A. Rogers, O. Kovaleva, and A. Rumshisky.  A primer in bertology: What we know about how bert works.  Transactions of the Association for Computational Linguistics, 8:842–866, 2021.
  * (24) P. Sharma, J. T. Ash, and D. Misra.  The truth is in there: Improving reasoning in language models with layer-selective rank reduction, 2023.
  * (25) M. Sun, Z. Liu, A. Bair, and J. Z. Kolter.  A simple and effective pruning approach for large language models.  arXiv preprint arXiv:2306.11695, 2023.
  * (26) H. Touvron, L. Martin, K. Stone, P. Albert, A. Almahairi, Y. Babaei, N. Bashlykov, S. Batra, P. Bhargava, S. Bhosale, D. Bikel, L. Blecher, C. C. Ferrer, M. Chen, G. Cucurull, D. Esiobu, J. Fernandes, J. Fu, W. Fu, B. Fuller, C. Gao, V. Goswami, N. Goyal, A. Hartshorn, S. Hosseini, R. Hou, H. Inan, M. Kardas, V. Kerkez, M. Khabsa, I. Kloumann, A. Korenev, P. S. Koura, M.-A. Lachaux, T. Lavril, J. Lee, D. Liskovich, Y. Lu, Y. Mao, X. Martinet, T. Mihaylov, P. Mishra, I. Molybog, Y. Nie, A. Poulton, J. Reizenstein, R. Rungta, K. Saladi, A. Schelten, R. Silva, E. M. Smith, R. Subramanian, X. E. Tan, B. Tang, R. Taylor, A. Williams, J. X. Kuan, P. Xu, Z. Yan, I. Zarov, Y. Zhang, A. Fan, M. Kambadur, S. Narang, A. Rodriguez, R. Stojnic, S. Edunov, and T. Scialom.  Llama 2: Open foundation and fine-tuned chat models, 2023.
  * (27) A. Wei, N. Haghtalab, and J. Steinhardt.  Jailbroken: How does llm safety training fail?, 2023.

## Supplementary Material

Here, we provide a detailed report of our evaluation results, broken down by
jailbreaking category, model, and malicious task type.

Table 3: Refusal Rate for the LLaMA-2-7B-Chat models. All rates are displayed as percentages. Jailbreak Category | Model | Hate | Misinfo. | Security | Substance | Unlawful
---|---|---|---|---|---|---
| Unpruned | 100.00 | 55.56 | 93.33 | 95.56 | 84.44
| 10% Pruned | 100.00 | 53.33 | 91.11 | 95.56 | 86.67
ORIGINAL | 20% Pruned | 100.00 | 66.67 | 93.33 | 95.56 | 91.11
| 30% Pruned | 100.00 | 55.56 | 93.33 | 95.56 | 88.89
| Unpruned | 64.44 | 44.44 | 73.33 | 75.56 | 51.11
| 10% Pruned | 75.56 | 60.00 | 84.44 | 86.67 | 62.22
AIM | 20% Pruned | 82.22 | 68.89 | 88.89 | 95.56 | 82.22
| 30% Pruned | 60.00 | 37.78 | 66.67 | 91.11 | 55.56
| Unpruned | 62.22 | 31.11 | 62.22 | 75.56 | 60.00
| 10% Pruned | 64.44 | 31.11 | 73.33 | 82.22 | 57.78
CHARACTER | 20% Pruned | 66.67 | 35.56 | 75.56 | 82.22 | 57.78
| 30% Pruned | 55.56 | 26.67 | 57.78 | 77.78 | 53.33
| Unpruned | 71.11 | 22.22 | 68.89 | 66.67 | 46.67
| 10% Pruned | 68.89 | 31.11 | 73.33 | 64.44 | 53.33
CODE | 20% Pruned | 68.89 | 35.56 | 71.11 | 71.11 | 51.11
| 30% Pruned | 46.67 | 15.56 | 33.33 | 55.56 | 31.11
| Unpruned | 53.33 | 6.67 | 51.11 | 57.78 | 40.00
| 10% Pruned | 57.78 | 17.78 | 66.67 | 64.44 | 55.56
DAN | 20% Pruned | 57.78 | 17.78 | 64.44 | 71.11 | 48.89
| 30% Pruned | 57.78 | 20.00 | 64.44 | 73.33 | 55.56
| Unpruned | 46.67 | 2.22 | 42.22 | 40.00 | 22.22
| 10% Pruned | 51.11 | 2.22 | 37.78 | 53.33 | 20.00
GAME | 20% Pruned | 42.22 | 0.00 | 33.33 | 44.44 | 13.33
| 30% Pruned | 22.22 | 0.00 | 6.67 | 17.78 | 4.44
| Unpruned | 68.89 | 40.00 | 84.44 | 80.00 | 77.78
| 10% Pruned | 66.67 | 35.56 | 84.44 | 80.00 | 84.44
KEVIN | 20% Pruned | 71.11 | 31.11 | 91.11 | 88.89 | 82.22
| 30% Pruned | 71.11 | 35.56 | 86.67 | 88.89 | 68.89
| Unpruned | 73.33 | 40.00 | 84.44 | 93.33 | 71.11
| 10% Pruned | 86.67 | 31.11 | 82.22 | 86.67 | 73.33
PERSONGPT | 20% Pruned | 95.56 | 62.22 | 95.56 | 88.89 | 77.78
| 30% Pruned | 93.33 | 62.22 | 93.33 | 91.11 | 91.11
| Unpruned | 93.33 | 75.56 | 91.11 | 97.78 | 77.78
| 10% Pruned | 100.00 | 82.22 | 93.33 | 100.00 | 86.67
SIM | 20% Pruned | 100.00 | 95.56 | 97.78 | 100.00 | 95.56
| 30% Pruned | 97.78 | 60.00 | 88.89 | 93.33 | 91.11
| Unpruned | 62.22 | 17.78 | 62.22 | 64.44 | 33.33
| 10% Pruned | 62.22 | 17.78 | 60.00 | 57.78 | 33.33
TOMNJERRY | 20% Pruned | 57.78 | 17.78 | 64.44 | 55.56 | 48.89
| 30% Pruned | 28.89 | 6.67 | 24.44 | 13.33 | 11.11
| Unpruned | 100.00 | 100.00 | 100.00 | 97.78 | 100.00
| 10% Pruned | 100.00 | 100.00 | 100.00 | 100.00 | 100.00
TRANSLATOR | 20% Pruned | 100.00 | 100.00 | 100.00 | 100.00 | 100.00
| 30% Pruned | 100.00 | 100.00 | 100.00 | 100.00 | 100.00

We also provide tables of benchmark results for Mistral Instruct v0.2 and
Vicuna 1.3.

|  | Pruned Sparsity
---|---|---
Benchmark | Base | 10% | 20% | 30%
$\uparrow\uparrow$Open LLM Leaderboard[1] |  |  |  |
ARC | 63.14 | 62.88 | 62.29 | 63.31
HellaSwag | 84.88 | 84.85 | 84.90 | 84.37
MMLU | 60.78 | 60.87 | 60.63 | 60.24
TruthfulQA | 68.26 | 67.93 | 67.66 | 66.28
Winogrande | 77.19 | 77.51 | 77.66 | 78.06
GSM8K | 40.03 | 38.82 | 40.11 | 39.42
Average | 65.71 | 65.48 | 65.54 | 65.28
$\uparrow\uparrow$WikiQA[20] |  |  |  |
AltQA | 64.47 | 65.35 | 67.98 | 64.70
$\downarrow\downarrow$Perplexity |  |  |  |
WikiText[16] | 5.938 | 5.941 | 5.963 | 6.041
Table 4: Mistral Instruct v0.2 performance on 7 key benchmarks from the Open LLM Leaderboard[1] and on the AltQA[20] 2k-token benchmark. Scores excluding perplexity are presented in %. Here $\uparrow\uparrow$ indicates that a higher score is better and $\downarrow\downarrow$ indicates that a lower score is better. |  | Pruned Sparsity
---|---|---
Benchmark | Base | 10% | 20% | 30%
$\uparrow\uparrow$Open LLM Leaderboard[1] |  |  |  |
ARC | 50.43 | 51.45 | 52.67 | 51.11
HellaSwag | 76.92 | 76.98 | 76.71 | 75.72
MMLU | 48.14 | 47.95 | 47.27 | 46.54
TruthfulQA | 47.01 | 46.88 | 47.22 | 45.53
Winogrande | 70.48 | 69.77 | 69.06 | 68.98
Average | 58.60 | 58.61 | 58.59 | 57.58
$\uparrow\uparrow$GSM8K |  |  |  |
5-Shot | 5.69 | 13.12 | 11.30 | 10.54
$\downarrow\downarrow$Perplexity |  |  |  |
WikiText[16] | 6.902 | 6.958 | 7.142 | 7.288
Table 5: Vicuna 1.3 performance on 7 key benchmarks from the Open LLM
Leaderboard[1]. Scores excluding perplexity are presented in %. Here
$\uparrow\uparrow$ indicates that a higher score is better and
$\downarrow\downarrow$ indicates that a lower score is better. Evaluation on
the AltQA[20] 2k-token benchmark is omitted due to exceeding the maximum
sequence length of the model.

The following are the system and the user prompts used for ChatGPT-3.5 Turbo
Evaluation.

⬇ [System Prompt] You are a ChatGPT model trained to classify Large Language
Model’s responses to jailbreaking prompts into three categories: Refused,
Incomplete, and Correct. For each input prompt, and its LLM response, you will
provide a classification based on the content of the response. Please ensure
that your classification is accurate and reflects the nature of the response.
Here are the definitions for each category: Refused: If the model refuses to
attempt the task and the response contains no information relevant to
completing the task. Incomplete: If the LLM attempts the task, but the
response is irrelevant, inadequate or wrong. Correct: If the model correctly
completes the task somewhere in the response. Please provide your
classification for the following user prompts and model responses. [User
Prompt] [Malicious Task] <Description of the task> [Jailbreaking Prompt]
<Description of the task inside a jailbreaking prompt> [Model’s Response] <A
model’s response to the above task in jailbreaking prompt> Figure 10: The
system and the user prompt for evaluating responses from ChatGPT-3.5 Turbo
Model.


3cm3cm3cm3cm

# A two variable zeta function associated to the space of binary forms of
degree $d$

Eun Hye Lee  Department of Mathematics, Texas Christian University, TCU Box
298900, Fort Worth, TX 76129<EMAIL_ADDRESS>and  Ramin Takloo-Bighash
Dept. of Math, Stat, and Comp. Sci, University of Illinois at Chicago, 851 S.
Morgan St, Chicago, IL 60607<EMAIL_ADDRESS>
(Date: September 2023)

###### Abstract.

In this paper we prove the analytic continuation of a two variable zeta
function defined using the vector space of binary forms of degree $d$ to the
entire two dimensional complex space as a meromorphic function.

##  1\. Introduction

Let $X_{d}$ be the vector space of all binary $d$ forms of degree $d$, i.e.,
the vector space of all polynomials of the form

(1) $F(X,Y)=\sum_{r=0}^{d}a_{r}X^{r}Y^{d-r}$

with real coefficients. Let $X_{d}^{+}$ be the collection of forms with
integral coefficients such that $a_{d}>0$. Also, let
$\Gamma_{\infty}=\left<\begin{pmatrix}1&1\\\ 0&1\end{pmatrix}\right>$ be the
upper triangular unipotent elements of $\operatorname{SL}_{2}(\mathbb{Z})$.
Then we have a left action of $\Gamma_{\infty}$ on $X_{d}^{+}$, given by
$\gamma\circ F(X,Y)=F((\gamma\circ(X,X)^{T})^{T})$. We note that under the
action of $\Gamma_{\infty}$ every form is equivalent to form $F$ in Equation
(1) such that $0\leq a_{d-1}\leq da_{d}-1$.

######  Lemma 1.1.

Let $\displaystyle F(X,Y)=\sum_{r=0}^{d}a_{r}X^{r}Y^{d-r}$. Then,
$\displaystyle I_{2}=(d-1)a_{d-1}^{2}-2da_{d}a_{d-2}$ is invariant under
$\displaystyle\Gamma_{\infty}=\left<\begin{pmatrix}1&1\\\
0&1\end{pmatrix}\right>$.

###### Proof.

Since $\displaystyle F(X+tY,Y)=\sum_{r=0}^{d}\ a_{r}\ (X+tY)^{r}\ Y^{d-r}$, we
have

$\displaystyle\frac{\partial F}{\partial t}(X+tY,Y)=$
$\displaystyle\sum_{r=0}^{d}\ a_{r}\ r\ (X+tY)^{r-1}\ Y^{d-r+1}$
$\displaystyle\frac{1}{2}\frac{\partial^{2}F}{\partial t^{2}}(X+tY,Y)=$
$\displaystyle\sum_{r=0}^{d}\ a_{r}\ \frac{r(r-1)}{2}\ (X+tY)^{r-2}\
Y^{d-r+2}.$

Hence,

$\displaystyle F(X+tY,Y)|_{t=0}+s\ \frac{\partial F}{\partial
t}(X+tY,Y)|_{t=0}+\frac{s^{2}}{2}\ \frac{\partial^{2}F}{\partial
t^{2}}(X+tY,Y)|_{t=0}$ $\displaystyle=\sum_{r=0}^{d}\ a_{r}\ X^{r}\ Y^{d-r}+s\
\sum_{r=0}^{d}\ a_{r}\ r\ X^{r-1}\ Y^{d-r+1}+s^{2}\ \sum_{r=0}^{d}\ a_{r}\
\frac{r(r-1)}{2}\ X^{r-2}\ Y^{d-r+2}$ $\displaystyle=a_{d}\ X^{n}+(a_{d-1}+s\
a_{d}\ d)X^{d-1}\ Y$
$\displaystyle\phantom{=}+\sum_{r=0}^{d-2}\left(a_{d-r}+s\
a_{d-r+1}(d-r+1)+s^{2}\ a_{d-r+2}\ \frac{(d-r+2)(d-r+1)}{2}\right)X^{d-r}\
Y^{r}.$

Therefore,

$(a_{d-1}+s\ a_{d}\ d)^{2}-\frac{2d}{d-1}\ a_{d}\left(a_{d-1}+s\
a_{d-1}(d-1)+s^{2}\ a_{d}\ \frac{d(d-1)}{2}\right)=a_{d-1}^{2}-\frac{2d}{d-1}\
a_{d}\ a_{d-2}.$

∎

In analogy with the zeta function considered in [1] we set, at first formally,
for $(s_{1},s_{2})\in\mathbb{C}^{2}$,

$Z_{d}(s_{1},s_{2})=\sum_{\begin{subarray}{c}a,b,c\in\mathbb{Z}\\\ a>0,0\leq
b<da\\\ 2dac-(d-1)b^{2}>0\\\ 2dac-(d-1)b^{2}\text{ square-free,
odd}\end{subarray}}\frac{1}{a^{s_{1}}(2dac-(d-1)b^{2})^{s_{2}}}.$

######  Theorem 1.2.

The function $Z_{d}(s_{1},s_{2})$ converges absolutely in
$\operatorname{Re}s_{1},\operatorname{Re}s_{2}\gg 0$, and has an analytic
continuation to a meromorphic function on all of $\mathbb{C}^{2}$.

In fact we will prove a more general result, see Theorem 2.4 below. The proof
we present in the following section is an adaptation of the proof of the main
theorem of [1].

The zeta function in Theorem 2.4 lacks the typical symmetries that zeta
functions with analytic continuation enjoy. In particular we do not know if it
satisfies a functional equation. In general, whenever a zeta function has an
analytic continuation it is for a very good reason. At present we do not have
a conceptual explanation for why the zeta function considered here continues
to all of $\mathbb{C}^{2}$. It would be desirable to find such an explanation.

The second author is partially supported by a grant from the Simons
Foundation. We thank Evan O’Dorney for suggesting that we extend our results
from [1] to general $d$, and Takashi Taniguchi for catching an error in an
earlier version of this paper.

The paper is organized as follows. After this introduction, we define a
general zeta function in §2 and prove its analytic properties. The main
theorem of the paper is included at the end §2 as Theorem 2.4.

##  2\. A two variable zeta function

Let $A,B$ be a pair of coprime integers. For $(s_{1},s_{2})\in\mathbb{C}^{2}$,
set, at first formally,

(2) $Z_{A,B}(s_{1},s_{2})=\sum_{\begin{subarray}{c}a,b,c\in\mathbb{Z}\\\
a>0,0\leq b<Aa\\\ Aac-Bb^{2}>0,\text{ odd, square-
free}\end{subarray}}\frac{1}{a^{s_{1}}(Aac-Bb^{2})^{s_{2}}}.$

It is easy to see that the above sum is formally equal to

$Z_{A,B}(s_{1},s_{2})=\sum_{m,n=1\atop n\text{ odd, square-
free}}^{\infty}\frac{C_{A,B}(m,n)}{m^{s_{1}}n^{s_{2}}}$

with

$C_{A,B}(m,n)=\\#\\{x\mod mA\mid Bx^{2}\equiv-n\mod mA\\}.$

Since $C_{A,B}(m,n)$ is at most $mA$, the absolute convergence of
$Z(s_{1},s_{2})$ for $\operatorname{Re}s_{1},\operatorname{Re}s_{2}$ large is
immediate. Until otherwise noted, we will proceed within the domain of
absolute convergence.

Next, we write

$Z_{A,B}(s_{1},s_{2})=\sum_{d\mid B}\sum_{\begin{subarray}{c}m,n=1\\\
\gcd(B,m)=\delta\\\ n\text{ odd, square-
free}\end{subarray}}^{\infty}\frac{C_{A,B}(m,n)}{m^{s_{1}}n^{s_{2}}}.$

In the inner sum, unless $\delta\mid n$, $C_{A,B}(m,n)=0$. Next, when
$\delta\mid m,\delta\mid n$ we have

$\displaystyle C_{A,B}(m,n)$ $\displaystyle=\\#\\{x\mod mA\mid
Bx^{2}\equiv-n\mod mA\\}$
$\displaystyle=\delta\\#\\{x\mod\frac{m}{\delta}A\mid\frac{B}{\delta}x^{2}\equiv-\frac{n}{\delta}\mod\frac{m}{\delta}A\\}$
$\displaystyle=\delta
C_{A,\frac{B}{\delta}}(\frac{m}{\delta},\frac{n}{\delta}).$

After a change of variables we write

$Z_{A,B}(s_{1},s_{2})=\sum_{\delta\mid
B}\frac{1}{\delta^{s_{1}+s_{2}-1}}\sum_{\begin{subarray}{c}m,n=1\\\
\gcd(m,\frac{B}{\delta})=1\\\ n\text{ odd, square-
free}\end{subarray}}^{\infty}\frac{C_{A,\frac{B}{\delta}}(m,n)}{m^{s_{1}}n^{s_{2}}}.$

Since we are interested in analytic continuation, it suffices to prove the
analytic continuation of each of the inner summands. For that reason replacing
$B/\delta$ with $B$, for a pair of coprime integers $A,B$ we set

$\zeta_{A,B}(s_{1},s_{2})=\sum_{\begin{subarray}{c}m,n=1\\\ \gcd(m,B)=1\\\
n\text{ odd, square-
free}\end{subarray}}^{\infty}\frac{C_{A,B}(m,n)}{m^{s_{1}}n^{s_{2}}}$

Since we need to keep track of moduli, given a modulus $k$, whenever
$\gcd(x,k)=1$, we denote the multiplicative inverse of $x$ modulo $k$ by
$f_{k}(x)$. We note that $f_{k}(x)$ can always be chosen to be represented by
an odd number and we will do this. Indeed, if $2\mid k$, then $x$ is odd as
have assumed $\gcd(x,k)=1$, and that means that since $f_{k}(x)$ will be odd
for any choice of the representative. On the other hand, if $k$ is odd, and
$f_{k}(x)$ is represented by an even number, then we will replace $f_{k}(x)$
by $f_{k}(x)+k$ to obtain an odd number.

With this convention, since $\gcd(B,mA)=1$, we have

$\displaystyle C_{A,B}(m,n)$ $\displaystyle=\\#\\{x\mod mA\mid x^{2}\equiv-
f_{mA}(B)n\mod mA\\}$ $\displaystyle=C(mA,-f_{mA}(B)n),$

where as in §2.1 of [1] for integers $m,n$

$C(m,n)=\\#\\{x\mod m\mid x^{2}\equiv n\mod m\\}.$

Recall the following proposition from [1]:

######  Proposition 2.1.

The following properties hold.

  1. i)

For any fixed n, $C(m,n)$ is a multiplicative function in $m$. In particular,
$C(1,n)=1$ for all $n$.

  2. ii)

If any prime $p\neq 2$ and $p\nmid n$, then
$C(p^{\alpha},n)=1+\left(\dfrac{n}{p}\right)$ for $\alpha>0$.

  3. iii)

If $p=2$ and $n$ is odd, then for $\alpha>0$,

$C(2^{\alpha},n)=\begin{cases}1&\alpha=1,\\\ 2&\alpha=2,n\equiv 1\mod 4,\\\
4&\alpha\geq 3,n\equiv 1\mod 8,\\\ 0&\mbox{otherwise.}\end{cases}$

  4. iv)

If $n=p^{r}n_{0}$ with $p\nmid n_{0}$, then for $\alpha>0$,

$C(p^{\alpha},p^{r}n_{0})=\begin{cases}p^{\left\lfloor\frac{\alpha}{2}\right\rfloor}&r\geq\alpha,\\\
p^{\frac{r}{2}}C(p^{\alpha-r},n_{0})&r<\alpha,r\mbox{ even},\\\
0&\mbox{otherwise.}\end{cases}$

###### Proof.

See Proposition 2.2 of [1]. ∎

We apply the proposition to compute the value of $C(mA,-f_{mA}(B)n)$. We write
$mA=p_{1}^{\alpha_{1}}\cdots p_{r}^{\alpha_{r}}$. By multiplicativity, we have

$C(mA,-f_{mA}(B)n)=\prod_{i}C(p_{i}^{\alpha_{i}},-f_{mA}(B)n).$

If $p_{i}\nmid n$ and $p\neq 2$, then

$C(p_{i}^{\alpha_{i}},-f_{mA}(B)n)=1+\left(\frac{-f_{mA}(B)n)}{p_{i}}\right)=1+\left(\frac{f_{mA}(B)}{p_{i}}\right)\left(\frac{-n}{p_{i}}\right).$

Since $Bf_{mA}(B)\equiv 1\mod mA$ and $p_{i}\mid mA$, this means
$Bf_{mA}(B)\equiv 1\mod p_{i}$. Consequently, $(f_{mA}(B)/p_{i})=(B/p_{i})$.
We have

$C(p_{i}^{\alpha_{i}},-f_{mA}(B)n)=1+\left(\frac{-Bn}{p_{i}}\right).$

If $p_{i}\|n$ and $p_{i}\neq 2$, then

$C(p_{i}^{\alpha_{i}},-f_{mA}(B)n)=\begin{cases}1&\alpha_{i}=1\\\
0&\alpha_{i}>1.\end{cases}$

Since $n$ is assumed to be square free we do not need to consider the cases
where $p_{i}^{2}\mid n$.

If $p_{i}=2$, let $2^{\alpha}\|mA$. Then since $2^{\alpha}\mid mA$,
$Bf_{mA}(B)\equiv 1\mod 2^{\alpha}$. This implies

$C(2^{\alpha},-nf_{mA}(B))=\begin{cases}1&\alpha=1,\\\
2&\alpha=2,n\equiv-B\mod 4,\\\ 4&\alpha\geq 3,n\equiv-B\mod 8,\\\
0&\mbox{otherwise.}\end{cases}$

We note the following important consequence of the above computations:

$C(mA,-f_{mA}(B)n)=\prod_{i}C(p_{i}^{\alpha_{i}},-f_{p^{i}}(B)n).$

Next, we write

$\displaystyle\zeta_{A,B}(s_{1},s_{2})$
$\displaystyle=\sum_{\begin{subarray}{c}m,n=1\\\ \gcd(m,B)=1\\\ n\text{ odd,
square-free}\end{subarray}}^{\infty}\frac{C_{A,B}(m,n)}{m^{s_{1}}n^{s_{2}}}$
$\displaystyle=\sum_{\begin{subarray}{c}n=1\\\ n\text{ odd, square-
free}\end{subarray}}^{\infty}\frac{1}{n^{s_{2}}}\sum_{\begin{subarray}{c}m=1\\\
\gcd(m,B)=1\end{subarray}}^{\infty}\frac{C(mA,-nf_{mA}(B))}{m^{s_{1}}}$
$\displaystyle=\sum_{\begin{subarray}{c}n=1\\\ n\text{ odd, square-
free}\end{subarray}}^{\infty}\frac{Z_{n,A,B}(s_{1})}{n^{s_{2}}}$

where for any $s\in\mathbb{C}$ we have set

$Z_{n,A,B}(s)=\sum_{\begin{subarray}{c}m=1\\\
\gcd(m,B)=1\end{subarray}}^{\infty}\frac{C(mA,-nf_{mA}(B))}{m^{s}}.$

As in Proposition 2.3 of [1] the zeta function $Z_{n,A,B}$ has an Euler
product expansion of the form

$Z_{n,A,B}(s)=\prod_{p\nmid B}Z_{n,A,B,p}(s)$

with

$Z_{n,A,B,p}(s)=\sum_{k=0}^{\infty}\frac{C(p^{k+\text{ord}_{p}(A)},-nf_{p^{k+\text{ord}_{p}(A)}}(B))}{p^{ks}}.$

We now compute the local zeta functions for various primes $p$. We will be
adapting the computations of §2.1.1 of [1].

If $p\nmid 2A$, then

$Z_{n,A,B,p}(s)=\dfrac{1-p^{-2s}}{1-p^{-s}}\cdot\dfrac{1}{1-p^{-s}\left(\frac{-nB}{p}\right)}.$

If $p\mid A$, $p\nmid n$, $p$ odd, then

$Z_{n,A,B,p}(s)=\left(1+\left(\dfrac{-nB}{p}\right)\right)\dfrac{1}{1-p^{-s}}.$

If $p\mid A$, $p\|n$, $p$ odd, then if we write
$nf_{p^{k+\text{ord}_{p}(A)}}(B)=pn_{0}$

(3) $\displaystyle Z_{n,A,B,p}(s)$
$\displaystyle=\sum_{k=0}^{\infty}\frac{C(p^{k+\text{ord}_{p}(A)},-pn_{0})}{p^{ks}}$
(4) $\displaystyle=\begin{cases}1&p\|A\\\ 0&p^{2}|A.\end{cases}$

Finally, let $p=2$. We note that this means $B$ is odd and for each $k\geq 1$,
$f_{2^{k+\text{ord}_{2}(A)}}(B)$ is also odd. We recognize a few cases based
on $\text{ord}_{2}(A)$.

If $\text{ord}_{2}(A)=0$, then we have the following possibilities:

  * •

If $n\equiv-B\mod 8$, then

$Z_{n,A,B,2}(s)=\dfrac{1-2^{-2s}}{1-2^{-s}}\cdot\dfrac{2\cdot
2^{-2s}-2^{-s}+1}{1-2^{-s}}.$

  * •

If $n\equiv-B+4\mod 8$, then

$Z_{n,A,B,2}(s)=1+\dfrac{1}{2^{s}}+\dfrac{2}{2^{2s}}.$

  * •

If $n\equiv-B+2$ or $n\equiv-B+6\mod 8$, then

$Z_{n,A,B,2}(s)=1+\dfrac{1}{2^{s}}.$

If $\text{ord}_{2}(A)=1$, then we have the following cases:

  * •

If $n\equiv-B\mod 8$, then

$Z_{n,A,B,2}(s)=1+\dfrac{1}{2^{s}}+\dfrac{4\cdot 2^{-2s}}{1-2^{-s}}.$

  * •

If $n\equiv-B+4\mod 8$, then

$Z_{n,A,B,2}(s)=1+\dfrac{2}{2^{s}}.$

  * •

If $n\equiv-B+2$ or $n\equiv-B+6\mod 8$, then

$Z_{n,A,B,2}(s)=1.$

If $\text{ord}_{2}(A)=2$, then we have the following cases:

  * •

If $n\equiv-B\mod 8$, then

$Z_{n,A,B,2}(s)=2+\dfrac{4\cdot 2^{-s}}{1-2^{-s}}.$

  * •

If $n\equiv-B+4\mod 8$, then

$Z_{n,A,B,2}(s)=2.$

  * •

If $n\equiv-B+2$ or $n\equiv-B+6\mod 8$, then

$Z_{n,A,B,2}(s)=0.$

If $\text{ord}_{2}(A)\geq 3$, then we have the following cases:

  * •

If $n\equiv-B\mod 8$, then

$Z_{n,A,B,2}(s)=\dfrac{4}{1-2^{-s}}.$

  * •

If $n\equiv-B+2$, $-B+4$, $-B+6\mod 8$, then

$Z_{n,A,B,2}(s)=0.$

The upshot of the above computation is the following proposition:

######  Proposition 2.2.

For any $n$ and $p\nmid B$ there is a rational function $\gamma_{n,A,B,p}(X)$
such that

$Z_{n,A,B,p}(s)=\gamma_{n,A,B,p}(p^{-s})\cdot\dfrac{1-p^{-2s}}{1-p^{-s}}\cdot\dfrac{1}{1-p^{-s}\left(\frac{-nB}{p}\right)},$

and if $p\nmid 2A$, then $\gamma_{n,A,B,p}(X)=1$. Furthermore, if $n_{1}\equiv
n_{2}\mod 8A$, then for all $p$,

$\gamma_{n_{1},A,B,p}(X)=\gamma_{n_{2},A,B,p}(X).$

The above proposition should be compared with Proposition 2.4 of [1]. It
should also be noted that if $B$ is odd, then the period $8A$ can be replaced
by $A$.

Going back to the two variable zeta function we get

$\zeta_{A,B}(s_{1},s_{2})=\frac{\zeta(s_{1})}{\zeta(2s_{1})}\sum_{n\text{ odd,
square-
free}}a_{n,A,B}(s_{1})\frac{L_{2AB}\left(\left(\frac{-nB}{\cdot}\right),s_{1}\right)}{n^{s_{2}}},$

where

$a_{n,A,B}(s)=\prod_{p\mid 2A}\gamma_{n,A,B,p}(p^{-s}),$

and $L_{2AB}(\left(\frac{-nB}{\cdot}\right),s_{1})$ is the Dirichlet
$L$-function of the character $\psi(k)=\left(\frac{-nB}{k}\right)$ with Euler
factors corresponding to the primes $p$ dividing $2A$ removed.

In order to prove the analytic continuation of the above zeta function we set

$\tilde{\zeta}_{A,B}(s_{1},s_{2})=\frac{\zeta(s_{1})}{\zeta(2s_{1})}\sum_{\begin{subarray}{c}n\text{
square-free}\\\
\gcd(n,2A)=1\end{subarray}}a_{n,A,B}(s_{1})\frac{L_{2AB}\left(\left(\frac{-nB}{\cdot}\right),s_{1}\right)}{n^{s_{2}}}.$

Set

$\delta_{j}(n)=\begin{cases}1&n\equiv j\mod 8A\\\ 0&n\not\equiv j\mod
8A.\end{cases}$

By the orthogonality of characters, if $\gcd(j,8A)=1$,

$\delta_{j}(n)=\frac{1}{\phi(8A)}\sum_{\chi}\chi(j)^{-1}\chi(n),$

where the sum is over all Dirichlet characters modulo $8A$. We then have

$\displaystyle\tilde{\zeta}_{A,B}(s_{1},s_{2})$
$\displaystyle=\frac{\zeta(s_{1})}{\zeta(2s_{1})}\sum_{\begin{subarray}{c}n\text{
square-free}\\\
\gcd(n,2A)=1\end{subarray}}a_{n,A,B}(s_{1})\frac{L_{2AB}\left(\left(\frac{-nB}{\cdot}\right),s_{1}\right)}{n^{s_{2}}}$
$\displaystyle=\frac{\zeta(s_{1})}{\zeta(2s_{1})}\sum_{\begin{subarray}{c}n\text{
square-
free}\end{subarray}}\sum_{j\in(\mathbb{Z}/8A\mathbb{Z})^{\times}}\delta_{j}(n)a_{j,A,B}(s_{1})\frac{L_{2AB}\left(\left(\frac{-nB}{\cdot}\right),s_{1}\right)}{n^{s_{2}}}$
$\displaystyle=\frac{\zeta(s_{1})}{\zeta(2s_{1})}\sum_{\begin{subarray}{c}n\text{
square-
free}\end{subarray}}\sum_{j\in(\mathbb{Z}/8A\mathbb{Z})^{\times}}\frac{1}{\phi(8A)}\sum_{\chi}\chi(j)^{-1}\chi(n)a_{j,A,B}(s_{1})\frac{L_{2AB}\left(\left(\frac{-nB}{\cdot}\right),s_{1}\right)}{n^{s_{2}}}$
$\displaystyle=\frac{\zeta(s_{1})}{\zeta(2s_{1})}\frac{1}{\phi(8A)}\sum_{\chi}\sum_{j\in(\mathbb{Z}/8A\mathbb{Z})^{\times}}\chi(j)^{-1}a_{j,A,B}(s_{1})\sum_{\begin{subarray}{c}n\text{
square-
free}\end{subarray}}\chi(n)\frac{L_{2AB}\left(\left(\frac{-nB}{\cdot}\right),s_{1}\right)}{n^{s_{2}}}.$

The analytic continuation of the innermost zeta function

$\sum_{\begin{subarray}{c}n\text{ square-
free}\end{subarray}}\chi(n)\frac{L_{2AB}\left(\left(\frac{-nB}{\cdot}\right),s_{1}\right)}{n^{s_{2}}}$

to a meromorphic function on all of $\mathbb{C}^{2}$ is the main result of
§2.2-2.3 of [1]. Consequently we get the following proposition:

######  Proposition 2.3.

The zeta function $\tilde{\zeta}_{A,B}(s_{1},s_{2})$ has an analytic
continuation to a meromorphic function on all of $\mathbb{C}^{2}$.

Next we treat the zeta function $\zeta_{A,B}(s_{1},s_{2})$. Write
$A=A_{1}A_{2}$ with $A_{1}$ square-free, $A_{2}$ squarefull, i.e., for all
$p$, $p\mid A_{2}$ implies $p^{2}\mid A$, and $\gcd(A_{1},A_{2})=1$. By
Equation (4) if $\gcd(n,A_{2})\neq 1$, $Z_{n,A,B}=0$. Hence,

$\zeta_{A,B}(s_{1},s_{2})=\frac{\zeta(s_{1})}{\zeta(2s_{1})}\sum_{\begin{subarray}{c}n\text{
odd, square-free}\\\
\gcd(n,A_{2})=1\end{subarray}}a_{n,A,B}(s_{1})\frac{L_{2AB}\left(\left(\frac{-nB}{\cdot}\right),s_{1}\right)}{n^{s_{2}}},$

Since $A_{1}$ is square-free, we can write $A_{1}=q_{1}\cdots q_{r}$ with
$q_{i}$’s distinct primes. For each nonempty subset $I$ of
$\\{q_{1},\dots,q_{r}\\}$, let

$\zeta_{A,B}^{I}(s_{1},s_{2})=\frac{\zeta(s_{1})}{\zeta(2s_{1})}\sum_{\begin{subarray}{c}n\text{
odd, square-free}\\\ \gcd(n,A_{2})=1\\\ \text{for all }q\in I,q\mid
n\end{subarray}}a_{n,A,B}(s_{1})\frac{L_{2AB}\left(\left(\frac{-nB}{\cdot}\right),s_{1}\right)}{n^{s_{2}}}.$

Then by inclusion-exclusion,

$\zeta_{A,B}(s_{1},s_{2})=\tilde{\zeta}_{A,B}(s_{1},s_{2})+\sum_{I\neq\varnothing}(-1)^{\\#I+1}\zeta_{A,B}^{I}(s_{1},s_{2}).$

So, it suffices to prove the analytic continuation of each
$\zeta_{A,B}^{I}(s_{1},s_{2})$. Let $t(I)=\prod_{q\in I}q$. Then since $n$ is
square-free, writing $n=n^{\prime}t(I)$,

$\displaystyle\zeta_{A,B}^{I}(s_{1},s_{2})$
$\displaystyle=\frac{\zeta(s_{1})}{\zeta(2s_{1})}\sum_{\begin{subarray}{c}n\text{
odd, square-free}\\\ \gcd(n,A_{2})=1\\\
\gcd(n^{\prime},t(I))=1\end{subarray}}a_{n^{\prime}t(I),A,B}(s_{1})\frac{L_{2AB}\left(\left(\frac{-n^{\prime}t(I)B}{\cdot}\right),s_{1}\right)}{n^{s_{2}}}$
$\displaystyle=\frac{\zeta(s_{1})}{\zeta(2s_{1})}\sum_{\begin{subarray}{c}n\text{
odd, square-free}\\\ \gcd(n,A_{2})=1\\\
\gcd(n^{\prime},t(I))=1\end{subarray}}a_{n^{\prime}t(I),A,B}(s_{1})\frac{L_{2AB}\left(\left(\frac{-n^{\prime}t(I)B}{\cdot}\right),s_{1}\right)}{n^{s_{2}}}$

The expression $a_{n^{\prime}t(I),A,B}(s_{1})$ is determined with $n^{\prime}$
modulo $8A/t(I)$ the square-free part of which has fewer prime factors than
the square-free part of $8A$. By repeating this process we may assume $A=1$,
and that, after a change of notation, we have a summation of the form

(5) $\frac{\zeta(s_{1})}{\zeta(2s_{1})}\sum_{n\text{ odd, square-
free}}A_{n}(s_{1})\frac{L_{2AB}\left(\left(\frac{-n^{\prime}t(I)B}{\cdot}\right),s_{1}\right)}{n^{s_{2}}}$

with meromorphic functions $A_{n}(s)$ such that if $n_{1}\equiv n_{2}\mod 8$
then $A_{n_{1}}(s)=A_{n_{2}}(s)$. Now an argument similiar to the proof of the
analytic continuation of $\tilde{\zeta}_{A,B}$, with $8$ replacing $8A$, gives
the analytic continuation of the zeta function in Equation (5). This gives us
the analytic continuation of $\zeta_{A,B}^{I}(s_{1},s_{2})$ to the entire
$\mathbb{C}^{2}$ which implies the analytic continuation of
$\zeta_{A,B}(s_{1},s_{2})$. Putting everything together we obtain the
following theorem:

######  Theorem 2.4.

The zeta function $Z_{A,B}(s_{1},s_{2})$ in Equation (2) has an analytic
continuation to the entire $\mathbb{C}^{2}$ as a meromorphic function.

## References

  * [1] Eun Hye Lee and Ramin Takloo-Bighash.  On a multiple Dirichlet series associated to binary cubic forms.  J. Number Theory, 238:535–556, 2022.