Semiparametric Token-Sequence Co-Supervision (2024)

Abstract: In this work, we introduce a semiparametric token-sequence co-supervision training method. It trains a language model by simultaneously leveraging supervision from the traditional next token prediction loss which is calculated over the parametric token embedding space and the next sequence prediction loss which is calculated over the nonparametric sequence embedding space. The nonparametric sequence embedding space is constructed by a separate language model tasked to condense an input text into a single representative embedding. Our experiments demonstrate that a model trained via both supervisions consistently surpasses models trained via each supervision independently. Analysis suggests that this co-supervision encourages a broader generalization capability across the model. Especially, the robustness of parametric token space which is established during the pretraining step tends to effectively enhance the stability of nonparametric sequence embedding space, a new space established by another language model.

## introduction language models are typically trained through next-token prediction (ntp), where the model forecasts the distribution of the next token's embedding, given a current token embedding @xcite @xcite . this process relies on a language model head, which includes embeddings for the entire vocabulary. while this approach has demonstrated high performance, its reliance on predicting over a finite parametric vocabulary space restricts the models' expressivity @xcite @xcite . also, such supervision constrains the model's predictive capabilities to only the next to- * work performed during internship at kaist. † denotes equal contribution 1 code and data in https://github.com/kaistai/ semiparametric_token-sequence_co-supervision ken, whereas humans can anticipate sequences of varying granularities highlighting a significant gap. in this work, we aim to enhance the capabilities of language models by superposing parametric token embedding space and nonparametric sequence embedding space at the output space of a language model. drawing on previous research that highlights the adaptable nature of language models' parametric token embedding space @xcite @xcite , we theorize that the language model can integrate a new embedding space alongside the model's existing parametric space and can also leverage the stable foundation of its parametric space established during the pretraining phase when integrating the new embedding space. to this end, we introduce semiparametric tokensequence co-supervision, a novel training approach that trains a language model (gen) by incorporating supervision from both the traditional next token prediction (ntp), calculated over the parametric token embedding space, and next sequence prediction (nsp), which is calculated over nonparametric sequence embedding space as in figure 1 . this nonparametric sequence embedding space is constructed by another language model, emb seq , which compresses the entire input text into a singular, representative embedding. the supervision is calculated via contrastive learning over embedding from the nonparametric embedding space and the output distribution from gen. we experiment across 10 information-seeking datasets from the kilt @xcite and alce @xcite benchmarks. we compare a model trained via semiparametric token-sequence co-supervision, which employs multi-task training over ntp and nsp, against a model trained under each supervision individually. the findings reveal that models under co-supervision significantly outperform those trained with separate supervision, with an average performance improvewith next token prediction loss (ntp), semiparametric tokensequence co-supervision trains a language model in a multitask manner where supervision from parametric token embedding space (ntp) and supervision from nonparametric sequence embedding space (nsp) flow simultaneously. notably, models trained via token-sequence cosupervision exhibit significant improvements, particularly in out-of-domain datasets, with an average enhancement rate of 6.6 over in-domain datasets. also, this co-supervision fosters a strong interaction between the parametric token and nonparametric sequence spaces, with gen more inclined to generate responses by drawing upon knowledge from the nonparametric space. moreover, the inherent distribution of emb seq established during a pretraining step influences the overall performance where aligning emb seq and gen to the same pretrained lm thereby the same distribution contributes to a more stable training process. language models with nonparametric embeddings integrating nonparametric embeddings into language models has consistently demonstrated advantages. this approach enhances the expressiveness beyond the inherent capabilities of language models @xcite . also, leveraging the rich contextual knowledge through nonparametric embeddings effectively reduces instances of hallucination and improves the generation of accurate and factual content @xcite @xcite . moreover, it shows high performance for rare and unseen cases as such tend to not exist in model parametric space @xcite . semiparametric token-sequence co-supervision also leverages the nonparametric embedding space but is unique in that it trains the model to utilize both the parametric and nonparametric embedding space, where the nonparametric embedding space is not static at the training step, but is trainable making the nonparametric space more adaptable to the well-constructed parametric embedding space. ## semiparametric token-sequence co-supervision in section 3.1, we delve into our interpretation of the next token prediction (ntp), first laying the foundation for our hypothesis. subsequently, in section 3.2, we explore next sequence prediction (nsp), an extension of next token prediction to nonparametric sequence embedding space. in section 3.3, to test our hypothesis we introduce semiparametric token-sequence co-supervision, a training method with supervision from both parametric token embedding space (ntp) and nonparametric sequence embedding space (nsp) in a simultaneous manner. ## implementation details in this section, we share implementation details of experiments over information-seeking datasets. in section 4.1, we share details of the problem setup. in section 4.2, we describe details of how we train a language model (gen) with both supervision of next token prediction (ntp) and next sequence prediction (nsp) simultaneously, and the inference step of the trained model in section 4.3. more details are in appendix a. ## experiments setup in this section, we share the experimental setup of baseline, metric, training and evaluation dataset, and training details. more details of each paragraph are in appendix b. baseline for the analysis of how co-supervision affects model performance, we train a baseline model using the same approach as outlined in section 4.2, but with each supervision (ntp and nsp) separately. ntp is computed in the same manner for both semiparametric token-sequence cosupervision and the baseline, involving the flow of gradients through gen. however, nsp is calculated differently: while semiparametric token-sequence co-supervision computes it between the output embeddings of gen (eq 6) and emb seq (eq 7) thereby flowing the gradient through both models, the baseline computes it between the output embeddings of emb seq only, where c is the same and q is different: metric to measure how training a language model with both supervision of ntp and nsp shows different aspects over the model trained with each supervision separately, we measure model performance via three metrics. correctness and grounding performance measures the generation ability of the model and how well it utilizes knowledge from the external sequence. retrieval performance mainly measures how well the model finds relevant sequences. correctness (cor) evaluates how well the model generates a response thereby answering the given query for each task. for instance, in the case of question answering (qa), correctness is measured by answer accuracy ( training dataset all models undergo training using the dataset provided by @xcite , featuring a diverse range of instruction-following datasets. the dataset contains information-seeking questions paired with long-form responses, where a set of sequences are annotated within the responses (example form in section 4.1). as the dataset includes instances beyond our scope such as cases where the matching sequence is irrelevant to the response as @xcite aims to train model self-critique, a filtering process was applied. this refined the dataset to 42,932 instances suitable for our research objectives. evaluation dataset we conduct evaluations on ten long-form information-seeking datasets from two benchmarks kilt @xcite and alce @xcite . some of these datasets overlap with the training dataset, categorized as in-domain datasets, while others are considered out-of-domain datasets. while the alce setting shows closer alignment of our training dataset, as the benchmark does not contain annotation of relevant sequences, we adapt the kilt benchmark to conform to the alce setting. for this reason, we mainly focus analysis on the kilt benchmark, which we can measure all three metrics. training details we use pretrained llama2 7b @xcite ) as an initial model for both gen and emb seq . we use 8 nvidia a100 for the experiments. we also set the base hyperpa-rameter as epoch 3, batch size 8, learning rate of 2e-5 and a decayed rate gamma of 0.85 every 1 epoch, and adamw optimizer @xcite with no decay across all the experiments. for experimenting semiparametric tokensequence co-supervision, we set the weight λ as 0.01 while training and apply gradient clipping to the gen model, with a maximum norm equal to 1. ## experimental results & analysis ## conclusion in this paper, we propose a semiparametric tokensequence co-supervision training method, which trains a single autoregressive language model with both supervision from parametric token embedding space and nonparametric sequence embedding space in a simultaneous manner. experiments over 10 information-seeking datasets show that such co-supervision consistently outperforms models trained with each supervision separately of average +14.2, demonstrating that constructing a common space through co-supervision fosters the generalization and robustness of the language model. such a method is not only limited to sequence space but can be expandable to any embedding space which we leave as future work.