Efficient Structured Prediction with Transformer Encoders (2024)

Abstract: Finetuning is a useful method for adapting Transformer-based text encoders to new tasks but can be computationally expensive for structured prediction tasks that require tuning at the token level. Furthermore, finetuning is inherently inefficient in updating all base model parameters, which prevents parameter sharing across tasks. To address these issues, we propose a method for efficient task adaptation of frozen Transformer encoders based on the local contribution of their intermediate layers to token representations. Our adapter uses a novel attention mechanism to aggregate intermediate layers and tailor the resulting representations to a target task. Experiments on several structured prediction tasks demonstrate that our method outperforms previous approaches, retaining over 99% of the finetuning performance at a fraction of the training cost. Our proposed method offers an efficient solution for adapting frozen Transformer encoders to new tasks, improving performance and enabling parameter sharing across different tasks.

## introduction the text encoder models evolved from the transformer architecture @xcite have extensively influenced natural language processing. the standard workflow of these models is based on transfer learning from a model pre-trained on vast amounts of text to a target task. finetuning is a commonly used technique adjusting the parameters of a pre-trained encoder to a target task using the standard backpropagation algorithm. despite its simplicity and tremendous success, finetuning can be computationally expensive, particularly for structured prediction tasks in which the parameters are updated at the token level, as opposed to document classification tasks in which the parameters are updated for each document. 1 in addition, due to in-place parameter updates, finetuning limits the reusability of transformer encoders, particularly in cloud environments where resources are shared between users. additionally, finetuning a transformer encoder for a specific task or a limited number of tasks does not necessarily perform well on other tasks that are not similar to the target task. this is because of catastrophic forgetting, which reduces the generalizability of neural network performance on outof-domain data @xcite . consequently, a finetuned transformer encoder becomes a massive computational block specified for a target task, which limits the scalability, modularity, and composi-tionality of the base encoder @xcite . recent attempts to resolve the shortcomings of finetuning are based on the adapter mechanism @xcite that injects learning blocks into a frozen encoder to facilitate knowledge transfer from pre-training to target tasks @xcite @xcite . while this approach can effectively replicate the finetuning performance @xcite , it still requires careful considerations to balance the encoder's sharability and inference efficiency in a cloud environment, as it tends to sacrifice inference efficiency for sharability, and vice versa, as we will empirically show in this study. we adopt a different strategy based on the early studies of @xcite ; @xcite ; @xcite for transformer encoder adaptation. in contrast to the adapter approach @xcite , which adds trainable parameters to the encoder, this method pipes the encoder into an aggregation block that adapts the representations obtained from the encoder's intermediate layers to a target task through linear interpolation. accordingly, this approach does not necessitate changing the base encoder architecture, making it easier to be shared when compared to the adapter solution. although the layer aggregation approach is easy to implement, it requires further consideration when applied to structured prediction. the primary reason for this is that the method assumes that the layers' linear weights are a function of the target task solely and independent of input tokens. this means that a layer weight remains constant for all tokens during inference. while this assumption might be acceptable for an encoder that is finetuned for document classification, as (1) finetuning allows for training the substantial parametric capacity of the original encoder, and (2) classification often relies on a single token vector that represents the entire document, it may not hold true for a frozen model employed for structured prediction. in such cases, it is necessary to model the word dynamics based on the available intermediate representations because the model parameters remain fixed throughout training. furthermore, @xcite suggest that the linear combination of intermediate representations in a frozen model can match the performance of a finetuned model only if the pre-training and target tasks are sufficiently similar. this implies the need for additional customization of the aggregated representations to account for any disparities between the pre-training and target tasks. our paper presents an encoder adaptation model that effectively combines the efficiency of a frozen model with the effectiveness of a finetuned model by addressing the weaknesses of the linear aggregation method. our approach includes two key mechanisms: an aggregation block and a tailoring block. the aggregation block models the dynamics of words by utilizing the intermediate representations of the encoder and introduces an attention mechanism that trains token representations based on both the target task and the local contribution of intermediate layers to tokens. the tailoring block reduces the impact of dissimilarity between the pretraining and target tasks by refining the aggregated representation, thus allowing for more effective knowledge transfer from pre-training to the target task. we are not the first to propose the intermediate layer aggregation at the token level. @xcite have also explored a similar approach. they train the token-layer attention weights based on a task-specific query vector used to measure the similarity between intermediate representations locally. nonetheless, our experiments demonstrate the practical benefits of our approach in downstream tasks in structured prediction and document classification. we evaluated the effectiveness of our adaptation mechanism on two major classes of transformer encoders: bert-based @xcite and gpt-based @xcite models. our evaluation of the proposed adaptation mechanism builds upon the promising results of ablation studies of the two key blocks. our experimental results confirm that the adaptation mechanism can perform as well as finetuning while being approximately 13 times more efficient in training time and consuming only 0.3% of the memory required for finetuning, with almost no harm to the inference efficiency. compared to other approaches, our technique performs significantly better on the majority of structured prediction tasks and remains on par with the best-performing models for document classification. ## related work the initial investigation of bert showed that it captures a rich hierarchy of linguistic knowledge in its intermediate layers @xcite @xcite ; this leads to the effective use of bert by linearly aggregating the middle layers based on a target task (e.g. sentiment analysis @xcite , morphosyntactic parsing @xcite , gender debiasing and coreference resolution @xcite , and cross-lingual transfer learning @xcite ). building on this approach, @xcite expand the task-oriented layer aggregation to encompass both token and task aspects. this extension is achieved through the introduction of an attention fusion model that leverages the local features of a token. we extend @xcite 's method by introducing an attention mechanism that incorporate global views of intermediate representations. other techniques that combine intermediate representations include @xcite , who apply the squeeze and excitation technique @xcite , and @xcite , who use a bidirectional gru layer to calculate the linear weights of the intermediate representations. from an architectural point of view, our adapter mechanism is based on a dynamic aggregation of the intermediate representations and preserves the parallel encoding functionality of the original encoder. this is in contrast to the linear method of @xcite , which is based on a static weighting of the intermediate representations, as well as the methods of @xcite , @xcite and @xcite , which add overhead sequential units to the encoder, hampering parallel computing. when it comes to improving efficiency, the literature has proposed strategies such as knowledge distillation @xcite , attention pruning @xcite , model quantization @xcite , low-rank adaptation @xcite , shallow finetuning @xcite and prompt tuning @xcite . our method aligns with the adaptation techniques category, augmenting a frozen model with a few learning blocks to facilitate knowledge transfer to a target task. a related approach by @xcite injects adapter layers into a frozen bert model to enable model sharing in a cloud environment, specifically for efficient sequential multitask learning. @xcite address the catastrophic forgetting issue and balancing of different tasks in @xcite 's ap-proach. additionally, @xcite enhance the efficiency of the adaptation technique by incorporating a low-rank approximation of the model's key operations. in connection to this work, @xcite introduce an efficient method centered around the low-rank factorization of the expected changes in attention matrices in a transformer encoder. our approach diverges from that of @xcite , @xcite , and @xcite in several ways. in terms of model integration, we envelop the adapter around a pre-trained encoder model, unlike the strategies employed by @xcite al. (2019), pfeiffer et al. (2021), and hu et al. (2022) , who embed the adapter blocks within the original encoder architecture. this distinction allows us to exclude the encoder during training, resulting in a substantial enhancement of training efficiency and facilitating model sharing. in addition, the number of trainable parameters in the models of @xcite and @xcite scales with the number of intermediate layers of the original encoder, as opposed to our model in which the number of trainable parameters is almost independent of the number of middle layers. moreover, our encoder adaptation mechanism uses a dedicated tailoring block to explicitly address the differences between the pretraining and target tasks. ## encoder adaptation an encoder b consisting of 𝑙 -1 intermediate transformer layers plus one input embedding layer transforms an input document 𝑠 = (𝑡 1 , . . . , 𝑡 𝑛 ) into a threedimensional tensor: an adapter function takes an encoding tensor 𝐵 and merges its layer views (i.e., ) into a matrix: the resulting matrix includes distilled information of the intermediate representations adapted to a target task. we propose a parametric adaptation function consisting of two blocks: an aggregation block that merges the proposed adapter architecture. the input tensor 𝐵 contains the intermediate representations taken from a transformer-based encoder for a sequence of tokens, and the output is a token embedding matrix. the layer views and a tailoring block that adjusts the aggregated representations to a target task. the final adapted output is constructed by a residual connection, enabling the model to control the tailoring influence on the aggregation. the architecture is shown in figure 1 . ## experiments we study the adapter performance on downstream tasks and investigate the contribution of the aggregation and tailoring blocks to the performance gain. the experiments are based on the cased versions of bert-large @xcite and roberta-large @xcite ) models as representatives of encoder-only models, and different variants of the gpt2, including gpt2-small, medium, and large, as representatives of decoder-only models. all models are provided by huggingface. the test benchmark includes the two major types of classification tasks in nlp, including structured prediction and document classification. the evaluation benchmark for structured prediction includes tasks defined on the following datasets: ## performance this section summarizes the results collected from the structured prediction and document classification tasks. we compare our adapter performance, named adapted, with other techniques, namely: we freeze the encoder parameters in all of the above techniques except for the finetune, in which all encoder parameters are updated during training. the outputs taken from these encoders are subsequently employed for the target tasks outlined in section 4. the following sections provide detailed insights into the modules employed for each task. ## ablation study this section studies the importance of the aggregation and tailoring blocks to our adapter architecture. figure 4 represents the improvements made by each block over the baseline frozen models. a significant part of the improvement in the structured prediction tasks comes from the aggregation that accounts for 85% of the average improvements over frozen bert models. however, both blocks contribute almost equally to the average improvement in document classification. this indicates the importance of the token-wised layer aggregation for structured prediction tasks that search for the interconnections between tokens. on the other hand, the tailoring block contributes significantly to the document classification tasks whose objectives differ from the encoder's pretraining objective (i.e., masked token prediction). we also see that the necessity for tailoring in the structured prediction becomes more evident as the task complexity increases. the results suggest that the required information for syntactic tasks such as pos tagging, chunking, and parsing is already available in the intermediate representations and we only need to aggregate the information properly. however, more complex tasks that rely on both syntactic and semantic inference (e.g., deprel, ner, and re) can benefit from both aggregation and tailoring blocks. next, we study the importance of the residual connection between the aggregation and tailoring blocks. ## conclusion we have introduced a task adaptation mechanism for transformer encoders to address the reusability and efficiency issues of finetuning in structured prediction. the proposed model aggregates the intermediate representations of a frozen encoder based on the input tokens and tailors them to a target task. empirical results confirmed that the adaptation mechanism improves the training efficiency significantly while being on par with the finetuning performance. further ablation studies confirmed the importance of both the aggregation and tailoring blocks. in future work, we want to study the adapter performance within different languages and analyze attention weights in more detail across different tasks in multilingual benchmarks.