Soft Alignment Objectives for Robust Adaptation of Language Generation (2023)

Abstract: Domain adaptation allows generative language models to address specific flaws caused by the domain shift of their application. However, the traditional adaptation by further training on in-domain data rapidly weakens the model’s ability to generalize to other domains, making the open-ended deployments of the adapted models prone to errors. This work introduces novel training objectives built upon a semantic similarity of the predicted tokens to the reference. Our results show that (1) avoiding the common assumption of a single correct prediction by constructing the training target from tokens’ semantic similarity can largely mitigate catastrophic forgetting of adaptation, while (2) preserving the adaptation in-domain quality, (3) with negligible additions to compute costs. In the broader context, the objectives grounded in a continuous token similarity pioneer the exploration of the middle ground between the efficient but naive exact-match token-level objectives and expressive but computationally- and resource-intensive sequential objectives.

## introduction large language models (llms) based on instances of encoder-decoder architecture @xcite provide a strong standard for generative applications of nlp, such as summarization or machine translation, mainly thanks to their outstanding ability to fluently model language. these models might face issues with adequacy of the generated text (ustaszewski, 2019) when applied in data domain(s) different from the training domain, but such errors can be partially mitigated using domain adaptation @xcite . identically to the pre-training phase, the adaptation is commonly carried out using maximum likelihood estimation (mle) objective with teacher forcing @xcite . the popularity of this approach can be rightfully attributed to its outstanding data and computing efficiency. however, model adaptation using mle notoriously comes for a price of over-specialization to the target domain, also referred to as catastrophic forgetting @xcite , characterized by a continuous decay of model performance on the inputs from the other domains than the adaptation domain. we hypothesize that catastrophic forgetting might be related to mle's naïve single-truth assumption, penalizing models' uncertainty over the possibly valid predictions, such as the synonyms. in domain adaptation, a repeated penalization of possibly valid tokens that are uncommon in the adapted domain might drive the model to unlearn the original features robust to meaning-invariant formulations. we propose to counteract the single-truth assumption of mle by constructing targets that respect mutual tokens' similarity through the alignment of output tokens to the reference (figure 1 ). consequentially, the expected target distribution is spread over the tokens that can be accurately aligned to the reference, based on the representations provided by a domain-agnostic embedding model. we find that using such objectives in domain adaptation can eliminate a major portion of model performance loss on out-of-domain (ood), caused by the adaptation techniques while reaching comparable or higher qualitative gains on the adapted domain. our main contributions are the following. (i) we present a framework for training generative language models with an alternative training signal based on token similarity provided by an arbitrary embedding model. a similar methodology can be applied for more robust training and adaptation of any language model. (ii) we introduce efficient and accurate training objectives that alleviate catastrophic forgetting of low-resource domain adaptation in nmt without losing adaptation quality. (iii) we further investigate the covariates that impact the robustness of generative llm. among others, we find that a more robust model can be obtained merely by exposing a generative model to its own predictions during the training. this paper is structured as follows. section 2 surveys and compares our work to the existing work in training and adapting robust generative llms. section 3 introduces two main objectives that we experiment with: tokenalign and seqalign. section 4 describes our experimental methodology and ablation analyses and section 5 summarizes our findings, highlighting the broader implications. ## background language generation is the modus operandi for a wide set of problems requiring an open-ended sequence of tokens as the answer. machine translation is the representative of this group that we focus on, but other tasks such as summarization @xcite , vision captioning @xcite , question answering @xcite or in-context learning @xcite are also applications of the described framework. in the commonly-used auto-regressive generation, for each pair of input and reference texts, i.e. sequences of tokens x j and y j , a language model θ(y j,i |x j , y j,1..i-1 ) is trained to generate output sequence by maximizing the probability of generating the i-th token y @xmath0, y j,1..i-1 )) matching the reference token y ji while minimizing the probability of generating other tokens of the vocabulary, conditionally to the input text x j and previous reference tokens y j,1..i-1 : this goal is implemented in the commonly-used approach that we refer to as the maximum likelihood estimation objective (mle), which minimizes a cross-entropy (ce) of the predicted distribution of θ(x j , y j,1..i-1 ) to the expected distribution, which is a one-hot encoding e ji of the true reference token y ji over the model vocabulary: while the adaptation brings benefits in modelling domain-specific terminology @xcite , or in avoiding inadequate generation artefacts such as repetitions or hallucinations @xcite , it comes at a price of generalization to other domains; the adapted models improve on the adapted domain but gradually perform worse on other domains. previous work in domain adaptation presents methods addressing the mitigation of catastrophic forgetting. @xcite enhance model robustness by mixing the pre-training and adaptation samples in continuous training, assuming that the full pre-training dataset is available, which is commonly not the case. @xcite regularize the training objective with fischer information matrix. dakwale and monz (2017) also use the regularization in training, instead based on the predictions of the original model. similarly, @xcite use the ensemble of the original and trained model in prediction. in this line, we experiment with the ensemble approach using transformers but find it underperforms other methods in low-resource adaptation. @xcite find that using parameterefficient fine-tuning methods, such as using adapters @xcite can increase the robustness of the adapted model. previous work also applied adapters in the fine-tuning of generative llms @xcite , but do not evaluate the distributional robustness of the final models; therefore, we include adapters as another baseline, but find it also struggling in lower-resource cases, due to the random initialisation of its bottleneck representations. we find that this problem can be avoided using lora @xcite , which instantiates tuned parameters as additions to attention matrices initialised close to zero values, therefore commencing the adaptation with the originally-performing model. another problem that commonly arises in the training of generative lms is referred to as exposure bias: while in the teacher-forced training, the model's i-th prediction θ(x j ) i is conditioned by the correctly-generated previous tokens from the reference y j,1..i-1 , in practice, the model conditions its predictions on its own outputs θ(x j ) 1..i-1 . we speculate that this discrepancy might be magnified under a domain shift where the model could not have learned to follow the reference closely. exposure bias was addressed by sequential objectives, such as minimum risk training (mrt) @xcite ) that optimize the model by the evaluation of complete output sequence @xcite @xcite . apart from the specifics of reinforcement learning, such as fragility to the optimization settings @xcite , these methods are also more resource-demanding as they require a sequence of predictions for a single update, limiting their applicability in low-resource adaptation. previous work of @xcite also shows that gains of sequential methods in adaptation might be similar to a random training signal. inspired by this finding, we also assess the gains and ood robustness of our methods against a random-feedback sequential baseline ( §4.3). closer to us, previous work uses alternative training signal based on comparing model hypotheses to the reference. @xcite build soft alignment between fully-generated hypotheses based on hidden states of bidirectional lstm encoder-decoder and weigh the predicted probability distribution by such alignment in the training objective. similarly, @xcite complement mle and sentencelevel objective with the objective minimizing a dot-product of the best-matching hidden representations of tokens of a hypothesis and a reference. @xcite and later @xcite introduce the matching scheme that uses the optimal transport cost @xcite of the embeddings of reference to the hypothesis as their objective loss. referenced work reaches improvements in conventional high-resource training scenarios, whereas our goal is to propose a method for training robust generative models for challenging low-resource settings. this also motivates a primary difference in the design of our methods; that is, to use domainagnostic representations for constructing training targets, instead of the model's own representations, which are subject of over-specialization in adaptation. ## soft alignment objectives this section describes the details of alignmentbased objectives foot_0 that we introduce in this work. ## experiments we evaluate the impact of the proposed training objectives in the domain adaptation experiments in machine translation, where the distributional robustness in adaptation may bring well-measurable benefits. we compare our results with the adaptation using the commonly-used mle objective ( §2), and selected parameter-efficient methods shown to mitigate forgetting. we use the novel objectives as the weighted complements of the mle objective (eq. ( 2 )), thus optimising both the objectives in parallel: ## results alignment-based objectives improve robustness; both tokenalign and seqalign consistently improve the model robustness (ood) over the mle in all the evaluated cases and on average deliver more robust models compared to all other methods. in addition, comparing tokenalign to instances of mle, we also see the advances in the adaptation quality (id), in four out of five cases where mle is able to deliver any id improvements. in ood evaluations, seqalign is slightly more robust than tokenalign, presenting a more robust, yet also technically more complex alternative. while the average results confirm our main hypothesis that circumventing mle's assumption of a single-truth prediction can improve the model's distributional robustness, we see a large variance in the performance of our methods similar to mle. the in-domain results of seqalign also dispute our assumption that self-generation of prefixes could compensate for the scarcity of natural in-domain data; seqalign's id performance on the two smallest domains is inferior to both mle instances, while it is very efficient in the higher-resource law/dgt. avoiding teacher-forcing improves robustness; a comparison of the results of srand and mle in alignment-based targets complement avoiding teacher-forcing; robustness improvements of seqalign over srand ( decontextualization does not carry a large qualitative drop; both objectives grounding their targets in decontextualized embeddings (token-align and seqalign-dec) show relatively good average results on both id and ood ( loss formulation impacts model robustness; a comparison of seqalign-dec and seqalign-ce in ## conclusion our work sets out to explore the alternatives between the efficient yet naïve mle objective and expressive but resource-demanding sequential objectives, by building the training signal from the semantic token representations. we build an alignment mechanism applicable with an arbitrary representation model and propose objectives that utilize a domain-agnostic embedding model as its target. we find that using semantically-grounded targets in adaptation persists robustness of the model much better than other methods, without compromises in in-domain performance. we additionally explore the impact of selected design choices on the robustness of generative llms in the ablation experiments. among others, we find that a major part of the model's robustness can be persisted merely by including the model's own outputs among the inputs, attributing a part of adaptation forgetting to exposure bias. future work might also build upon the qualitative assessment of the impact of decontextualization, resolving the computational overhead of applying the contextualized embeddings in dynamic contexts. we look forward to future work that will explore the potential of applying semantically-grounded objectives in a more robust and data-efficient training of llms for many other applications, including the pre-training stages. while our experiments do not evaluate such settings, we note that our methods complement the model-centric ones, including recent parameterefficient training strategies @xcite . given the encouraging results of lora (