\documentclass[english]{jnlp_3.5}
\usepackage{jnlpbbl_2.2}
\usepackage[dvipdfmx]{graphicx}
\usepackage{amsmath}
\usepackage{amssymb}
\usepackage{url}
\usepackage{hangcaption_jnlp}
%%%%\usepackage{udline}
%%%%\setulminsep{1.2ex}{0.2ex}
%%%%\let\underline=\ul
\allowdisplaybreaks


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% author's definition
%%%%\usepackage{times}
%%%%\usepackage{latexsym}
%%%%\usepackage{mathtools}
%%%%\usepackage{amssymb}
%%%%\usepackage{svg}
%%%%\usepackage{subfig}
%%%%\usepackage{pgfplots}
%%%%\pgfplotsset{compat=1.5}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

\Volume{29}
\Number{3}
\Month{September}
\Year{2022}
\Issuetitle{Special Issue: XXXXXXXX}

\received{2021}{12}{1}
\revised{2022}{2}{24}
\accepted{2022}{4}{8}

\setcounter{page}{762}

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\esubject{General Paper}
\etitle{A Peek Into the Memory of T5: \\ Investigating the Factual Knowledge Memory in a Closed-Book QA Setting and Finding Responsible Parts}
\eauthor{Tareq Alkhaldi\affiref{Author_1} \and Chenhui Chu\affiref{Author_1} \and Sadao Kurohashi\affiref{Author_1}}
\eabstract{
Recent research shows that Transformer-based language models (LMs) store considerable factual knowledge from the unstructured text datasets on which they are pre-trained. The existence and amount of such knowledge have been investigated by probing pre-trained Transformers to answer questions without accessing any external context or knowledge (also called closed-book question answering (QA)). However, this factual knowledge is spread over the parameters inexplicably. The parts of the model most responsible for finding an answer only from a question are unclear. This study aims to understand which parts are responsible for the Transformer-based T5 reaching an answer in a closed-book QA setting. Furthermore, we introduce a head importance scoring method and compare it with other methods on three datasets. We investigate important parts by looking inside the attention heads in a novel manner. We also investigate why some heads are more critical than others and suggest a good identification approach. We demonstrate that some model parts are more important than others in retaining knowledge through a series of pruning experiments. We also investigate the roles of encoder and decoder in a closed-book setting.
} 
\ekeywords{Analysis, Attention Components, Question Answering, Transformers}


\headauthor{Alkhaldi et al.}
\headtitle{A Peek Into the Memory of T5}

\affilabel{Author_1}{}{Graduate School of Informatics, Kyoto University}


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\begin{document}

\maketitle


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\section{Introduction}

Neural language models (LMs) store knowledge from the unstructured text on which they are pre-trained \cite{roberts2020much,petroni2019language,jiang2020can}. Such a phenomenon of storing knowledge is interesting and valuable because LMs are trained on the unstructured text and are easy to query. Previous work investigated the existence and amount of such knowledge in Transformer-based models, such as BERT \cite{devlin2018bert} and T5 \cite{raffel2019exploring}, and how to effectively query models to find the stored knowledge. 


Querying these models is a form of question answering (QA), and it is considered a direct interface to the knowledge stored in LMs. Given that using context to answer questions makes it difficult to know if the model has retrieved this knowledge from within or from the context, we focus on analyzing this phenomenon using closed-book (QA) (Section \ref{sec:pre}). This setting is used by \citeA{roberts2020much} to feed T5 models open-domain questions with no context to measure the amount of knowledge stored. In open-book QA, models are either fed context that contains the answer \cite{rajpurkar2016squad} or allowed to use external knowledge sources \cite{chen2017reading}. However, in closed-book QA, the model only depends on the question and ``look up information'' stored in its parameters to find the answer. Therefore, closed-book QA is essential to estimate the amount of stored knowledge. For example, only feeding ``What is the capital of Japan?'' is sufficient for the model to generate the answer.

Although the amount of knowledge stored after pre-training is estimated using closed-book QA \cite{roberts2020much}, how this knowledge is stored and which parts of the T5 model are responsible for reaching answers remains poorly understood. This study aims to understand how knowledge is stored inside T5 in a closed-book QA setting while also studying the importance and role of different parts of the model (Section \ref{sec:understanding}). We also present head importance scoring methods for pruning the model to verify the distribution of knowledge (Section \ref{sec:pruning}). Understanding the inner workings and the crucial parts of the model allows for several benefits, such as guided pruning, where the size of LMs is reduced while preserving learned knowledge, aiding better design choices, and providing hints about how to integrate external knowledge with LMs better.

The contributions of this study are three-fold:
\begin{itemize}
\item After a novel investigation of the components of attention, we find that the query component is the most important, as it is almost the only component that changes when back-propagation is calculated.
\item We find that the encoder finds the embedding space of the desired answer entity, and the decoder generates the textual representation of that entity.
\item We compare different attention head importance scoring methods on three datasets and find that the stored knowledge is not evenly distributed in the model.
\end{itemize}


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\section{Preliminary}
\label{sec:pre}

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Task}

This study aims to understand how QA Transformer-based T5 models generate answers to factual questions in a closed-book setting. It also attempts to find whether parts of the model are more important than others in reaching answers and then identify them.

In a closed-book QA setting, the model is given a question and is expected to provide an answer without accessing context or external knowledge. The model must ``memorize'' the answers given during the pre-training or fine-tuning. How Transformers work in this setting is principally different than other tasks. In other settings, the model generally uses attention to focus on some parts of the provided context. For instance, the model can extract answers from a context paragraph in a normal QA setting \cite{devlin2018bert} or align words between source and target languages as in translation tasks. However, because only the question is provided in a closed-book QA setting, the model can only focus on the question tokens and how it decides the answer is unclear.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Datasets}
\label{sec:datasets}

Following \citeA{roberts2020much}, we use three open-domain QA datasets: Natural Questions (NQ) \cite{kwiatkowski2019natural}, WebQuestions (WQ) \cite{berant2013semantic}, and TriviaQA (TQA) \cite{joshi2017triviaqa}. NQ and WQ are both datasets of questions from web queries, while TQA is a collection of questions from quiz league websites. All questions are accompanied with context of different forms to help find the answer. However, to use the datasets as closed-book QA datasets, all context is ignored, and only the questions are considered.

Our analysis uses the T5 \cite{raffel2019exploring} models provided by \citeA{roberts2020much}. In particular, we use three variants of T5: T5-Large, T5-11b, and T5v1.1-xxl. Details about them are explained in Section \ref{sec:t5specs}.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Transformer Architecture}
\label{sec:t_arch}

We briefly describe the Transformer architecture \cite{vaswani2017attention}, which comprises the T5 model while focusing on the attention mechanism therein.

The Transformer is a model architecture with blocks of two stacked sub-layers: multi-head self-attention and feed-forward (FF) networks. Given an input sequence \textbf{x} = $\mathbf{x}_1, ..., \mathbf{x}_n \in \mathbb{R}^{d_{model}} $, for each attention head $h$, a query $Q^h$, a key $K^h$, and a value $V^h$ are computed from the input $\mathbf{x}$ as follows:
\begin{equation}
\begin{aligned}
Q^h & = \mathbf{x} W_q^h  \\ 
K^h & = \mathbf{x} W_k^h \\
V^h & = \mathbf{x} W_v^h  
\end{aligned}
\end{equation}
where $W_q^h$, $W_k^h$, and $W_v^h$ are parameter matrices for the head $h$. $Q^h$, $K^h$ and $V^h$ matrices contain $\mathbf{q}^h$, $\mathbf{k}^h$, and $\mathbf{v}^h$ vectors of different input tokens. Then, the attention scores $\alpha^h$ are computed as:
\begin{equation}\label{eq:alpha}
\alpha^h = \text{softmax}(\dfrac{Q^h{K^h}^T}{\sqrt{d_k}})
\end{equation}
where $d_k$ is the key dimensionality. Then, the contextualized $Z^h$ is computed by scaling $V^h$ by attention scores as follows:
\begin{equation}\label{eq:zh}
Z^h = \alpha^h V^h
\end{equation}
with $Z^h$ containing $\mathbf{z}^h$ vectors of the input tokens. The heads are then combined to produce $Z$ as follows:
\begin{equation}\label{eq:z}
Z = \text{concat}_{h=1}^H(Z^h) W_o
\end{equation}
where concat$_{h=1}^H$ is a function that concatenates $Z^h$ of all $H$ heads and $W_o$ is a parameter matrix.

Next, the group of attended vectors $Z$ is passed to the next component, a FF network. \citeA{vaswani2017attention} proposed using a two-layer network with ReLU activation. The final model has N layers of such blocks. When using the block in the decoder instead of the encoder, an additional sub-layer of cross-attention is added after the self-attention. This additional sub-layer receives its $K$ and $V$ from the output of the encoder.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\section{Understanding Transformer-Based T5 in Closed-Book QA}
\label{sec:understanding}

The main component of Transformers is the attention mechanism; therefore, we start by looking into its components. Then, we verify the findings, and finally investigate the role of the encoder/decoder as a whole.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Dissecting Attention Components}
\label{sec:hyp1}

Initially, we look at the gradients (or error back-propagation) resulting from the mismatch between the model prediction and the golden answer. They represent where the model attempts to fix itself to predict the correct answer. Areas with high gradients represent important parts responsible for reaching the answer. We gather these gradients and attempt to find interesting patterns.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% fig. 1
\begin{figure}[b]
\begin{center}
\includegraphics{29-3ia1f1.pdf}
\end{center}
\hangcaption{Visualization of the averaged gradients of the attention and FF parameters for T5-Large over the validation split of the NQ dataset. } 
\label{fig:vis_raw_avg}
\end{figure}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsubsection*{Gathering Gradients}
\label{sec:gather}

We evaluate the model on the validation set of NQ while calculating loss and applying a backward pass to obtain the gradients \textit{without updating the parameters}. The loss is calculated against the gold labels for all tokens in one go using the built-in loss function in the T5 implementation of the Transformers library \cite{wolf-etal-2020-transformers}. We reduce the calculated gradients for each example (consisting of a question and a golden answer) by averaging each parameter matrix. The attention parameter matrices reduced are:
\[
\begin{aligned}
W^h_q, W^h_k, W^h_v & \in \mathbb{R}^{d_k \times d_{model}} \\
W_o & \in \mathbb{R}^{H \cdot d_k \times d_{model}} 
\end{aligned}
\]
where $W^h_q$, $W^h_k$, $W^h_v$ and $W_o$ are from the attention sub-layer as explained in Section \ref{sec:t_arch}. For completeness, we also analyze the FF matrices:
\[
W_{ff1}^T, W_{ff2} \in \mathbb{R}^{ff_{size} \times d_{model}}
\]
where $W_{ff1}$ and $W_{ff2}$ are the learnable matrices of the two-layer FF networks of each block. We then average the reduced results from all the examples.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsubsection*{Visualizing Gradients}
\label{sec:vis_grads}

We use the T5-Large provided by \citeA{roberts2020much} with the NQ dataset to gather gradients, as explained previously. Fig.~\ref{fig:vis_raw_avg} visualizes the reduced gradients of the three attention types: Encoder.SelfAttention, Decoder.SelfAttention and Decoder.CrossAttention (shortened to EncSA, DecSA, and DecXA, respectively), each having 24 layers. Each row represents a layer and consists of the $16$ heads of T5-Large. We also visualize the FF parameters of the encoder and decoder.

Interestingly, as shown in Fig.~\ref{fig:vis_raw_avg}, almost only $W_q$ is changed in all attention types. While later layers are changed more in the EncSA and DecXA, the middle layers have the most change for DecSA. We notice minor changes in $W_v$ from Fig.~\ref{fig:vis_raw_avg} and almost no change in the FF network parameters. This implies that, when the model predicts a wrong answer and the gradients are applied to fix the prediction, almost only $\mathbf{q}$ vectors are changed. We apply a similar analysis to other variants of T5 (T5-11b and T5v1.1-xxl) and find similar results.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsubsection*{Confirming the Importance of Q}

Having shown the importance of \textit{Q} for correcting wrong predictions, we conduct another experiment to confirm the importance and the applicability to correctly predicted answers. We add a normal noise $N(0.0, \sigma)$ to \textbf{q}, \textbf{k}, \textbf{v}, or FF vectors and observe the performance change, where $\sigma$ is the standard deviation. A larger $\sigma$ represents a more substantial noise.

We evaluate T5-large on the NQ dataset with noise added to one attention component or FF at a time. We show the results in Fig.~\ref{fig:noise2} and compare them to the baseline of T5-large with no noise, which is 29.03 exact match (EM). Applying noise to \textbf{q} vectors degrades the performance substantially faster than other components, affirming the importance of the query component for wrongly and correctly predicted answers. The FF networks were most resilient to noise because of the drop-out training.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% fig. 2
\begin{figure}[b]
\begin{center}
\includegraphics{29-3ia1f2.pdf}
\end{center}
\hangcaption{Evaluation of T5-large on NQ dataset with noise of magnitude $\lambda$ added to different attention components and FF output.}
\label{fig:noise2}
\end{figure}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

We assume the following: (1) The behavior of not changing $W_k$, $W_v$, or $W_o$ is related to the fact that changes required to fix predictions are typically small. (2) A drastic change in the expected answer invokes more changes in, for example, $W_v$. These assumptions can be answered by considering the type of predictions and golden answers.

We use SpaCy \cite{spacy} to apply named entity recognition (NER) on all answers and compare the types of pairs of \{prediction, gold\} answers. The majority are of the same type (e.g., (GPE, GPE) or (DATE, DATE)). We then shuffle the golden answers, replacing each with a random golden answer of a different type. This would introduce many type mismatches, and the model would require bigger gradients/changes to fix the predictions. Then, we evaluate and gather the gradients again, as in Section \ref{sec:gather} and show them in Fig.~\ref{fig:vis_raw_roberts_avg}. Interestingly, the only change between Figs.~\ref{fig:vis_raw_avg} and \ref{fig:vis_raw_roberts_avg} was slightly brighter $W_q$, with almost no change to the other attention parameters. 

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% fig. 3
\begin{figure}[t]
\begin{center}
\includegraphics{29-3ia1f3.pdf}
\end{center}
\hangcaption{Visualization of the averaged gradients of the attention parameters for T5-Large over the validation split of the NQ dataset after shuffling all target answers to random answers of a different named entity type. } %The 16 heads are shown for each parameter.
\label{fig:vis_raw_roberts_avg}
\end{figure}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

To verify whether $W_k$ and $W_v$ are changed during fine-tuning, we fine-tune a raw version of T5 on the NQ dataset and apply the gradient analysis explained earlier. While still only $W_q$ was highly changed, at 50 and 250 steps, the decoder had significantly more gradients than the encoder, whereas, at 1k and 2k, the encoder had higher gradients. The results in Fig.~\ref{fig:vis_raw_avg} are after fine-tuning ends. The lack of change in $W_k$ and $W_v$ during fine-tuning implies that they were changed (and probably learned optimally) during the pre-training phase.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsubsection*{Difference Between \textit{Q} and \textit{K}}

Although $Q^h$ and $K^h$ are computed by the same input and formula (except the weights), 
\pagebreak
only $W_q$ receives high gradients. This gradient difference comes from the matrix multiplication order. To calculate the weight $a^h$, we do $Q^h{K^h}^T$, where the resulting matrix has rows that correspond to the queries and the columns to keys. Then, we take softmax on the last dimension, that is, by row. The softmax lets the row sum to 1.0, which would be the weights sum of our current query word attending to all other key words. The model learns this distinction between queries and keys and learns to change the queries rather than the keys.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsubsection*{How Knowledge Is Stored Inside T5}

As experiments (Figs.~\ref{fig:vis_raw_avg} and \ref{fig:noise2}) show that \textbf{q} vectors are the most crucial component for fixing predictions or reaching answers, we conclude that the following is how T5 stores knowledge.

In Eqs. \ref{eq:alpha} and \ref{eq:zh}, the role of \textbf{q} vectors is only to calculate a score that represents how much \textbf{v} to take from each token. Multiplying a scaling score with a vector does not change its direction, but only the norm of that vector. The \textbf{z} vectors are only composed from combining differently weighted \textbf{v} vectors (Eq. \ref{eq:zh}). For easier imagination, \textbf{v} vectors can be considered pointers or arrows of different lengths that point somewhere in the embedding space. When \textbf{z} vectors are combined cumulatively (over layers), they point to a place in the embedding space that represents the desired answer. In summary, $Q^h$ is a map that guides the $V^h$ arrows/pointers to this final point in the embedding space.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Role of Encoder/Decoder}

Thus far, we have discussed attention head components as parts of the encoder--decoder model. Encoders and decoders usually have straightforward roles in other tasks, such as focusing on source/target languages and performing alignment in translation tasks, or selecting text spans in open-book QA. However, in a closed-book QA setting, their role is unclear. To investigate their role in a closed-book QA setting, we remove attention heads from one attention type at a time (e.g. EncSA) or FFs from the encoder or decoder and check the differences in predictions before and after a significant drop in performance. The procedure for removing heads is introduced in Section \ref{sec:pruning}. We apply the same procedure for removing FFs.
Comparing the results by removing EncSA heads and DecSA heads presented in Table \ref{tab:encdec_diff}, the role of the encoder is finding the correct entity while the decoder generates the textual representation of that entity. We did not witness special patterns when removing heads only from DecXA.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% table 1
\begin{table}[t]
\input{01table01.tex}
\hangcaption{Patterns of the differences in predictions when only pruning EncSA vs DecSA. Each row's first and second lines show the prediction before and after pruning.}
\label{tab:encdec_diff}
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% table 2
\begin{table}[t]
\input{01table02.tex}
\hangcaption{Quantitative comparison between the predictions before and after substantial pruning (not accuracy against gold answers).}
\label{tab:encdec_quantitative}
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

Table \ref{tab:encdec_quantitative} shows a quantitative comparison of how predictions change before and after substantial pruning (with about 10~EM points drop). All predictions after pruning can either match predictions before pruning (EM), differ but with textual overlap, or become different with no overlap. Pruning the decoder heads does not change many predictions (high EM with predictions before pruning). It produces more wrong but textually overlapping predictions because the entity space has been decided in the encoder. However, harming the encoder heads results in predicting about 30\% more non-textually overlapping predictions (i.e., different entities). Pruning the FFs in the encoder and decoder has a more negligible effect than pruning heads, but still follows the same trend.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\section{Head Importance for Pruning}
\label{sec:pruning}

In a closed-book QA setting, the model must generate the correct answer with only access to the question tokens. The ``knowledge'' about the correct answer is within the model's parameters. We hypothesize that such ``knowledge'' is not distributed equally in the model, but there are areas that are significantly more important than others in retaining that knowledge.

To verify that some heads of the model are more important than others in retaining knowledge, we perform a series of pruning tests on attention heads according to different importance scoring methods. We explore pruning least-important heads first or most-important heads first. The former provides a practical advantage where pruning models and reducing their parameters count while preserving their knowledge allows for faster and lighter models. Because FF networks do not store factual knowledge, as discussed in Section \ref{sec:gather}, we exclude them from further pruning experiments.

We explain the pruning settings and methods, visualize different importance maps, and compare them via pruning experiments on different datasets.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Pruning Flow}

An importance score is given for each attention head. Then, the least/most 10\% important heads are masked iteratively by adding a mask variable $\xi^h$ following \citeA{michel2019sixteen}, changing Eq. \ref{eq:zh} as follows:
\begin{equation}\label{eq:masking}
Z^h = \xi^h \alpha^h V^h
\end{equation}
where the values of $\xi^h$ are either $0$ or $1$. After each masking round, the model is re-evaluated and a new head importance map is calculated for the next masking round until we reach a threshold.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Importance Scoring Methods}
\label{sec:prune_methods}

For each of the three attention types (EncSA, DecSA, and DecXA), a head importance map is generated using one of the following methods:


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsubsection{Random}

Heads are given random importance scores. This represents a counter hypothesis that knowledge is distributed equally in all of the model.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsubsection{Gradients of Mask Gates (\textit{attn})}
\label{sec:attn_method}

The importance score for a head is the gradient of the mask variable $\xi^h$. This method, introduced by \citeA{michel2019sixteen}, focuses on the value of the attention score $\alpha^h$. When the attention score is high, the gradient of $\xi^h$ is also high; therefore, this method is called \textit{attn}.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsubsection{Summed Gradients of Attention Parameters (\textit{grad\_avg})}
\label{sec:grad_method}

Looking at the analysis result in Section \ref{sec:hyp1}, we introduce a new method of calculating the importance score by summing the reduced gradients of the parameter matrices in attention heads. More specifically, for each head $h$, we calculate scores $grad_{avg}^h$ as follows:
\[
    grad_{avg}^h = \text{avg}(g(W_q^h)) + \text{avg}(g(W_k^h)) + \text{avg}(g(W_v^h))
\]
where $g()$ is the function that obtains the gradients of a parameter matrix, and avg is a function that obtains the average of a matrix.\footnote{We also experimented with max for scoring; however, we excluded it because its performance was poor.} Since almost only $W_q$ has high gradients among attention components (as shown in Fig.~\ref{fig:vis_raw_avg}), adding the gradients of $W_k$ and $W_v$ has almost no effect on the final importance map. Therefore, we can consider that the importance of an attention head is equivalent to the importance of only the query component of that head. However, in this method, we still add the gradients of all components for completeness.\footnote{We calculated the KL-divergence between importance maps with the gradients of all components, and maps of only $W_q$ gradients. As the KL-divergence was 0.0001, confirming similar distributions, we exclude the latter maps.}


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsubsection{Norm of $\mathbf{z}^h$ Vectors (\textit{norm})}
\label{sec:norm_method}

This method does not depend on the gradients but on the norm of $\mathbf{z}^h$ vectors. \citeA{kobayashi-etal-2020-attention} used the norm for analyzing alignment in machine translation. We use it as a scoring method; for each head $h$, the importance score $norm^h$ is calculated as follows:
\[
norm^h =  \sum_{t=1}^n || \mathbf{z}^h_t ||  
\]
where $\mathbf{z}^h_t$ is the contextualized vector $\mathbf{z}^h$ of token $t$, and $||\cdot||$ denotes the Euclidean norm. 


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Pruning Settings}

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsubsection{Importance Sorting}
\label{sec:imp_sorting}

\begin{itemize}
\item Local: To mask 10\% of the heads, the importance scores are sorted per attention type \textit{independently}. Then, the 10\% are divided equally between the EncSA, DecSA, and DecXA attention types.
\item Global: The maps of the three attention types are concatenated, and all the importance scores are sorted before selecting 10\% to mask. Here, more heads might be masked from DecSA than EncSA, for example.
\end{itemize}


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsubsection{Normalization}
\label{sec:nln}

\begin{itemize}
    \item Layer normalization: \citeA{michel2019sixteen} normalize importance scores by layer using the $\ell_2$ norm. We use this as the default.
    \item No layer normalization: We skip the layer normalization and denote the setting as \textit{nln}.
\end{itemize}


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsubsection{Used Pre-trained Models}
\label{sec:t5specs}

We use three variants of T5 released by \citeA{roberts2020much}\footnote{The pre-trained models can be downloaded from: \url{https://github.com/google-research/google-research/tree/master/t5_closed_book_qa}}: T5-large, T5-11b, and T5v1.1-xxl. They all have $24$ layers but differ in several aspects, as shown in Table \ref{tab:models_diff}. T5v1.1 also differs from other models as it has no parameter sharing between embedding and classifier layers, and it was pre-trained on C4 \cite{raffel2019exploring} only without mixing in the downstream tasks. It also uses GatedGelu instead of Relu for activations. Both T5v1.1-xxl and T5-11b have around 11 billion parameters.

All models had additional pre-training with salient span masking \cite{guu2020realm} before being fine-tuned. In this objective, salient spans (named entities and dates) are mined beforehand and then used in the pre-training as masks. A fine-tuned model is released for each of the three datasets (NQ, WQ, and TQA), except T5-large, which only has a version fine-tuned on NQ.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% table 3
\begin{table}[b]
\input{01table03.tex}
\hangcaption{Main differences between used models. d\_model is the dimension of the word embeddings, \#heads is the number of heads, d\_kv is the hidden dimension of the query, key, and value vectors, and d\_ff is the dimension of the FF networks after the attention sub-layer.}
\label{tab:models_diff}
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Visualizing Head Importance Maps}
\label{sec:vis_impmaps}

Fig.~\ref{fig:impmaps} shows the head importance maps generated with T5-large from the validation set of the NQ dataset for each scoring method. In ``nln'' all methods assign more importance to later layers in the EncSA, but have different assessments about the DecSA. For DecXA, all methods agree on the importance of certain heads in the last layer. Although there are differences in general between methods, the importance patterns are very similar across datasets, as shown in Fig.~\ref{fig:datasets_diff_ours_11bxxl}. This is likely because even though models were fine-tuned separately on the datasets, they all share the same pre-training, which is the actual source of the stored knowledge. This similarity is verified by pruning experiments in Section \ref{sec:cross_maps}.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% fig. 4
\begin{figure}[t]
\begin{center}
\includegraphics{29-3ia1f4.pdf}
\end{center}
\hangcaption{Head Importance score maps for different scoring methods with T5-large on the NQ dataset. ``nln'' means no layer normalization and ``ln'' means layer-normalized.} 
\label{fig:impmaps}
\end{figure}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% fig. 5
\begin{figure}[t]
\begin{center}
\includegraphics{29-3ia1f5.pdf}
\end{center}
\hangcaption{Comparing head importance visualizations of the encoders of T5-11b and T5v1.1-xxl models between different datasets, using \textit{grad\_avg} with ``nln''}
\label{fig:datasets_diff_ours_11bxxl}
\end{figure}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Experiments}
\label{sec:experiments}

We experiment with different scoring methods and settings 
\pagebreak
with T5-large on the NQ dataset for speed. Then, we apply the best settings on the huge T5-11b and T5v1.1-xxl variants on the other datasets. We focus on ``least-important heads first'' in our experiments as it has a more practical benefit.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsubsection{Least-Important Heads First (LIF)}
\label{sec:lif}

We prune the least important heads first and do not expect to see a large change in performance. We stop pruning when the EM score reaches below $90$\% of the original score. In Fig.~\ref{fig:graphs1}, we show scoring methods compared to the random baseline. The global sorting setting is usually ineffective and achieves below the random baseline; therefore, we ignore it in further experiments. Layer normalization also hurts the importance scores; however, we notice different trends on huge models. We discuss the reasons later. Furthermore, since the \textit{attn} method's performance is closer to the random baseline than other methods, we exclude the \textit{attn} method from comparison. The \textit{attn} method showed relatively reasonable pruning performance in machine translation tasks \cite{michel2019sixteen} because focusing on which source/target language tokens are attended is essential. \citeA{kobayashi-etal-2020-attention} found that only focusing on attention score ignores the norm aspect of vectors.\footnote{Experiments of \textit{attn} on other datasets and variants were also conducted, but because of poor performance we exclude it from graphs for visibility.}

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% fig. 6
\begin{figure}[t]
\begin{center}
\includegraphics{29-3ia1f6.pdf}
\end{center}
\hangcaption{Comparing pruning experiments with different settings and methods against the random baseline. We compare \textit{attn}, \textit{grad\_avg} and \textit{norm} methods introduced in Section \ref{sec:prune_methods}. In the legend, ``\_local'' and ``\_global'' mean local and global sorting explained in Section \ref{sec:imp_sorting}, while ``\_nln'' means ``no layer normalization'' as explained in Section \ref{sec:nln}. } %The 16 heads are shown for each parameter.
\label{fig:graphs1}
\end{figure}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


We run the remaining settings with the huge variants (T5-11b and T5v1.1-xxl) on all datasets and show the results in Fig.~\ref{fig:all_huge_lif}. While the best pruning method in T5-large is the non-layer-normalized \textit{grad\_avg} as shown in Fig.~\ref{fig:graphs1}, T5v1.1-xxl brings the performance of the \textit{norm} method closer to \textit{grad\_avg}. In T5-11b, the \textit{norm} method outperforms the non-layer-normalized \textit{grad\_avg} and slightly outperforms the \textit{grad\_avg}.

We think the reason is that the standard deviation between examples for the \textit{norm} scores of heads is mostly lower than that for the \textit{grad\_avg} method, especially in the areas of low importance. However, small models have a more uniform importance map when normalized, as shown in Fig.~\ref{fig:impmaps}. Such distribution leads the pruning towards a more random selection process. In models with many heads, layer-normalization produces less uniform maps. 

Figs.~\ref{fig:graphs1} and \ref{fig:all_huge_lif} show that for \textit{grad\_avg}, layer normalization becomes more important as the number of heads in a model increases. We find that non-normalized importance maps tend to assign more importance to later layers as models increase in size, and least important heads just become ``heads of earlier layers.'' However, when normalized by layer, the least important heads from all layers are pruned. This is not a problem in small models because when the number of heads is small (e.g., $16$), heads of medium importance are not condensed only in later layers. We focus on heads of medium importance here because we notice in Fig.~\ref{fig:all_huge_lif} that the sharpest drop between normalized and non-normalized experiments is when $50$--$60$\% of heads are remaining. 

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% fig. 7
\begin{figure}[t]
\begin{center}
\includegraphics{29-3ia1f7.pdf}
\end{center}
\hangcaption{Least-important first pruning experiments on three datasets with T5-11b and T5v1.1-xxl models, using the best performing settings and methods.}
\label{fig:all_huge_lif}
\end{figure}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

To understand the intuition of why performance does not decrease much at the beginning of the pruning, we point out that removing the least important heads in the \textit{norm} method is akin to removing the least impactful $\mathbf{z}$ vectors, as vectors with a small norm value are closer to a zero vector. Intuitively, smaller vectors contribute less than longer ones in reaching the embedded location of an answer when combined. For the \textit{grad\_avg} method, the least changed heads mean correcting the vectors in those heads was not necessary to find the correct answer. Therefore, they are probably not crucial for finding the answer.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% fig. 8
\begin{figure}[b]
\begin{center}
\includegraphics{29-3ia1f8.pdf}
\end{center}
\caption{Pruning most important heads first with different methods using T5-large on the NQ dataset.} %The 16 heads are shown for each parameter.
\label{fig:large_mif}
\end{figure}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% fig. 9
\begin{figure}[b]
\begin{center}
\includegraphics{29-3ia1f9.pdf}
\end{center}
\hangcaption{Most-important first pruning experiments on three datasets with T5-11b and T5v1.1-xxl models, using the best-performing settings and methods.} %The 16 heads are shown for each parameter.
\label{fig:all_huge_mif}
\end{figure}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsubsection{Most-important Heads First (MIF)}
\label{sec:mif}

In this mode, we prune important heads first and expect to see a rapid drop in performance. We stop pruning when 50\% of heads are pruned. We show a method comparison in Fig.~\ref{fig:large_mif} only with ``nln'' setting, as we find that normalizing by layer loses valuable information about important heads. The \textit{norm} method finds important heads better in T5-large; however, we find that, although this behavior is consistent between datasets, it differs with model changes. In T5-11b, \textit{grad\_avg}  and \textit{norm} have similar performance, while in T5v1.1-xxl, \textit{grad\_avg} drops faster than \textit{norm}. We show the graphs in Fig.~\ref{fig:all_huge_mif}. An exception to the pattern of T5-11b is the TriviaQA dataset, where the \textit{norm} method performs worse. We think this is because the size of the evaluation set is over three times that of the other two datasets. This causes the \textit{norm} scores of heads to be averaged more flatly over the 128 heads, with the data distribution being another possible factor. 

As the main change in behavior (i.e., \textit{grad\_avg} and \textit{norm} methods swapping places in which one drops the performance faster) is between T5-large and T5v1.1-xxl, we look at all the differences explained in Section \ref{sec:t5specs}. Then, list the main differences directly influencing the scoring methods between the models and discuss their likelihood of being the responsible change:

\begin{itemize}
    \item Hidden model size (d\_model): T5-large has d\_model of $1,024$, while T5v1.1-xxl has $4,096$. While $\mathbf{Z}^h \in \mathbb{R}^{d\_kv}$ did not change between the models, the \textit{grad\_avg} was performed on $64 \times 4,096$ instead of $64 \times 1,024$, giving a better hint at what is important because of viewing more gradients, making this change likely responsible.
    \item The number of heads: T5-large has 16 heads, while T5v1.1-xxl has 64 heads. This increase in heads may not be the reason for the change because T5-11b has 128 heads, yet the two methods perform similarly.
    \item Feed-forward layers size (d\_ff): For T5-large it is $4,096$, while for T5v1.1-xxl it is $10,240$. This change is not responsible for two reasons: T5-11b has a more significant FF layer size of $65,536$ with no drastic behavior change, and the gradients in the FF parameters are negligible, as shown in Fig.~\ref{fig:vis_raw_avg}.
    
\end{itemize}

Therefore, we conjecture that the reason for the behavior change is the d\_model size. However, many experiments (of changing each factor and observing the performance change) are required to prove all hypotheses. Because their computational costs are unaffordable for a university research lab, this verification is left for future work.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsubsection{Head Similarity Between Datasets}
\label{sec:cross_maps}

Fig.~\ref{fig:datasets_diff_ours_11bxxl} showed that models have similar head importance maps across different datasets because of the underlying pre-training phase. We prove this similarity in this section.

We generate three importance maps for each model, one for each dataset (nqo\_map, wqo\_map, tqao\_map) using the \textit{grad\_avg} method explained in Section \ref{sec:grad_method}. For example, nqo\_map of T5-11b is a head importance map generated by evaluating the T5-11b on the NQ dataset, while the wqo\_map is generated with the WQ dataset, and so on. We then apply important-first pruning on a model and a dataset using each of the generated maps (the one generated on the dataset itself and on other datasets as well) and compare the results against a baseline of random head importance, random\_map. The results are shown in Fig.~\ref{fig:cross_map_comparison}. The important heads for one dataset are also important for the others because of the shared knowledge in the pre-training phase.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% fig. 10
\begin{figure}[b]
\begin{center}
\includegraphics{29-3ia1f10.pdf}
\end{center}
\hangcaption{Most-important first pruning experiments on three datasets with T5-11b and T5v1.1-xxl models, using \textit{grad\_avg} with ``nln'' and the importance maps that were calculated on different datasets.} 
\label{fig:cross_map_comparison}
\end{figure}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% fig. 11
\begin{figure}[b]
\begin{center}
\includegraphics{29-3ia1f11.pdf}
\end{center}
\hangcaption{Distribution of knowledge inside T5-large attention heads depending on the named entity type of the target answer, calculated using the \textit{grad\_avg} method.}
\label{fig:ner_types}
\end{figure}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Distribution of Knowledge per Answer Type}

In Section \ref{sec:experiments}, we showed that the knowledge stored in T5 is not evenly distributed and that some attention heads store more knowledge than others. However, a question arises: are there heads that are specialized in certain types of questions? To investigate the distribution for different types of questions (who, where, when, etc.), we analyze the named entity type of the target answer (person, location, date, etc.).

We used SpaCy to do NER on all target answers and we separated the \textit{grad\_avg} importance map with ``nln'' calculated in Section \ref{sec:vis_impmaps} by named entity type. The results are shown in Fig.~\ref{fig:ner_types}. Different named entities have different patterns of importance. To determine the similarity of importance maps between entity types, we sum the importance of each type per area (EncSA, DecSA or DecXA).\footnote{We calculated divergence per head and layer, but the results were not meaningful, as one head or layer may not hold enough information about many types to deduce similarity.} Then, we calculate the Kullback-Leibler (KL) divergence between different entity types and show the results in Fig.~\ref{fig:kl_div_ners}. We sort the entity types by similarity and notice the correlation with the KL-divergence. In general, similar entity types are stored in similar locations. For example, geopolitical entity (GPE), location (LOC) and facilities (FAC) have similar distribution because they are essentially locations. CARDINAL, QUANTITY, and ORDINAL are similar because they are all related to numbers. DATE and TIME are both related to time.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% fig. 12
\begin{figure}[t]
\begin{center}
\includegraphics{29-3ia1f12.pdf}
\end{center}
\hangcaption{KL-divergence of importance maps per area between named entity types. Lower divergence represents higher similarity. }
\label{fig:kl_div_ners}
\end{figure}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

We note that the results in Section \ref{sec:pruning} might be dependent on the three experimented datasets. In future work, we want to experiment with other datasets, including ones with different domains.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\section{Related Work}

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Confirming the Existence of Knowledge}

Several works investigating knowledge probing are surveyed by \citeA{liu2021pre}: \citeA{petroni2019language} have introduced the LAMA probe that tests the factual and common-sense knowledge in LMs. Using four knowledge sources, they convert the fact triples or question-answer pairs into cloze templates, and they evaluate models on how highly they rank the ground truth token. They found that factual knowledge can be recovered ``surprisingly well'' from pre-trained LMs. Some work explored continuous template learning instead of discrete prompt template search. Also tested on the LAMA benchmark, \citeA{zhong-etal-2021-factual} proposed a model OptiPrompt that optimizes in continuous embedding space to find prompts, rather than being limited to a space of discrete tokens.

\citeA{roberts2020much} estimated how much knowledge is stored in LMs by evaluating T5 models on open-domain QA in a closed-book setting. Compared to other competing models, they show that LMs can achieve competitive results on open-domain QA benchmarks without access to context or external knowledge. Different from previous work, we investigate how Transformer-based T5 stores knowledge for closed-book QA. 


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Methods of Importance Scoring}

\citeA{michel2019sixteen} found that a percentage of heads can be removed with no significant impact on the performance. Their criteria for assigning the importance of heads are explained in Section \ref{sec:attn_method}. They effectively give a higher importance score if the attention score is high, which, as presented in Section \ref{sec:experiments}, is insufficient. \citeA{voita2019analyzing} also focused on the translation task and used several scoring methods: head ``confidence'' is the average of its maximum attention score over tokens in a sentence. As this method focuses on the ``maximum'' attention score, it will have the similar problems as the \textit{attn} method in Section \ref{sec:attn_method}. They also used a method called layer-wise relevance propagation (LRP), which has similar behavior to the ``confidence'' method. They also performed pruning experiments based on stochastic gates and relaxed $L_0$ penalty criteria. However, their pruning method requires fine-tuning to prune heads, and their results coincide with that of their LRP analysis. Hence, we exclude their methods from our comparison. In this work, we introduce other methods and compare them with the \textit{attn} method.

While we only focus on attention head pruning, \citeA{prasanna2020bert} experimented with other techniques, such as magnitude pruning, where sparse weights are pruned according to their magnitude and whole layers pruning. However, these techniques are not informative as they are either quite fine-grained or overly general.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\section{Conclusion}

This study provided an insight into how Transformer-based T5 operates in a closed-book QA setting. We showed what components of the T5 are important in retrieving the factual knowledge within. We explored the importance of heads in retaining knowledge via pruning while using different settings and scoring methods. In the future, we will investigate other Transformer-based LMs, and how they work on NLP tasks with context, such as open-book QA.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\bibliographystyle{jnlpbbl_1.6}
\bibliography{01refs}


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\begin{biography}

\bioauthor[:]{Tareq Alkhaldi}{
 received his B.S. in Information Technology from Arab American University in 2010, and his M.S. in Informatics from Kyoto University in 2018. He is currently working towards a doctoral degree at Kyoto University. His research interests include natural language processing, and in particular, question answering and knowledge representation.
}

\bioauthor[:]{Chenhui Chu}{ received his B.S. in software engineering from Chongqing University in 2008, and his M.S. and Ph.D. in Informatics from Kyoto University in 2012 and 2015, respectively. He is currently a program-specific associate professor at Kyoto University. His research interests include natural language processing, particularly machine translation and multimodal machine learning.
}

\bioauthor[:]{Sadao Kurohashi}{ received the B.S., M.S., and Ph.D. in Electrical Engineering from Kyoto University in 1989, 1991 and 1994, respectively. He has been a visiting researcher of IRCS, University of Pennsylvania in 1994. He is currently a professor of the Graduate School of Informatics at Kyoto University. His research interests include natural language processing, knowledge acquisition/representation, and information retrieval. He received the 10th anniversary best paper award from the Journal of Natural Language Processing in 2004, 2009 Funai IT promotion award, and 2009 IBM faculty award.}

\end{biography}


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\biodate




\end{document}

