Revisiting the Architectures like Pointer Networks to Efficiently Improve the Next Word Distribution, Summarization Factuality, and Beyond
Year: 2023

Abstract:
Is the output softmax layer, which is adopted by most language models (LMs), always the best way to compute the next word probability? Given so many attention layers in a modern transformer-based LM, are the pointer networks redundant nowadays? In this study, we discover that the answers to both questions are no. This is because the softmax bottleneck sometimes prevents the LMs from predicting the desired distribution and the pointer networks can be used to break the bottleneck efficiently. Based on the finding, we propose several softmax alternatives by simplifying the pointer networks and accelerating the word-byword rerankers. In GPT-2, our proposals are significantly better and more efficient than mixture of softmax, a state-of-the-art softmax alternative. In summarization experiments, without significantly decreasing its training/testing speed, our best method based on T5-Small improves factCC score by 2 points in CNN/DM and XSUM dataset, and improves MAUVE scores by 30% in BookSum paragraph-level dataset.

heading: 1 Introduction
text: When recurrent neural networks such as LSTM (Hochreiter and Schmidhuber, 1997) are the mainstream language model (LM) architecture, pointer networks, or so-called copy mechanisms (Gu et al., 2016), have been shown to improve the state-of-the-art LMs for next word prediction (Merity et al., 2017) and summarizations (See et al., 2017) by a large margin. However, after transformer (Vaswani et al., 2017) becomes the dominating LM architectures, the pointer networks are rarely used in the state-of-the-art pretrained LMs.One major reason is that the attention mechanism in every transformer layer can learn to copy the words from the context, so it
∗indicates equal contribution †The work is done while the author was at UMass
seems to be redundant to add a copying mechanism on top of the transformer.
In this paper, we demonstrate that the architectures like pointer networks can still substantially improve the state-of-the-art transformer LM architectures such as GPT-2 (Radford et al., 2019) and T5 (Raffel et al., 2020) mainly due to breaking the bottleneck of their final softmax layer (Yang et al., 2018; Chang and McCallum, 2022).
In Figure 1, we illustrate a simple example from Chang and McCallum (2022) to explain the softmax bottleneck and why the pointer networks could alleviate the problem. When predicting the next
12707
GPT-2
woman man
queen king
king or woman
Context Partition / Reranker Partition
hct,V
Softmax
Global Word Embeddings
hct,S
… king … woman first , the jester decided on the
Context / Top k Vocabulary
Dot Product Dot Product
word, most LMs would try to output a hidden state hct,V that is close to all the next word possibilities. For example, when the next word should be either king or woman with similar probabilities, the ideal hidden state is supposed to be the average of the global output word embeddings of king and woman. However, there might be other interfering words (queen and man in this case) between the ideal next word candidates, which force the LM to output the wrong distribution.
To solve this problem, we can let the LMs predict the probability of copying the words in the context separately by paying attention to the previous hidden states (Gu et al., 2016) and we call this kind of architecture pointer networks in this paper. That is, we can compute the dot products with the hidden states of king he,k and the hidden states of woman he,w rather than with their global output word embeddings in order to estimate the probabilities of copying these two words in the context. Our experiments show that the pointer networks consistently improve the performance of GPT-2 in next word prediction and the quality of summarization from T5 and BART.
Contrary to the mainstream explanation in previous pointer network literature, we discover that most of the improvements in our experiments do not come from the attention mechanism. To study these improvements, we propose a very simple pointer network variant that does not use any previous hidden states and we show that the proposed method can achieve similar improvements.
As shown in Figure 2, we simply project the last hidden state into two embeddings. One embedding
hct,S is to compute the dot product with the context words, and hct,V is for the dot product of the other words. Then, the GPT-2 can output the hidden state for context words hct,S as the average embedding of the king and woman without interfered by the words of man and queen that are handled by hct,V . We call this method context partition. In addition to words in the context, we can also use another embedding for the top-k likely next words. This can be viewed as a very simple and efficient alternative to a reranker, so we call it reranker partition.
In our experiments, we show that the context partition performs similarly to pointer networks while combining a pointer network, context partition, and reranker partition would significantly outperform each individual method. Compared to the state-ofthe-art solutions for alleviating the softmax bottleneck such as mixture of softmax (Yang et al., 2018; Chang and McCallum, 2022), our proposed method is more efficient while achieving lower perplexity on GPT-2. Furthermore, we find that adding a very expensive word-by-word reranker only improves our method slightly, which suggested the difficulty of further improving the final softmax layer over the proposed alternatives.
In the text completion task using GPT-2, we find that the proposed softmax alternatives reduce hallucination by copying more proper nouns from the context even though we did not provide any partof-speech information during training. In summarization, our methods and pointer networks output a more specific summary, increase the factuality, and consistently improve 9 metrics, especially in the smaller language models. Finally, we show that the softmax bottleneck problem is not completely solved in GPT-3.5 in the limitation section.

heading: 1.1 Main Contributions
text: • We propose a series of efficient softmax alter-
natives that unify the ideas of pointer network, reranker, multiple embeddings, and vocabulary partitioning.1
• We evaluate the proposed softmax alternatives in text completion tasks and summarization tasks using various metrics to identify where our methods improve the most.
• Our experiments indicate pointer networks and our proposed alternatives can still improve the modern transformer-based LMs. By breaking
1Our codes are released at https://github.com/iesl/ Softmax-CPR
the softmax bottleneck, our methods learn to sometimes copy the context words to reduce generation hallucination and sometimes exclude the context words to reduce the repetition. Besides, we find that the softmax bottleneck problem won’t be completely solved by the huge size of GPT-3.5.

heading: 2 Background
text: Before introducing our method, we would first briefly review the problem we are solving and its state-of-the-art solutions.

heading: 2.1 Softmax Bottleneck Problem
text: Most LMs use a softmax layer to compute the final probability of predicting the word x:
PM (x|ct) = exp(Logit(x, ct))∑ x′ exp(Logit(x′, ct)) , (1)
where ct is the context words. Typically, the logit Logit(x, ct) = (hMct )
Twx, hMct is the M th-layer hidden state given the input context ct and wx is the output word embeddings for x.
One problem is that the output word embeddings wx are global and independent to the context. After pretraining, the similar words would have similar output word embeddings. However, the similarity structure in the word embedding space might prevent LMs from outputting the desired distribution. The parallelogram structure among the embeddings of king, queen, woman, and man is a simple example. Chang and McCallum (2022) generalize this observation and show that some words in a small subspace would create some multi-mode distributions that a LM cannot output using a single hidden state hct in the softmax layer.

heading: 2.2 Mixture of Softmax Method
text: To overcome the bottleneck, one natural solution is to have multiple hidden states and each hidden state corresponds to a group of possible words (Yang et al., 2018). For example, we can have one hidden state for king and another hidden state for woman.
One major concern of this mixture of softmax (MoS) approach is the computational overhead. MoS needs to compute the final softmax multiple times and merge their resulting distributions. That is, we need to compute the dot products between every hidden state and all the words in the vocabulary, which is expensive especially when the vocabulary size is large.

heading: 2.3 Multiple Input State Enhancement
text: In MoS, the multiple hidden states come from the linear projections of the last hidden state. Chang and McCallum (2022) point out that the total degree of freedom among the multiple hidden states is limited by the dimensionality of the hidden state.
To allow LMs to move multiple hidden states more freely, Chang and McCallum (2022) propose to concatenate the projection of a block of hidden state with the last hidden state hMct so as to increase its dimensionality:
qct = h M ct ⊕GELU ( Lh(⊕i,mhM−mct−i ) ) , (2)
where GELU is the non-linear transformation used in GPT-2 and Lh is a linear transformation that allows us to consider more hidden states without significantly increasing the model size. ⊕i,mhM−mct−i is the concatenation of a block of hidden states. We set the block size to be 3x3 in our GPT-2 experiments and 1x3 in our summarization experiments (i.e., considering the last 3 hidden states in the last layer as shown in Figure 3).

heading: 3 Methods
text: To break the softmax bottleneck more efficiently compared to MoS, our overall strategy is simple. If we can identify a small partition of words that are very likely to become the next word, we can just compute the dot products between a hidden state and the embeddings of these likely words instead of all the words as in MoS. For example, if we can identify king and woman are much more likely to appear than queen and man, we can only compute the dot product between a hidden state and the embeddings of king and woman without being interfered by other words.
Specifically, when we compute the next word probability in Equation 1, the logit of the word x given the context ct
Logit(x, ct) = { fTct,Sex if x ∈ S fTct,V wx O/W , (3)
where fct,S = L f S(qct) and fct,V = L f V (qct) are the linear projections of the hidden state concatenation qct in Equation 2. As shown in Table 1, different softmax alternatives have different ways of constructing this set S and use different word embeddings ex.
To simplify our explanation, we will focus on the decoder-only LM (i.e., GPT-2) first and extend our method to encoder-decoder LM (i.e., T5 and BART).

heading: 3.1 GPT-2
text: We will explain each softmax alternative individually and their connections to previous work such as pointer networks or rerankers.

heading: 3.1.1 Pointer Network (P) as Local Word Embedding
text: Similar to Pointer Sentinel (PS) (Merity et al., 2017), we treat the words in the context differently (S = {x|x ∈ ct}) and let their word embeddings ex come from the previous hidden states:
ex = fx,ct,LD =
∑t i=1 1cit=x
LfLD(qcit)∑t i=1 1cit=x , (4)
where cit is the ith input words in the context ct, LfLD is a linear layer, and 1cit=x = 1 if c i t = x.
As a result, we can use the GPT-2 model to not only predict the hidden state fct,S = fct,PD =
LfPD(qct) and fct,V but also predict the word embedding of context words ex. Unlike the global word embedding wx, the local word embedding ex is context-dependent, so the LM can break the softmax bottleneck by adjusting the similarity of words based on the context. For example, GPT-2 could increase the similarity between eking and ewoman to output the high probability for both words easily.
We call this version of pointer network local decoder (LD) embedding, which has some minor differences compared to PS (Merity et al., 2017) and other variants. For example, we merge their logits while PS merges their probabilities. PS does not do normalization when computing ex. In our experiments, we would show that these pointer network variants all have very similar improvements in modern LMs.

heading: 3.1.2 Context Partition (C)
text: To understand the source of the improvements from pointer networks, we simplify their architectures by setting the word embedding ex = wx and the partition S is still the set of context words. Although much simpler, the LM with this context partition method can still break the softmax bottleneck by properly coordinating the hidden state specifically for the context words fct,S = fct,C = L f C(qct) and the hidden state for other words fct,V . Compared to the pointer network, one advantage of context partition is that the LM can still leverage the learned global word similarity when estimating
the probabilities of context words.

heading: 3.1.3 Reranker Partition (R)
text: In some cases, the possible next words might not be mentioned in the context. For example, in the context My favorite actor is Ryan [MASK], the next word could be Reynolds, Gosling, or the last names of other Ryan. Hence, using only the context partition does not completely solve the multimodal distribution problem.
Inspired by the idea of the reranker, we set S to be the top k words with the highest logits fTct,V wx. In practice, finding an ideal k could be difficult. When k is small, the reranker partition might not include the very likely next word. When k is large, the reranker partition might not be able to separate the output candidates and the interfering words. To alleviate the problem, we can have multiple reranker partitions and use different hidden state embeddings (e.g., fct,R1 and fct,R2) for different partitions.

heading: 3.1.4 Hybrid Approach (CPR)
text: Local embeddings in the pointer networks and global embeddings in the context partition are complementary. Using local embeddings is representational powerful while using global embedding can leverage the global similarity of words. Hence, we can combine the two methods by summing their dot products.
For the methods that use different S, we can simply determine an order of computing the dot products and let the later dot products overwrite the existing values. In our experiments, we always use the order illustrated in Figure 3. That is, we compute the logits (LogitCPR(x, ct)) by
   fTct,Cwx + f T ct,PD fx,ct,LD if x ∈ ct fTct,R1wx if x ∈ W (k1)− ct fTct,R2wx if x ∈ W (k2)−W (k1)− ct fTct,V wx O/W , (5)
where W (k2) is the top k2 words with the highest fTct,V wx and W (k1) is the top k1 words with the highest max(fTct,V wx,f T ct,R2 wx).

heading: 3.2 T5 and BART
text: In the encoder-decoder architectures, our local decoder embedding, context partition, and reranker partitions are still applicable. Besides, we can leverage the words in the encoder input to further improve the performance.

heading: 3.2.1 Encoder Partition (E) and Local Encoder Embedding (P)
text: Similar to the context partition, the encoder partition handles the words in the encoder input I differently by setting S = {x|x ∈ I} and using the global word embedding ex = wx.
As in Equation 4, we can also let the hidden states in the last layer pass through another linear layer LfLE() to predict the embeddings of the words in the encoder input. The method is called local encoder (LE) embedding.

heading: 3.2.2 Hybrid Approach (CEPR)
text: Similar to GPT-2, we combine local encoder embedding and encoder partition for computing the probabilities of the words that are in the encoder context but not in the decoder context. As shown in Figure 3, we compute LogitCEPR(x, ct) by
   fTct,Cwx + f T ct,PD fx,ct,LD if x ∈ ct fTct,Ewx + f T ct,PE
fx,I,LE if x ∈ I − ct fTct,R1wx if x ∈ W (k1)− ct − I fTct,V wx O/W , (6)
which is the same as Equation 5 except that we add the encoder partition and local encoder embedding, and we remove the second reranker partition.

heading: 4 Experiments
text: The pointer network was a popular technique in language modeling (Merity et al., 2017) and summarization (See et al., 2017). Thus, we also focus on these two fundamental applications.

heading: 4.1 GPT-2
text: We follow the setup in Chang and McCallum (2022) to continue training GPT-2 on Wikipedia 2021 and OpenWebText (Radford et al., 2019).

heading: 4.1.1 Perplexity Comparison
text: In Table 2, we first compare their predictions on the next word distribution using the testing data perplexity, which is a standard metric in the LM architecture studies. In the table, Mi refers to multiple input state enhancement, which is proposed to break the softmax bottleneck more effectively (please see details in Section 2.3 and Chang and McCallum (2022)).
As we can see, Softmax + CPR:20,100 + Mi, which combines all the efficient approaches (i.e., context partition, reranker partition, and local decoder embedding), results in better performance
and faster inference speed than the mixture of softmax (MoS) (Yang et al., 2018; Chang and McCallum, 2022). The inference speed is measured by our pure PyTorch implementation, which we believe could be further accelerated by implementing some new PyTorch operations using CUDA code.
If only using one method, the context partition (Softmax + C + Mi) is better than the reranker partitions (Softmax + R:20,100 + Mi) while performing similarly compared to local decoder word embedding (Softmax + P + Mi), Pointer Generator (PG + Mi) (See et al., 2017), and Pointer Sentinel (PS + Mi) (Merity et al., 2017).2 Their similar performances indicate that the improvement of pointer networks come from breaking the softmax bottleneck. The significantly better performance of PS + Mi compared to PS further supports the finding.
To know how well our method breaks the softmax bottleneck, we implement a word-by-word reranker model on GPT-2, which appends the most likely 100 words to the context when predicting each next word (see Appendix C.3 for more details). In Table 3, we show that our efficient softmax alternative Softmax + CPR:20,100 + Mi achieves significantly lower perplexity. Furthermore, the word-by-word reranker is at least 10 times slower during training. Combining word-by-word reranker with our method only improves the perplexity very
2Notice that the pointer networks from the previous work were originally designed for RNN. To add them on top of the transformer based LMs and make it more comparable to our methods, we simplify their architectures a little. Please see Appendix C.2 for more details.
slightly, which suggests the challenges of further improving LM by breaking softmax bottleneck.

heading: 4.1.2 Generated Text Comparison
text: Next, we would like to understand how the distribution improvement affects the text generation. We sample some contexts in the test set of Wikipedia 2021 and compare the generated text quality of the different models given the contexts. The quality is measured by the ROUGE-1 F1 scores between the generated text and the actual continuation. To know how much the different models copy from the context, we also report the ROUGE-1 scores between the generation and the contexts.
The results in Table 4 show that different meth-
ods have very similar overall ROUGE-1 scores. Nevertheless, compared to Softmax + Mi, Softmax + CPR:20,100 + Mi is 21% more likely to copy the proper nouns (i.e., entity names) from the context and 9% more likely to generate the proper nouns in the actual continuation. This suggests that our method could alleviate the common incoherence problem of entities in generated text (Shuster et al., 2022; Papalampidi et al., 2022; Zhang et al., 2022; Guan et al., 2022; Goyal et al., 2022b). In Table 8, we compare methods using more metrics to further support the conclusion.

heading: 4.1.3 Qualitative Analysis
text: In Table 5, we visualize some distributions to explain our improvements. The softmax layer of GPT2 is unable to properly learn to copy or exclude the word from the input context. For example, Softmax + Mi and MoS + Mi might output “There are plates, keys, scissors, toys, and balloons in front of me, and I pick up the phone”, which causes a hallucination problem, while Softmax + CPR:20,100 + Mi and Pointer Sentinel (PS) + Mi can output the mentioned options with similar probability by copying the words in the context. In addition, GPT-2, MoS, and PS + Mi are very likely to output “I like tennis, baseball, golf, basketball, and tennis”. This repetition problem happens because the next word should be some words similar to the listed sports names except for the sports that have been mentioned and the softmax layer has difficulties in outputting a donut-shape next word distribution in embedding space. In contrast, Softmax + CPR:20,100 + Mi can learn to exclude the listed sports by putting very negative logits on the context words, which yield the desired donut-shape distribution.

heading: 4.2 T5 and BART in Summarization
text: We test our methods on two popular encoderdecoder LMs, T5 (Raffel et al., 2020) and BART (Lewis et al., 2020). We fine-tune the pretrained LMs with different softmax alternatives on two news summarization datasets: CNN/DM (See et al., 2017) and XSUM (Narayan et al., 2018), one narrative summarization dataset: BookSum at paragraph level (Kryściński et al., 2021), and one dialogue summarization dataset: SAMSUM (Gliwa et al., 2019).
In the main paper, we evaluate the quality of summaries using four metrics. ROUGE-1 F1 (Lin, 2004) measures the unigram overlapping between the generated summary and the ground truth sum-
mary; CIDEr (Vedantam et al., 2015) adds a tfidf weighting on the n-gram overlapping score to emphasize correct prediction of rare phrases; factCC (Kryscinski et al., 2020) evaluates the factuality of the summary; MAUVE (Pillutla et al., 2021) compares the word distribution of summary and ground truth in a quantized embedding space. To further support our conclusions, we also compare the quality measured by several other metrics and their model sizes in Table 9 and Table 10.
The results are reported in Table 6. Similar to the GPT-2 experiments, the results are generally better as we combine more partitions and local embedding approaches. This demonstrates that we can directly fine-tune the LMs with our softmax alternatives without expensive pretraining.
Unlike the GPT-2 experiments, multiple input hidden state enhancement (Mi) is not very effective, so we mainly compare the methods without Mi (i.e., qct = h M ct , unlike Equation 2). We hypothesize one possible reason is that we haven’t pretrained the T5 and BART with our softmax alternatives.
Our improvements are larger in smaller models. This is probably because in a smaller word embedding space, there are more likely to be interfering words between the desired next word possibilities. Compared to our methods, the pointer networks perform well in BART-base but usually perform worse in other LMs. We need further investigations in the future to explore the reasons.
Compared to ROUGE-1 score, the improvement percentage of CIDEr is overall higher. One major problem of the summarization LMs is that the generated summary contains too many commonly used phrases (King et al., 2022) and our considerably higher CIDEr scores indicate the alleviation of the problem. Our improvement on the factCC is also significant (Cao and Wang, 2021). Finally, our MAUVE improvement percentage on BookSum Paragraph dataset could reach around 30% in T5-Small. We hypothesize this is because we often mention the global entity names in the news (e.g., Obama) while the meaning of names in stories (e.g., John) is often defined by the context.

heading: 5 Related Work
text: Repetition and hallucination are two common problems in language generation tasks. One common solution for repetition is to avoid outputting the words in the context, which is often called unlikelihood training (Welleck et al., 2020; Jiang et al.,
2022b; Su et al., 2022). However, when LM should mention some names in the context, this might exacerbate the hallucination problem. In contrast, our method can learn to copy and exclude the words in context as in Table 5.
To alleviate the hallucination problem or satisfy some constraints, many recent generation models rerank the generated text (Deng et al., 2020; Gabriel et al., 2021; Cobbe et al., 2021; Ravaut et al., 2022; Krishna et al., 2022; Glass et al., 2022; An et al., 2022; Arora et al., 2022; Adolphs et al., 2022; Meng et al., 2022; Mireshghallah et al., 2022; Kumar et al., 2022; Wan and Bansal, 2022; Jiang et al., 2022a). Although being effective, the
rerankers usually slow down significantly the training and/or inference speed (as our word-by-word reranker baseline) and might occupy extra memory resources.
Our analyses demonstrate that parts of the hallucination and repetition problem come from the softmax bottleneck. The findings provide an explanation for the effectiveness of prior studies such as the above reranker approaches and pointer networks (Li et al., 2021; Zhong et al., 2022; Ma et al., 2023). Another example is encouraging the word embeddings to be isotropy (Wang et al., 2020; Su et al., 2022). Their improvement might also come from reducing linear dependency of the candidate
word embeddings. Nevertheless, their side effect of breaking the similarity structure in the word embedding space might hurt the generation quality in some cases. Concurrently to our work, Wan et al. (2023) also use the softmax bottleneck theory (Chang and McCallum, 2022) to explain the improvement of a pointer network. Their empirical results also support our conclusion that softmax bottleneck is a major reason that causes the factuality problem of LMs.
Our work is motivated and inspired by Chang and McCallum (2022). In their work, they also propose to use different hidden states for different vocabulary partitions, but their partitioning is global and needs to be combined with the mixture of softmax (MoS) approach, which adds a significant overhead compared to the standard softmax layer. Our dynamic partitioning methods not only perform better but greatly reduce the overhead by removing the reliance on MoS.

heading: 6 Conclusion
text: Since the transformer becomes the mainstream encoder and decoder for LMs, the output softmax layer seems to be the only reasonable option for computing the word probability distribution. Although being simple and efficient, the softmax layer is inherently limited while the existing solutions are relatively slow (Chang and McCallum, 2022). This work proposes a series of softmax alternatives that can improve the text generation models without increasing the computational costs significantly. Our experiments suggest that the main improvement of the pointer network on top of a transformer comes from breaking the softmax bottleneck. Our results also indicate that the alternatives could alleviate some problems of hallucination, repetition, and too generic generation. Furthermore, all of the proposed alternatives can be applied to the LMs that have already been pretrained using softmax without requiring retraining from scratch. For the practitioner, we recommend using all the partitioning methods together to get the best performance, or using only the simple context partition to keep the architecture simple while getting the majority of the gain.

heading: 7 Acknowledgement
text: We thank Nadar Akoury and the anonymous reviewers for their constructive feedback. This work was supported in part by the Center for Data Sci-
ence and the Center for Intelligent Information Retrieval, in part by the Chan Zuckerberg Initiative under the project Scientific Knowledge Base Construction, in part by the IBM Research AI through the AI Horizons Network, in part using high performance computing equipment obtained under a grant from the Collaborative R&D Fund managed by the Massachusetts Technology Collaborative, and in part by the National Science Foundation (NSF) grant numbers IIS-1922090 and IIS1763618. Any opinions, findings, conclusions, or recommendations expressed in this material are those of the authors and do not necessarily reflect those of the sponsor.

heading: 8 Limitations
text: In our experiments, we find that the improvement of our methods tend to be larger in relatively smaller language models. Due to our limited access of computational resources, we are not able to try our methods on larger LMs. To know if a larger LM still suffers from the softmax bottleneck problem, we input the examples we used in Table 5 to GPT-3.5 and report their results in Figure 4.
We find that although GPT-3.5 greatly reduces the chance of hallucination compared to GPT-2, the next word distribution is still not ideal. For example, in Figure 4a, although the incorrect answer queen receives only a small probability, GPT-3.5 puts around 67% probability on woman. Similarly, even though GPT-3.5 is unlikely to hallucinate the sentence: There are plates, keys, scissors, toys, and balloons in front of me, and I pick up the phone as GPT-2, Figure 4b and Figure 4d show that the output distribution is still heavily biased toward one of the options and the most likely next word could change if the order of the options in the context changes. These results suggest that increasing model size indeed alleviates the softmax bottleneck problem but the problem is not completely solved even if a huge hidden state size (12k) and model size (175B) are used (Brown et al., 2020). We expect that adding our methods to the large LMs could rectify the biased distributions as shown in our experiments on smaller LMs (Table 5). Therefore, although improving smaller LMs has already had wide applications in practice, trying our methods on a larger LM is a promising next step, which we haven’t been able to do.
The current implementation of our methods also has some room for improvements. Our codes cur-
rently contain some unnecessary computation to circumvent the restrictions of PyTorch library, so we should be able to further accelerate it by writing CUDA code. Furthermore, our codes haven’t supported the pretraining of BART or T5. We expect that completing the future work could make our method faster and better.
Since the focus of this paper is improving the architecture of general transformer decoder, our evaluation of each application is not as comprehensive as the studies for a particular application. For example, although we test our methods using many metrics and the metrics show a consistent trend, there are many other factuality metrics we haven’t tried (Li et al., 2022). We also haven’t conducted human evaluation to further verify our conclusion because conducting human evaluation properly is challenging (Karpinska et al., 2021) and time-consuming. In addition, if we include more words in a context partition, the performance might be better at the cost of extra computational overhead. We leave the analyses of the tradeoff as future work.

heading: 9 Ethics Statement
text: In our experiments, we find that our methods usually copy more words from the context or encoder input. The tendency might have some potential issues. For example, our improvements might be reduced on the languages with more morphology. Furthermore, in some summarization applications, increasing the factuality by increasing the extractiveness might not be ideal (Ladhak et al., 2022; Goyal et al., 2022a).
As described in Section 2.1, one major limitation of the popular softmax layer is its global word
embeddings. The problem would become more serious when there are more tokens whose meanings are locally defined (e.g., names in the BookSum dataset). Our methods would be more useful in those circumstances and might alleviate some biases described in Shwartz et al. (2020) and Ladhak et al. (2023). Moreover, the meaning of tokens are also locally defined in many other applications such as variables in code or math problems, the new terminologies in a scientific paper, or the products in a sequential recommendation problem. We believe that our methods could become an efficient alternative of reranker (Cobbe et al., 2021; Welleck et al., 2022) and create impacts in those areas.
Finally, our results show that when there are some uncertainties in the next word (e.g., could be king or woman), existing LMs could have some difficulties of copying the words from the context and our methods alleviate the problem. Thus, our methods should also be able to improve the lexically controllable language generation models that put the desired keywords into the context such as Goldfarb-Tarrant et al. (2019) and Lu et al. (2021).

heading: A Appendix Overview
text: In the appendix, we first analyze our methods using more metrics in Appendix B and describe what we learn from the results. Next, we provide some details of our methods and baselines in Appendix C. Finally, we specify some experiment setups and hyperparameters in Appendix D.

heading: B More Results and Analysis
text: In this section, we will report more results and provide more detailed analyses accordingly to investigate the advantages of different methods.
B.1 GPT-2 Experiments
Kaplan et al. (2020); Henighan et al. (2020) demonstrate that the loss decreases linearly as the log of the model size increases. Therefore, a new architecture needs to perform better than the old architecture with a similar model size to verify that the improvement does not come from memorizing more information through the extra parameters. From the loss versus log(model size) curve in Figure 5, we can see that our proposed methods are significantly better than MoS and slightly better than a pointer network baseline as the model becomes larger.
We use the following metrics to measure the text generated by GPT-2.
• ROUGE-1 (R1): The prediction F1 for unigram in the actual continuation.
• ROUGE-1 Context (R1C): The prediction F1 for unigram in the context.
• ROUGE-1 Proper (R1P): The same as ROUGE-1 except that only the proper nouns are considered. We measure this metric because the correctness of the entity name prediction is critical to the factuality of the generation.
• ROUGE-1 Proper Context (R1PC): The same as ROUGE-1 Context (R1C) except that only the proper nouns are considered.
• ROUGE-2 (R2): The prediction F1 for bigram in the actual continuation.
• Proper Noun Ratio (P Ratio): The average number of proper nouns in the generation divided by the average number of proper nouns in the actual continuation. The LMs usually generate fewer proper nouns compared to the actual continuation (See et al., 2019), so the values are usually lower than 1. The P Ratio closer to 1 is better.
• CIDEr (Vedantam et al., 2015): A metric for measuring the quality and specificity of the generation.
• NIST (Doddington, 2002): Similar to CIDEr. CIDEr uses tf-idf to weigh the n-gram while NIST measures the information gain.
The results are reported in Table 8. In terms of R1, R2, CIDEr, and NIST, our proposed methods such as Softmax + C + Mi and Softmax + CPR:20,100 + Mi are significantly better than the pointer network baselines PS + Mi and PG + Mi. Comparing with Softmax + CPR:20,100 + Mi, PS + Mi has a significantly higher P Ratio and R1PC but similar R1P. This indicates that PS + Mi copies
more proper nouns from the context while there is a similar number of proper nouns that are in actual continuation, so Softmax + CPR:20,100 + Mi actually has a higher accuracy on the proper noun prediction.
In text corpus such as Wikipedia, we do not know the ground truth next word distribution and which context leads to multiple probable next words, so we cannot quantitatively analyze the improvement on the ambiguous contexts. To alleviate the concern, we test our methods on the synthetic dataset constructed by Chang and McCallum (2022). The dataset is built using templates and Google analogy dataset (Mikolov et al., 2013), so we know the ground truth next word distribution. The dataset consists of the ambiguous contexts such as I went to Paris and Germany before, and I love one of the places more, which is, where the next word is either the diagonal words of the parallelogram such as Paris and Germany or the edge words such as Paris and France. For the details of the experimental setup, please refer to Chang and McCallum (2022).
In Table 7, we can see that Softmax +
CPR:20,100 + Mi achieves the lowest perplexity in all subsets and outperforms the Softmax + Mi baseline by a large margin, especially in the diagonal subset where the ground truth word embedding distribution has multiple modes. Notice that the performance of MoS + Mi is worse than what reported in Chang and McCallum (2022) probably because we shared the input and output word embeddings.
B.2 Summarization
Compared to Figure 5, Figure 6 shows that our methods improve the loss of T5 in CNN/DM more than GPT-2 in Wikipedia.
In Table 9 and Table 10, we compare the different summarization models by their model size, evaluation losses, inference time, and other metrics which we use in subsection B.1. The pointer network baselines and our methods significantly improve most metrics over the softmax baseline, which is used ubiquitously in nearly all LMs. Although our method generally improves less on the T5-Base model, the percentages of additional parameters and inference time overhead are much
smaller. Although our methods tend to improve less in larger language model, we still improve BART Large very significantly in NIST, CIDEr, and MAUVE, and Mi seems to become more effective in BART Large.
The testing set of SAMSUM dataset only has 819 samples, so some metrics such as R1 and R2 are not as stable as other three datasets. PG (See et al., 2017) for T5-Small and T5-Base perform much worse in SAMSUM dataset. We hypothesize that it is because the dialog input in SAMSUM dataset is very different from the pretraining data of T5, which makes training PG unstable.
In most datasets and models, the R Ratio from our method is significantly closer to 1 than the softmax baseline, which means the average number of
proper nouns in our summaries is much closer to the average number of proper nouns in the humanwritten summary. For example, in BookSum Paragraph, we improve its R Ratio by 26%, which partially explains our large MAUVE improvement in Table 6. Notice that our methods do not always output more proper nouns. For example, for BART Base in CNN/DM dataset, our methods reduce the R Ratio of the softmax baseline, which is larger than 1. This shows that our methods could learn when we should copy the proper nouns according to the training data.

heading: C Method Details
text: We describe some details of our methods and baselines in this section.
C.1 Proposed Methods To allow us to start from existing LMs that are pretrained using softmax, we keep the modified softmax layer initially working almost the same as the original softmax layer. We initialize the linear transformation weights of LfPD(), L f LD(), L f PE(), and LfLE() as 10 −10 · I. The other linear weights Lf. () are initialized as the identity matrix I. In the local decoder embedding method Softmax + P + Mi, the initialization would give the 0 logit to all context words. To solve the issue, we revise Equation 3 a little and compute LogitP (x, ct) by
{ fTct,V wx + f T ct,PD
fx,ct,LD if x ∈ ct fTct,V wx O/W . (7)
That is, we initially rely on the original softmax layer to compute all the logits and let the term
fTct,PDfx,ct,LD gradually influences the logits of the context words.
In MoS + CPR:20,100 + Mi, our proposed method only revises the logit in one of the softmax.
C.2 Pointer Network Baselines
The pointer networks are originally designed for RNN, so we are unable to use exactly the same formula proposed in the papers. Nevertheless, we try our best to adapt the pointer networks for the transformer encoder while keeping the gist of the formulas. In all methods, to let the results more comparable to our methods, we use fct,PE and LfLE to determine the probability of copying the words from the context, and use fTct,V wx to determine the probability of generating all the words in
the vocabulary. In CopyNet (Gu et al., 2016), we compute the probability of outputting the word x as
Prob(x|I, ct) ∝ exp ( fTct,V wx )
+
|I|∑
j=1
1Ij=xexp ( fTct,PEL f LE(h M Ij ) + b ) . (8)
Notice that CopyNet needs to sum up the exponential of dot products, which often causes overflow problems in GPT-2. We can set b to be a large negative value initially to solve the problem, but its perplexity is much worse than the other two pointer network variants. Thus, we choose to skip the CopyNet in the GPT-2 experiments.
In Pointer Generator (See et al., 2017), we compute the probability of x using
Prob(x|I, ct) = pgen exp
( fTct,V wx )
ZV
+(1− pgen) |I|∑
j=1
1Ij=xPE(j|I, ct), (9)
where PE(j|I, ct) = exp ( vT tanh(fct,PE+L f LE(h M Ij )+b) )
ZE , pgen =
σ(qThMct + bptr), the normalization term ZV = ∑ x∈V exp ( fTct,V wx ) , and ZE = |I|∑ j=1 exp ( vT tanh(fct,PE + L f LE(h M Ij ) + b) ) .
We skip the coverage mechanism in the pointer generator paper to make it more comparable to other methods. In T5 experiments, its training loss is sometimes very large, so we set bptr as 3 initially to keep the pgen close to 1 (i.e., turn pointer part off initially). In other experiments, we set bptr = 0.
In Pointer Sentinel (Merity et al., 2017), the probability of x is computed by
Prob(x|I, ct) = g exp
( fTct,V wx )
ZV
+
|I|∑
j=1
1Ij=x
exp ( fTct,PE tanh(L f LE(h M Ij )) + b )
Zp , (10)
g = exp(qThMct )
Zp , and Zp = exp(qThMct ) +∑|I|
j=1 exp ( fTct,PE tanh(L f LE(h M Ij )) + b ) .
In our experiments, we find that the pointer network variants usually have similar performance (except that PG sometimes performs much worse
in summarization due to some training stability issues). This suggests that the differences in the pointer network variants often do not influence the performance significantly, which justifies our simplification of the formulas in the original paper and supports our conclusion that the improvement comes from breaking the softmax bottleneck.
Notice that in the above pointer network variants, the pointer part can only increase the probability of the context words from the generator part. As a result, it cannot alleviate the repetition problem in the last example of Table 5.
C.3 Word-by-word Reranker Baseline
We illustrate our word-by-word reranker (wbwR) in Figure 7. The method has two stages. In the first stage, we compute the logits using the projected hidden state fct,V and retrieve the top k words. At the second stage, we append the top k words to the input context along with the hidden state fct,R for reranking the context words.3 We use the same positional embeddings for all candidates to encourage the model to change the ranking of the words. Next, we use the hidden states corresponding to the candidates to compute their local word embeddings as fx,ct,LD. Finally, we re-estimate the probabilities of top k words by
{ fTct,V wx + f T ct,R
fx,ct,LD if x ∈ W (k) fTct,V wx O/W . (11)
To improve the quality of our top k candidates, the final loss is the addition of the wbwR loss at the second stage and the loss of the original softmax layer that only uses the logits from fTct,V wx at the first stage. When we combine the wbwR with Softmax + CPR:20,100 + Mi, we simply use Softmax + CPR:20,100 + Mi at the first stage and use the wbwR to overwrite the logits of Softmax + CPR:20,100 + Mi at the second stage.
Using this method, we can update the embeddings of the words that are not in the context and allow the candidates to interact with the input context to determine their probabilities as the classic two-stage reranker while keeping the model size roughly unchanged. Nevertheless, the method can only change the probability of the top k words and its computational overhead and memory requirement prevents us from using a very large k.
3The motivation is helping GPT-2 to output the local word embedding of a candidate closer to the fct,R if GPT-2 wants to increase the probability of the candidate.
Unlike the standard GPT-2, we cannot get the probability of all positions in one forward pass because the input contexts are different when computing the probability at each position and the input of the second stage reranker depends on the results of the previous forward pass at the first stage. To speed up, we reuse the computed hidden states and batchify the forward passes.
In our implementation, we first get the top k candidates corresponding to all tokens in the stage1 (just original GPT2) as the input of stage2 reranker. To avoid recalculating the hidden states of the context at stage2, we store the hidden states using the past-key-value in Hugging Face and only compute the hidden states corresponding to the top k candi-
date tokens at stage2. We divide the computation of the whole input sequence into several blocks as shown in Figure 8. In each block, we input a batch containing the last few tokens and top k candidates into the GPT-2 while reusing the hidden states of their common contexts from stage1. In this way, we can increase parallelism by increasing the block size if the GPU memory allows it.
Even though we spent substantial effort on optimizing the wbwR, the method is still too slow to be practically useful. Even if we use four RTX 8000 (a faster GPU with a larger memory), our wbwR implementation is still around 10 times slower than our proposed Softmax + CPR:20,100 + Mi that uses only one RTX 2080.

heading: D Experiment Details
text: For the reproducibility, we provide some experimental configuration in this section. Please see our codes (https: //github.com/iesl/Softmax-CPR) for more details.
D.1 GPT-2 Experimental Details
We mostly follow the experimental setup Chang and McCallum (2022) except that we share the input and output word embeddings as in the standard GPT-2 models. As in Chang and McCallum (2022), we use the last 2% of the corpus as the test set and the 2% before that as the validation set.4 In
4We do not shuffle the corpus before splitting the datasets. We found that our improvement could be even larger if we shuffle the corpus to let the training data distribution closer to
the word-by-word reranking experiment, we only use first 100k tokens in validation set to speed up the evaluation. To show that our model could be added to existing pretrained models, we continue training the pretrained GPT-2. For GPT-2 Small, we train for 1 epoch, and for GPT-2 Medium, we train for 0.4 epoch. We find that the performance improvements usually do not change significantly after training for 0.1 epoch.
As in Chang and McCallum (2022), we set the sequence length as 200, batch size as 4, and learning rate for AdamW as 1e− 5. Our methods only have two hyperparameters, k1 and k2, and we try values 20, 100, 200, 500 and select 20 and 100 using the validation data.
In the text completion experiment, we generate 360k continuations with a length of 50 given the prompts in Wikipedia. We first sample 40k sequences in the test data of Wikipedia 2021. Next, we use the first 20, 70, and 120 words in the sequence as our context and let the different models generate the next 50 words as continuations. The references are the actual next 50 words. All the methods use Top-K sampling and K=5.
D.2 Summarization Experimental Details
BookSum dataset (Kryściński et al., 2021) includes three summarization tasks: Summarizing a book, a chapter, and a paragraph. We test our methods using the paragraph summarization task due to the input length restriction of BART and T5. The dataset is constructed by automatically aligning the paragraphs in a chapter with the sentences in a chapter summary, which introduces noise to the dataset. Similarly, XSUM uses the first sentence in news instead of manually-written summary as the ground truth reference. The relatively noisier datasets such as XSUM and BookSum Paragraph, and smaller dataset like SAMSUM could test the stability of the methods. The sizes of the summarization datasets could be found in Table 11.
We conduct the summarization experiments
the testing data distribution.
based on a summarization example code from Hugging Face5. Most of our hyperparameters use the default value in the code. In our preliminary study, our improvement is not sensitive to the hyperparameter choice (e.g., the improvement gap is similar across different numbers of epochs). Thus, we do not tune the hyperparameters for each method or for each dataset unless we cannot reach a low training loss at the end.
In CNN/DM, XSUM, and SAMSUM datasets, We train models for 3 epochs. In BookSum datasets, We trained models for 5 epochs.6 The learning rate is set to be 5e− 05 except for BART Large model in BookSum, where we use 1e− 05 to stablize the training of all methods.
All the experiments use batch size 8 and AdamW with betas=(0.9,0.999), epsilon=1e − 06, weightdecay=1.2e − 6. During the generation, we used Top-K sampling (K=10) as our decoding method. The maximum summary length is set as 128 and maximum input length is 1024. We use warmup for the first 1000 steps in all the experiments, which allows us to change the architecture of T5 and BART more significantly (e.g., using Mi) without having a training stability issue.
The k in the reranker partition and the block size of multiple input hidden states (Mi) is coarsely tuned based on validation performance of CNN/DM. Unlike considering the top 100 words in the open-end text completion using GPT-2, we find that reranking the top 20 words is sufficient for our summarization models, probably because next words are easier to predict in the summarization task.
For our evaluation metrics, we use the default setting for ROUGE7 and set use_stemmer=True. When reporting the ROUGE scores, we follow the conventions to show their percentages. We use the default setting for MAUVE8, CIDER9, NIST10. For MAUVE, we insert a new line symbol after every sentence as in the original Hugging Face
5https://github.com/huggingface/transformers/ blob/main/examples/pytorch/summarization/run_ summarization.py
6The BookSum Paragraph is noisier, so we train longer to be safe. And we find that different numbers of epochs do not change the trend of the results.
7https://huggingface.co/spaces/ evaluate-metric/rouge
8https://huggingface.co/spaces/ evaluate-metric/mauve
9https://github.com/vrama91/cider 10https://www.nltk.org/api/nltk
summarization example code. For factCC metric11, we use the author-provided checkpoint to evaluate CNN/DM results since factCC is originally trained in CNN/DM. For the other three summarization datasets, we follow the author’s codes to constructed positive and negative data and continued training the CNN/DM factCC model on each dataset with one epoch respectively. Then we evaluate different summarization tasks with the corresponding factCC checkpoint.
D.3 Computational Environment and Software
We implement our methods by revising the Hugging Face library (Wolf et al., 2020). From Hugging Face, we load the pretrained LMs including GPT-2 Small12, GPT-2 Medium13, T5-Small14, T5Base15, BART Base16, and BART Large17. We use SpaCy (Honnibal et al., 2020) to detect the proper nouns.
For GPT-2 Medium, T5-Base, and BART Large, we use NVIDIA GeForce RTX 8000 to train the model and for other smaller models, we use NVIDIA GeForce RTX 2080. Most of experiments could be done within one week. In all the inference time experiments, we use NVIDIA GeForce GTX TITAN X, batch size 4 for GPT-2, and batch size 8 for BART and T5.
11https://github.com/salesforce/factCC 12https://huggingface.co/gpt2 13https://huggingface.co/gpt2-medium 14https://huggingface.co/t5-small 15https://huggingface.co/t5-base 16https://huggingface.co/facebook/bart-base 17https://huggingface.co/facebook/bart-large
ACL 2023 Responsible NLP Checklist

heading: A For every submission:

heading: 3 A1. Did you describe the limitations of your work?
text: section 8

heading: 3 A2. Did you discuss any potential risks of your work?
text: section 9

heading: 3 A3. Do the abstract and introduction summarize the paper’s main claims?
text: section 1

heading: 3 A4. Have you used AI writing assistants when working on this paper?
text: Grammerly
B 3 Did you use or create scientific artifacts? section 4

heading: 3 B1. Did you cite the creators of artifacts you used?
text: section 4
7 B2. Did you discuss the license or terms for use and / or distribution of any artifacts? All the datasets and language models in the papers could be used for research purpose
B3. Did you discuss if your use of existing artifact(s) was consistent with their intended use, provided that it was specified? For the artifacts you create, do you specify intended use and whether that is compatible with the original access conditions (in particular, derivatives of data accessed for research purposes should not be used outside of research contexts)? Not applicable. Left blank.
7 B4. Did you discuss the steps taken to check whether the data that was collected / used contains any information that names or uniquely identifies individual people or offensive content, and the steps taken to protect / anonymize it? Discussion of the privacy issues and offensive content in the language models is out of the scope of this paper
7 B5. Did you provide documentation of the artifacts, e.g., coverage of domains, languages, and linguistic phenomena, demographic groups represented, etc.? People can check those information in the paper that releases the datasets
3 B6. Did you report relevant statistics like the number of examples, details of train / test / dev splits, etc. for the data that you used / created? Even for commonly-used benchmark datasets, include the number of examples in train / validation / test splits, as these provide necessary context for a reader to understand experimental results. For example, small differences in accuracy on large test sets may be significant, while on small test sets they may not be. section D
C 3 Did you run computational experiments? section 4 and section B

heading: 3 C1. Did you report the number of parameters in the models used, the total computational budget
text: (e.g., GPU hours), and computing infrastructure used? Table 2, 8, and 9; section D.3
The Responsible NLP Checklist used at ACL 2023 is adopted from NAACL 2022, with the addition of a question on AI writing assistance.

heading: 3 C2. Did you discuss the experimental setup, including hyperparameter search and best-found
text: hyperparameter values? section D
7 C3. Did you report descriptive statistics about your results (e.g., error bars around results, summary statistics from sets of experiments), and is it transparent whether you are reporting the max, mean, etc. or just a single run? For GPT-2 experiments, the perplexity on a large test set is known to have a very small variance. For summarization experiments, we provide the results of lots of different metrics and most of the metrics show the consistent trend. Due to sufficiently large test set, the papers in two domains (self-supervised language modeling and summarization) often do not report their error bars.

heading: 3 C4. If you used existing packages (e.g., for preprocessing, for normalization, or for evaluation), did
text: you report the implementation, model, and parameter settings used (e.g., NLTK, Spacy, ROUGE, etc.)? section D
D 7 Did you use human annotators (e.g., crowdworkers) or research with human participants? Left blank.
D1. Did you report the full text of instructions given to participants, including e.g., screenshots, disclaimers of any risks to participants or annotators, etc.? No response.
D2. Did you report information about how you recruited (e.g., crowdsourcing platform, students) and paid participants, and discuss if such payment is adequate given the participants’ demographic (e.g., country of residence)? No response.
D3. Did you discuss whether and how consent was obtained from people whose data you’re using/curating? For example, if you collected data via crowdsourcing, did your instructions to crowdworkers explain how the data would be used? No response.
D4. Was the data collection protocol approved (or determined exempt) by an ethics review board? No response.
D5. Did you report the basic demographic and geographic characteristics of the annotator population that is the source of the data? No response.

heading: Future Work
text: We expect that completing the future work could make our method faster and better. Since the focus of this paper is improving the architecture of general transformer decoder, our evaluation of each application is not as comprehensive as the studies for a particular application. For example, although we test our methods using many metrics and the metrics show a consistent trend, there are many other factuality metrics we haven’t tried (Li et al. , 2022)