Looking Right is Sometimes Right: Investigating the Capabilities of Decoder-only LLM s for Sequence Labeling (2024)

Abstract: Pre-trained language models based on masked language modeling (MLM) excel in natural language understanding (NLU) tasks. While fine-tuned MLM-based encoders consistently outperform causal language modeling decoders of comparable size, recent decoder-only large language models (LLMs) perform on par with smaller MLM-based encoders. Although their performance improves with scale, LLMs fall short of achieving state-of-the-art results in information extraction (IE) tasks, many of which are formulated as sequence labeling (SL). We hypothesize that LLMs’ poor SL performance stems from causal masking, which prevents the model from attending to tokens on the right of the current token. Yet, how exactly and to what extent LLMs’ performance on SL can be improved remains unclear. We explore techniques for improving the SL performance of open LLMs on IE tasks by applying layer-wise removal of the causal mask (CM) during LLM fine-tuning. This approach yields performance gains competitive with state-of-the-art SL models, matching or outperforming the results of CM removal from all blocks. Our findings hold for diverse SL tasks, demonstrating that open LLMs with layer-dependent CM removal outperform strong MLM-based encoders and even instruction-tuned LLMs.

## introduction pre-trained language models (plms) built upon the transformer architecture have demonstrated exceptional performance across many natural language understanding (nlu) tasks @xcite @xcite @xcite . typically, achieving state-of-the-art (sota) results in tasks such as sequence classification and sequence labeling figure 1 : layer-wise causal mask removal from decoder block groups in a decoder-only llm. here, the causal mask is removed from the top eight decoder blocks of the llama2-7b model to enable bidirectionality during fine-tuning, which proves beneficial for many sl tasks. in transfer learning experiments that juxtapose models of a comparable number of parameters, mlm-based encoders consistently outperformed clm-based decoders on nlu tasks @xcite . 2 however, a shift in strategy emerged within the nlp community when encoder models ceased being scaled up to the same magnitude of parameters and pre-training data as their decoder counterparts. consequently, there has been a pronounced trend toward scaling decoder models to multiple billion parameters, leading to a proliferation of large language models (llms). combining llm text generation capabilities with various prompting strategies can boost the performance on many nlu tasks, eliminating the need for fine-tuning model parameters @xcite . despite llms' good performance on nlu tasks, there is still considerable room for improvement, even for the largest decoder models, such as chat-gpt, which fall far behind sota results on fundamental nlp tasks. this holds particularly for information extraction (ie) tasks @xcite , such as named entity recognition (ner), aspect-based sentiment analysis (absa), and event extraction (ee). these tasks are often formulated as sequence labeling (sl), and tackling sl by prompting llms proved quite difficult @xcite . however, it is unclear how exactly and to what extent the subpar performance of llms on sl tasks can be remedied. 3 as the community stopped scaling up encoders, llms are solidifying their position as the field's de facto standard. thus, improving llms' sl performance is the sole viable option. although many sota llms are accessible only through paywalls, the community has responded by training and publicly releasing open llms with multiple billions of parameters, such as llama2 @xcite , @xcite , and opt @xcite . parameter-efficient fine-tuning (peft) techniques and quantization, including qlora @xcite , facilitate experimentation with fine-tuning these models on task-specific data. these strategies allow leveraging open llms' hidden states for task-specific classifiers, enabling fine-tuning on consumer-grade hardware with comparable or better performance than full fine-tuning. the poor performance of llms on sl tasks can be traced back to the information encoded in the hidden states. using the hidden state of the decoder's last token serves as a reliable feature for classifying the input sequence. however, the same tactic does not suffice for sl because the causal mask (cm), integral to endowing decoders with text generation capabilities, limits bidirectional information flow, preventing the model from "looking right," i.e., attending to tokens in the sequence positioned to the right of the current token. for many sl tasks, the token's label depends on the succeeding tokens, and omitting this context often results in subpar performance. recently, @xcite found that completely removing the cm from the llama2 model during fine-tuning substantially enhances sl performance on ner. others also experimented with the cm removal and enabling bidirectionality concurrently with our work @xcite . in this paper, we explore techniques to enhance the sl performance of open llms on ie tasks. building on insight from @xcite , we hypothesize that removing the cm from all decoder layers may not be beneficial for all sl tasks. to confirm this, we experiment with layer-wise cm removal across llm blocks, observing gains competitive with sota models, in contrast to removing or keeping the cm in all layers. we extend our analysis to a series of ie tasks (ner, absa, and ee) and demonstrate that layer-wise "looking right" is highly task-dependent. we compare against strong encoder-only sequence taggers and instruction-tuned llms. open llms without cms in specific layers outperform these baselines in almost all scenarios. finally, we compare with encoder-only models, assessing the relative strengths of the decoder and encoder architectures regarding parameter scale, training data, and the particular sl task. we pre-train small encoder and decoder models from scratch with an identical number of parameters on the same data. our findings indicate the superiority of the same-scale encoder architecture over the decoder on sl tasks, which is consistent even when the cm is removed while fine-tuning the small pre-trained decoders. this suggests that the observed sl performance gain in llms upon cm removal stems from scaling up. our contributions are threefold: (1) we present evidence that layer-wise cm removal, followed by supervised fine-tuning, improves the performance of decoder-only llms on sl tasks compared to removing or keeping cm in all layers; (2) layer-wise cm removal configuration that yields the highest gains strongly depends on the peculiarities of the sl task; (3) we show that the cm removal effect brings no gains on the small scale. we believe our results will contribute to building more performant decoder-only sl models. pre-training with clm teaches the model to generate text and does not allow it to use the sequence context right to the token being processed. this is enabled via a cm, which prevents the model from attending to future tokens. while beneficial and necessary for coherent text generation, restricting the model to look right is detrimental if one wants to use the decoder to produce fully contextualized token-level embeddings with the decoder. although decoder-only models perform poorly on sequence labeling tasks (i.e., tasks where each individual token is assigned a label), they perform quite well on sequence classification tasks (i.e., tasks where the entire sequence is assigned one label). therefore, many decoder models implement sequence classification heads on top of decoder blocks but not sl heads @xcite . @xcite found that removing the cm from all decoder layers of llama2 @xcite improves the performance of named entity recognition (ner). however, @xcite did not examine in depth why this phenomenon occurs, whether it occurs with layer-wise cm removal, and whether it persists for other ie tasks. we build on this and inspect the effect of different unmasking configurations across layers of open decoder-only llms. furthermore, we train small decoders with clm and small encoders with mlm to demonstrate that the superiority of unmasked llms is a consequence of the model and training data scale. universal ie. many ie tasks, such as ner, subtasks of absa, and ee, can effectively be formulated as sl tasks. traditionally, these tasks were addressed with encoders, mapping each token's hidden state to task label space with a linear layer on top of the last encoder block @xcite @xcite . however, this approach necessitates training a separate model for each task. more recent methods aim to infer the underlying structure of multiple ie tasks simultaneously and develop universal information extraction (uie) models, employing encoder, decoder, or encoder-decoder architectures @xcite @xcite @xcite @xcite . while the practical advantages of uie are obvious, including more ie tasks would still require re-training the uie extractors. therefore, we stick to the one model per sl task strategy, relying on parameterefficient methods that yield compact and extensible models @xcite . in-context learning for ie. icl is a powerful method for scrutinizing llms, employed by prompting llms with task demonstrations. ie tasks can be tackled with icl, relying on the emergent properties of llms. nevertheless, @xcite report that the performance of llms leveraging icl lags behind the sota results of supervised ie models. they develop a model that learns to guide the prompt for improved llm performance. another notable example is gpt-ner @xcite , transforming the ner task into a generation one utilizing icl and chatgpt, while @xcite and @xcite apply prompt-based methods to structured prediction tasks such as ner and semantic role labeling. instruction tuning for ie. releasing the weights of encoder-decoder and decoder-only pretrained llms paved the way for instruction tuning (it) @xcite . this paradigm involves fine-tuning the model to address specific tasks in a supervised generative manner, providing the model with instructions for each data point. sl tasks can be adapted for compatibility with it. arguing that current prompt templates are primarily designed for sentence-level tasks and inappropriate for sl objectives, @xcite reformulate ner as a generation problem relying on improved prompt templates and it. drawing on the fusion of uie and it, @xcite reveal that multi-task it can give results comparable to bert in a supervised setting for ner. finally, @xcite find that instruction-tuned encoder-decoder plms excel at absa subtasks, leveraging the tk-instruct model @xcite . given that it is a powerful supervised paradigm for combining llms' generation abilities with labeled data, we compare sl models against it baselines. ## layer group unmasking the cm is a crucial component of clm-based decoders, preventing the model from attending to future tokens and facilitating autoregressive text generation. this constraint is enforced by defining the cm as a triangular matrix and adding this matrix to the dot product of the query and key attention matrices. the resulting sum is passed through the softmax function in the scaled dot-product attention mechanism, as introduced by @xcite . formally: here, q, k, and v are the query, key, and value attention matrices, respectively, and d k is the dimension of queries and keys. effectively, applying softmax for tokens with value -∞ in the cm results in attention scores being 0. as reported by @xcite , removing cm across all decoder layers of llama2 during finetuning increases the model's performance on the ner sl task by a large margin. this is surprising because although, during pre-training, the cm mask was in place and the model was restricted from "looking right," it learned to attend to future tokens by backpropagating on the training data over a few epochs with cm removed. building on this and to allow the llm to "sometimes look right," we select a subset of decoder blocks for which we replace all -∞ entries in the cm with zeros, effectively removing the cm (technically, we leave the positions of padding tokens to -∞, but replace all others with 0). since we group layers into groups of b blocks, we refer to these decisions as layer group unmasking configurations. considering that llama2 is made up of @xmath0, and each block can either remove or keep its cm, this gives a search space of 2 @xmath1. to reduce the search space to a manageable size, we group the n decoder blocks into @xmath2. this leaves us with 2 @xmath3. for ease of reference, we encode our unmasking configurations with binary four-digit codes ranging from 0000 (all four layer groups masked) to 1111 (all four layer groups unmasked), where each 0 and 1 denote a layer group that is masked or unmasked, respectively. unmasking configurations are interpreted left to right (the first digit pertains to the layer group closest to the model's input). ## conclusion the landscape of plms shifted from encoder to decoder dominance in various nlu tasks. while large decoder models show remarkable performance even without fine-tuning, challenges persist in ie tasks, typically formulated as sequence labeling. in a series of experiments, we showed that llms can yield performance competitive with sota on sequence labeling ie tasks when the causal mask is removed from specific decoder blocks. future work should investigate cm removal for individual layers, analyze task-specific differences in cm removal, and investigate the options for predicting the impact of cm removal without requiring extensive fine-tuning. promising alternatives to trying out all the unmasking configurations could involve leveraging dynamic programming or estimating a network's final performance without retraining, for example, using methods that rely on the fisher information, similar to the fit method by @xcite . additionally, performance estimation for cm removal from a layer could be improved using representation similarity analysis methods such as centered kernel alignment @xcite . these approaches might reduce the number of unmasking configurations that need to be evaluated to find the best one. comparisons with encoders stress the significance of architecture and scale in sl tasks. a noticeable gap exists between small-scale encoder and decoder models, where decoder models do not benefit from cm removal at a small scale.