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 cosupervision 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 model1.

heading: 1 Introduction
text: 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 (Touvron et al., 2023a; Brown et al., 2020; Zhang et al., 2022). 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 (Min et al., 2022; Yang et al., 2017; Pappas et al., 2020). Also, such supervision constrains the model’s predictive capabilities to only the next to-
∗Work performed during internship at KAIST. † Denotes equal contribution
1Code 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 (Li and Liang, 2021; Lester et al., 2021; Hao et al., 2023), 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, Embseq, 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 (Petroni et al., 2021) and ALCE (Gao et al., 2023) 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 improve-
3864
ment of 14.2. This demonstrates that constructing a common space through co-supervision fosters the generalization and robustness of the language model. The nonparametric space under semiparametric token-sequence co-supervision is more stable compared to models trained solely on NSP, suggesting that the robustness of the parametric space, established through pretraining, provides a solid foundation that enhances the stability of the nonparametric space. Also, unlike models trained only via NTP, the model trained via semiparametric token-sequence co-supervision tends to effectively use knowledge from the nonparametric space during generation, suggesting a shift from rote learning to active knowledge utilization.
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 Embseq established during a pretraining step influences the overall performance where aligning Embseq and Gen to the same pretrained LM thereby the same distribution contributes to a more stable training process.

heading: 2 Related Works
text: Aligning two different models Various studies have explored ways to align two different models. In multi-modal tasks, efforts have been made to connect pretrained vision-only and language-only models, either by employing cross-attention mechanisms (Alayrac et al., 2022) or by aligning the vision encoder’s output embedding with the language model’s input space (Liu et al., 2023). Also, aligning a language model with another multilingual language model enables the model to perform multilingual tasks by leveraging the distribution from the multilingual model (Bansal et al., 2024; Yoon et al., 2024). Moreover, recent research has focused on retrieval-augmented language models that align a retrieval model with a language model to mitigate the issue of hallucination in language models. Studies suggest that aligning these two models leads to improved performance compared to training them separately (Lin et al., 2023; Shi et al., 2023). semiparametric token-sequence cosupervision also aims to train on aligning two different models but is unique in that it focuses on aligning the two models at the output space of the language model rather than the input space.
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 (Khandelwal et al., 2019; Zhong et al., 2022). Also, leveraging the rich contextual knowledge through nonparametric embeddings effectively reduces instances of hallucination and improves the generation of accurate and factual content (Lewis et al., 2020; Borgeaud et al., 2022; Guu et al., 2020). Moreover, it shows high performance for rare and unseen cases as such tend to not exist in model parametric space (Lee et al., 2023b; Min et al., 2022). 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.

heading: 3 Semiparametric Token-Sequence Co-Supervision
text: 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.

heading: 3.1 Revisiting Next Token Prediction
text: We revisit the conventional approach of Next Token Prediction (NTP), which forms the foundation of most modern language models (LMs). NTP is a process to predict token t over the vocabulary set V when given the preceding tokens t1, . . . , tk to a language model:
argmaxt∈V P (t|t1, . . . , tk) (1)
In more detail, as shown in Figure 2, a language model (Gen) consists of a language model (Autoregressive LM) and language model head (LM Head). LM Head WV (∈ R|V |×h) is a linear layer where h denotes the model hidden state dimension. When
given sequence of tokens t1, . . . , tk to LM, it returns a hidden state vector qk (∈ Rh). The hidden state is calculated with LM Head which returns the probability distribution over vocabulary size. Thereby Equation 1 can be reformulated as:
argmaxt∈V WV qk (2)
Equation 2 can be interpreted as a retrieving stage, indicating that the parametric token embedding space WV (LM Head) which consists of model vocab set determines the corpus (the range of objects from which "what" we will retrieve) when given qk, the hidden state embedding as a query. As the conventional language modeling paradigm only provides supervision over a fixed vocabularysized token embedding space WV , the usage was limited to predicting the most probable next token embedding. However, with such interpretation, when given multiple supervision from various embedding spaces, the methodology is extendable to predicting not only token embedding WV but also various representatives in any other non-parametric embedding spaces.

heading: 3.2 Next Sequence Prediction
text: Broadening the scope of next token prediction (NTP), we explore the domain of sequence-level embedding space WC . In natural language processing, many tasks extend beyond merely predicting the next token; they necessitate the utilization of
sequence-level knowledge such as contexts from external memory from a corpus (Zhong et al., 2022; Hao et al., 2023; Lewis et al., 2020). This is where the next sequence prediction (NSP) becomes invaluable. Unlike traditional NTP, which operates over parametric token embedding space, NSP interacts with nonparametric sequence embedding space. NSP allows models to anticipate and generate answers based on given sequences on-demand, mirroring the human-like ability to anticipate sequences of varying granularities and its ability to refer to external sequences while answering a question.
When equipped with an embedding space WC representing embedding space constructed with a corpus set of sequences C, NSP is a process to predict the sequence s akin to Equation 2.
argmaxs∈CWCqk
Specifically, as shown in Figure 2, the embedding space WC is constructed by the last hidden state embedding of another autoregressive model Embseq; WC is a nonparametric embedding space constructed with a set of last hidden state embedding s where s = Embseq(s) for s ∈ C.
There are previous studies that explore supervision beyond the token level, targeting higher levels of text granularity (Devlin et al., 2018; Joshi et al., 2020). However, their motivations are distinct from semiparametric token-sequence cosupervision; our objective of additional supervision (NSP) is to leverage co-supervision to foster a unified space that bridges the nonparametric sequence embedding space with the parametric token embedding space. In contrast, earlier efforts, such as the next sentence prediction (NSP) feature in BERT (Devlin et al., 2018), focused on a simpler binary supervision task that determines the relevance of a succeeding sentence. Also to the best of our knowledge, semiparametric token-sequence co-supervision is the first approach to train autoregressive language model with co-supervision apart from only NTP.

heading: 3.3 Co-Supervision
text: We introduce semiparametric token-sequence cosupervision, a training approach that trains a language model (Gen) by incorporating supervision from both the traditional next token prediction (NTP) which is calculated over the parametric token embedding space, and next sequence predic-
tion (NSP) which is calculated over nonparametric sequence embedding space as shown in Figure 2.
For NTP, we apply an ordinary casual language modeling loss function. When given input X:
LNTP = − 1 |X| ∑
ti∈X logPGen(ti|t<i) (3)
For NSP, we apply the contrastive InfoNCE loss (Karpukhin et al., 2020; Izacard et al., 2021); when given a query embedding, positive sequence pair relevant to the query, and a pool of negative sequences unrelated to the query, Embseq and Gen are trained to maximize the similarity between the query embedding and the positive pair and minimize the similarity with the negatives pairs where the query embedding is last hidden state embedding of Gen and sequence embeddings are set of embeddings from the last hidden state of Embseq. As it is impractical to dump all embeddings of a corpus set every step, here we approximate softmax over all corpus by softmax over positive and negative pairs in the same batch (Karpukhin et al., 2020; Izacard et al., 2021); we consider other sequences in the same batch as negatives (in-batch negatives). Formally, given a query embedding q, positive pair passage embedding c+i , and negative pairs c−1 , · · · , c−M−1 where M is the number of sequences in a batch, NSP is calculated via:
LNSP(qi, c + i , c − 1 , . . . , c − M−1)
= − log e sim(qi,c + i )
esim(qi,c + i ) + ∑M−1 j=1 e sim(qi,c − j )
(4)
Thereby total loss over semiparametric tokensequence co-supervision is calculated as:
Lco-supervision = LNTP + λLNSP (5)
where λ is the weight parameter to match the loss scale between LNSP and LNTP as it flows through Gen together.

heading: 4 Implementation Details
text: 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.

heading: 4.1 Problem Setup
text: For details, we start with an explanation of notations and training instances. Gen is the trainable language model that trains over by co-supervision from both NTP and NSP. Embseq is the trainable language model that constructs the nonparametric sequence embedding space and which calculates over Gen output embedding for NSP loss.
Training instances are in a form of “q, n1, [NSP], c1, n2, · · ·ni, [NSP], ci, · · · ” (Asai et al., 2023) (Examples in Appendix 5). q is an input query. [NSP] is a special token indicating the model to calculate over sequence embedding space, in other words when the model itself considers external knowledge is necessary. ci is a relevant (positive) sequence when given sequences before [NSP] as input. For example, c1 is the relevant sequence when given q, n1 as an input. The number of [NSP] varies from 0 to multiple differing by how many times the query necessitates external knowledge during generation.

heading: 4.2 Training
text: Figure 2 shows the overview of how we train a language model Gen with semiparametric tokensequence co-supervision, where both LNTP (equation 3) and LNSP (equation 4) flows through a language model Gen simultaneously. For LNSP loss, we train another language model Embseq together which constructs a nonparametric sequence embedding space. Specifically, the query embedding (q) and sequence embeddings (c) to calculate LNSP are calculated by:
q = Gen (query[NSP]) [−1] (6) c = Embseq (
sequence
) [−1] (7)
which is the last layer token representation of [NSP] from Gen and the last layer token representation of end-of-sequence token (
) from Embseq, respectively. Thereby the gradient of LNSP flows through both Embseq and Gen. Whereas for LNTP, the gradient only flows through Gen.

heading: 4.3 Inference
text: During the inference step, we first dump all sequence embeddings with Embseq during the offline time. Given a set of sequences of size M , we feed each sequence into Embseq and extract representative embedding c from which we get a context embedding matrix of C ∈ Rh×M .
After dumping, Embseq is no longer necessary and we only need Gen. The generated response is in the same form as the training instances “q, n1, [NSP], c1, n2, · · ·ni, [NSP], ci, · · · ” (Section 4.1). When Gen generates [NSP], it treats the last layer representation of [NSP], generated after inputting sequences up to [NSP] (q, · · · , ni), as the query embedding q. This embedding q is then calculated over the context embedding matrix C to find the most relevant sequence (ci) for the query. ci is added after [NSP], allowing the generation process to proceed based on the sequence. When Gen generates a token other than [NSP], it functions the same as a standard language model, selecting the next token based on the highest probability through the language modeling head.

heading: 5 Experiments Setup
text: 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 Embseq (Eq 7) thereby flowing the gradient through both models, the baseline computes it between the output embeddings of Embseq only, where c is the same and q is different:
q = Embseq (query[NSP]) [−1] c = Embseq (
sequence
) [−1]
which limits the gradient of NSP to only Embseq.
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 (Table 6 in the Appendix). Retrieval performance measures whether the model finds relevant paragraphs to answer the given question which requires well-constructed nonparametric sequence space. Grounding (Gr) performance evaluates how well the model generates based on given external knowledge. Following the approach of previous works (Gao et al., 2023; Lee et al., 2023a), we use TRUE (Honovich et al., 2022), a T5-11B model finetuned on various NLI datasets, to see whether the external knowledge entails a part of the response that is generated based on the knowledge.
Training Dataset All models undergo training using the dataset provided by Asai et al. (2023), 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 Asai et al. (2023) 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 (Petroni et al., 2021) and ALCE (Gao et al., 2023). 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 (Touvron et al., 2023b) as an initial model for both Gen and Embseq. 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 (Loshchilov and Hutter, 2019) 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.

heading: 6 Experimental Results & Analysis
text: Table 1 and Table 2 show the overall performance on the KILT and ALCE benchmark, respectively.2. Both tables show that models trained with semiparametric token-sequence co-supervision consistently outperform models trained under each type of supervision separately. This suggests that this co-supervision encourages a broader generalization capability throughout the model. In this section, we delve into the impact of each type of supervision on the model’s performance and explore the effects of co-supervision. From now, for simplicity we name the model trained with semiparametric token-sequence co-supervision as NTP + NSP and the model trained under each supervision separately as NTP3. More details of each paragraphs are in Appendix C.
Co-supervision makes Embseq more robust Training a language model with both supervisions from LNTP and LNSP consistently shows higher retrieval performance (average of +16.6) over those trained only with LNSP. Especially, the performance gap tends to increase as corpus size increases and over out-of-domain (improvement rate of 35.04 for in-domain and 41.67 for out-of-domain). Such results suggest that the co-supervision makes Embseq, a language model that constructs the nonparametric sequence space through its output embeddings, more robust. We could further see that it shows more robust performance over other Embseq which is trained via NSP with more datasets and larger batch size (Appendix C.1). As previous research has shown that new embeddings (tokens) can easily adapt to wellestablished parametric token embedding space of the model (Hao et al., 2023; Schick et al., 2023), we
2ELI5 in the ALCE benchmark is a reformulated version from the KILT benchmark where it contains a smaller number of instances and the evaluation metric is different. For more details on reformulation, please refer to Gao et al. (2023)
3Please note that NTP is not only trained with NTP but also NSP but separately. We name it NTP for simplicity.
hypothesize such adaptability applies to nonparametric sequence embedding space; the robustness of the parametric token embedding space, established during the pretraining step provides a solid foundation that enhances the stability of the nonparametric space.
Co-supervision enhances generation generalization Models trained with semiparametric tokensequence co-supervision consistently outperform those trained under single supervision in terms of correctness and grounding performance across both KILT and ALCE benchmarks. Such results suggest that co-supervision enhances generation generalization of Gen. The performance gap is particularly noticeable in datasets such as FEVER, T-REx, and zsRE. Such results suggest that models under co-supervision demonstrate higher generalization ability showing a broader understanding across diverse input distributions, whereas models trained solely with next token prediction (NTP) struggle
with unique input formats. For example, T-REx and zsRE are slot-filling tasks, which present a unique input format “subject [SEP] relationship type” which is less likely to be encountered during standard training phases. Additionally, in FEVER, the issue of low grounding performance arises from a poor understanding of instruction. Despite being instructed to generate evidence-based responses, the model tends to generate simple answers without grounding to external knowledge.
Co-supervision encourages high interaction between the two spaces Figure 3 (Numbers are in Table 7 of Appendix) shows the degradation rate of correctness when removing the ones that are correct by the parametric knowledge (categorizing responses associated with incorrect paragraphs as wrong). Models trained with semiparametric token-sequence co-supervision show less degradation compared to the one with separate supervision, highlighting that co-supervising encourages interaction between token and sequence embed-
ding space as it is trained to share a common space. Also, it suggests that the model leverages the contextual knowledge within the nonparametric space for generating responses, as supported by previous research (Min et al., 2022; Lee et al., 2023b), rather than depending exclusively on its parametric knowledge base. Even with the restriction of Gen to the parametric space during inference, thereby excluding access to the nonparametric sequence embedding space, the model trained with co-supervision experiences a substantial decline in correctness (from 45.7 to 32.8) in contrast to the model solely supervised with NTP (from 32.8 to 32.0). This suggests that semiparametric token-sequence cosupervision promotes the utilization of knowledge from the nonparametric sequence space rather than relying solely on memorization, effectively leveraging information from its parametric space.
Flowing loss over the sequences when calculating LNTP in co-supervision tends to make the model memorize the knowledge rather than utilizing the knowledge from nonparametric embedding space When we do not mask sequences when calculating LNTP of token-sequence co-supervision, Gen tend to rely more on their memorized parametric knowledge whereas models trained with the sequences masked tend to utilize the external knowledge retrieved from the nonparametric sequence embedding space. Such a tendency aligns with the findings of Mallen et al. (2022), where the model tends to depend on retrieved knowledge more when it lacks familiarity with the information (long-tail knowledge). By calculating LNTP over the sequences, the model tends to encode the context knowledge in its parameters, leading to a reduced reliance on retrieved knowledge and more on its own knowledge.
How does the choice of Embseq affect performance? We investigate how using different pretrained models for Embseq impacts overall performance as the nonparametric sequence embedding space is constructed through the output embeddings of Embseq. We compare results over three different pretrained models for Embseq: GPT2-large (Radford et al., 2019), TinyLlama (Zhang et al., 2024), and Llama2-7B. As shown in Figure 4, Llama27B shows the highest performance. The result suggests that the specific distribution inherent to each pretrained language model influences the performance. As semiparametric token-sequence co-
supervision trains a model in a multi-task manner thereby constructing a common space of parametric token space and nonparametric sequence space through co-supervision, the training process appears to be most stable when Gen and Embseq are derived from the same model, thus sharing the same distribution. This observation underlines the significance of distribution compatibility between Gen and Embseq in enhancing model training and performance.
Affect of weight lambda (λ) when training semiparametric token-sequence co-supervision We investigate how different weights ({10−1, 10−2, 10−3}) between the next token prediction supervision (LNTP) and next sequence prediction supervision (LNSP) affect the model’s performance, which is the lambda weight in Equation 5. In our setup, as a single generation model Gen receives cosupervision from both the parametric token embed-
ding space and the nonparametric sequence embedding space, the weight determines the balance of influence between these two spaces on the model’s training. Figure 5 (numbers in Table 14 in the Appendix) illustrates that a weight of 10−3 results in poor retrieval performance, indicating challenges in grasping and stabilizing the nonparametric sequence embedding space. Moreover, at a weight of 10−1, the model’s generation ability tends to decline, suggesting that it rather ruins the wellformed parametric token embedding space. This analysis underscores the critical role of balancing supervision from both embedding spaces to optimize model performance across various metrics.

heading: 7 Conclusion
text: 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.

heading: 8 Limitation
text: Due to resource constraints, our experimentation did not extend to altering Gen with other pretrained LMs such as Mistral (Jiang et al., 2023). While we have identified several indicators suggesting that training Gen and Embseq through semiparametric token-sequence co-supervision enhances robustness compared to training each model separately, further exploration with increased computational resources would provide a more comprehensive understanding.

heading: Acknowledgments
text: We thank Hanseok Oh, Jiyeon Kim, Miyoung Ko, Sewon Min, and Patrick Lewis for helpful discussions and constructive feedback. This work was partly supported by Kakao Brain grant (2023, Aligning Language Model with
Knowledge Module, 60%) and Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korea government (MSIT) (No.2022-0-00264, Comprehensive Video Understanding and Generation with Knowledge-based Deep Logic Neural Network, 20%; No.2019-0-00075, Artificial Intelligence Graduate School Program (KAIST), 20%).

heading: A.1 Gathering in-batch negatives
text: Suppose we have B batched instances D1, . . . DB for each training step, and each Di has ci,1, . . . ci,li reference context(li ≥ 1). We can collect M =∑B
i=1 li context embeddings(after counting duplicates) and the same amount of query embedding, each referring to one positive target. We set the target context embedding as the positive one, while setting the rest M − 1 as the in-batch negative samples.
During the experiments, we gather all the context embeddings across all 8 GPUs to increase the number of in-batch negatives. We ensure at least 63 4, but the exact number differs across training steps (experimentally about 80 to 90).
4at least 1 context per instance, 8 per each GPU, a total of 8 GPUs infer 64 total reference context embeddings, yielding at least 64− 1 = 63 in-batch negatives.

heading: A.2 Generation by Grounding in Inference Step
text: To distinguish where the model generates by grounding on the retrieved sequence and where the model generates in a freeform, we add two special tokens [CS] and [CE] which each indicate the start of the generation by grounding on the retrieved sequence and the end, respectively. In other words, when the model generates a response in the form of “q, n1, [NSP], c1, g1, · · ·ni, [NSP], ci, gi, · · · ”, the part between [CS] and [CE] is the gi part and the rest is ni indicating freeform generation.

heading: B Experiments Setup

heading: B.1 Details of metrics
text: We assess the generation results in three axes: correctness, retrieval performance, and grounding performance
Correctness Correctness evaluates how well the model answers the given query for each task. For each dataset, we chose the metric to evaluate following the metric used in its official paper. Details for each dataset is in Table 6.
Retrieval Performance Retrieval performance measures whether the model retrieves relevant paragraphs to answer the given question. We measure in two aspects, whether the gold paragraph exists within retrieved paragraphs (Ret) and retrieval precision to assess how many of the model-retrieved paragraphs contain gold paragraphs (Ret-P). For example, when the model retrieves three different paragraphs {P1, P2, P3} while generating a response and only one of them {P2} is the gold paragraph, Ret will be 100 since there is a gold paragraph in the set of retrieved paragraphs while Ret-P is 13 since only one is correct. Please note that we measure the metric by considering both gold and retrieved as a set; when the same paragraph is retrieved twice, we consider it as one during the calculation.
Grounding Performance Grounding performance evaluates how well the model generates based on given external knowledge. Following the approach of previous works (Gao et al., 2023; Lee et al., 2023a), we use TRUE (Honovich et al., 2022), a T5-11B model finetuned on various NLI datasets, to see whether the external knowledge entails a part of the response that is generated based on the knowledge. To be more specific, as we
add special tokens [CS] and [CE] to distinguish between natural form generation and generation based on given external knowledge (grounding generation); [CS] indicates the starting point of grounding generation, and [CE] indicates the endpoint. (Appendix A.2) Thereby we parse the generated response into pairs (external knowledge, grounding generation) and calculate the average of whether the external knowledge entails the grounding generation in sentence wise of grounding generation.

heading: B.2 Details of training dataset
text: Filtering We apply 3 filtering from Asai et al. (2023) to make our training dataset.
1. Relevance and Support: We retained instances with relevant and fully or partially supported citations to ensure dataset integrity.
2. Citation Count: Instances without citations or with more than 3 citations were filtered out to balance complexity and informativeness.
3. Length Constraint: We removed instances exceeding 1024 tokens (input + output) to optimize computational efficiency.
These criteria yielded a final training dataset of 42932 instances, poised to provide high-quality examples for effective model training.
Statistics Among 42.9k training instances, 32.1k instances contain 1 [NSP], 6.7k instances contain 2 [NSP], and 4k instances contain 3 [NSP] in the output. The source of these instances is shown in Table 3.
Single [NSP] Example Question: "which river has the largest basin in world"
Answer: "[NSP] Amazon River and Crenarchaeota. The Amazon has over 1,100 tributaries, 12 of which are over long. Some of the more notable ones are: Amazon River The Amazon River (, ; Spanish and ) in South America is the largest river by discharge volume of water in the world, and by some definitions it is the longest. The headwaters of the Apurímac River on Nevado Mismi had been considered for nearly a century as the Amazon’s most distant source, until a 2014 study found it to be the Cordillera Rumi Cruz at the headwaters of the Mantaro River in Peru. The Mantaro and Question: which river has the largest basin in world [CS] Amazon [CE] "
Multi [NSP] Example Question: "Why do corporations lobby for laws that impoverish lower class citizens?"
Answer: "[NSP] Corporate capitalism power and influence corporations and large business interest groups have over government policy, including the policies of regulatory agencies and influencing political campaigns (see corporate welfare). Many social scientists have criticized corporations for failing to act in the interests of the people, and their existence seems to circumvent the principles of democracy, which assumes equal power relations between individuals in a society. In an April 29, 1938 message to the Congress, Franklin D. Roosevelt warned that the growth of private power could lead to fascism: Dwight D. Eisenhower criticized the notion of the confluence of corporate power and d̈e factoQ̈uestion: Why do corporations lobby for laws that impoverish lower class citizens? [CS] While it is true that corporations may lobby for laws or regulations that are favorable to their own interests, it is not necessarily the case that these laws or regulations would have negative impacts on lower class citizens. [CE] [NSP] Anti-corporate activism control the markets, that income inequality is increasing, and that corporations have gained too much power. As part of the political left, activists against corporate power and influence support a decreased income gap and improved economical equity. The defenders of corporations such as Ron Arnold highlight that governments do legislate in ways that restrict the actions of corporations (see Sarbanes-Oxley Act) and that lawbreaking companies and executives are routinely caught and punished, usually in the form
of monetary fines. In addition, from the perspective of business ethics it might be argued that chief executives are not inherently more evil than Question: Why do corporations lobby for laws that impoverish lower class citizens? [CS] For example, a corporation may lobby for laws that reduce regulation on its industry, which could potentially lead to lower costs and higher profits for the corporation, but could also have negative consequences for workers or consumers. [CE] It is important to recognize that the relationship between corporations, lobbying, and public policy is complex, and it is not always clear how specific laws or regulations will impact different groups of people.In general, it is important for citizens to stay informed about the activities of corporations and to advocate for policies that benefit the common good.",

heading: B.3 Details of evaluation dataset
text: Construction Step Following the evaluation setup of Gao et al. (2023), we retrieve the top 100 paragraphs from the Wikipedia corpus provided by KILT for each instance in the dataset. We retrieve paragraphs by utilizing a well-performing retriever model, contriever-msmarco (Izacard et al., 2021). When constructing the top 100 paragraphs, we initially populate the corpus set with gold annotations from the KILT benchmark and subsequently supplement the remainder with paragraphs retrieved from contriever-msmarco, ensuring that all gold paragraphs are in the top 100 paragraphs.
Dataset Statistics and Characteristics In Table 5, we present the statistics and characteristics of the datasets in KILT (Petroni et al., 2021) benchmark employed for evaluation. Datasets lacking evidence annotation, such as WNED-CWEB, WNEDWIKI, and AIDA CoNLL-YAGO, are excluded from the KILT benchmark.

heading: B.4 Training Details
text: The default experiment setting for both semiparametric token-sequence co-supervision and the baselines is set to train the pre-trained Llama2 7B (Touvron et al., 2023b) provided from huggingface (Wolf et al., 2019) as the initial model for both Gen and Embseq. We use 8 Nvidia A100 with 80GB memory for our experiments. We set the maximum token length to be 1,024 due to the memory constraint. We use FSDP (Zhao et al., 2023) to conduct multi-GPU distributed training. We set the base hy-
perparameter as epoch 3, batch size 8, and AdamW optimizer (Loshchilov and Hutter, 2019) with no decay, and learning rate of 2e-5 and a decayed rate gamma of 0.85 every 1 epoch. For semiparametric token-sequence co-supervision, we set the weight λ as 0.01 while training. We also apply gradient clipping to the Gen model, with a maximum norm equal to 1. We run inference of our trained models using 1 A6000 GPU with 48GB memory.
B.5 Instructions for evaluation
Table 4 shows the instructions used during evaluation. For those datasets where it is used in Asai et al. (2023) or Gao et al. (2023), we follow or reformulate the instruction so that the model provides the evidence to measure the grounding performance.

heading: C Experimental Results
text: C.1 Does the benefit of semiparametric token-sequence co-supervision still hold when replacing Embseq trained via NSP to more general retrieval model?
Table 12 shows the overall Top100 performance of the KILT benchmark when replacing Embseq to other general models trained via NSP, including the one that is widely known in out-of-domain, contriever-msmarco (Izacard et al., 2021). Llama without parameter sharing tends to show the strongest performance, where we could see that Embseq trained via co-supervision (NTP + NSP) tend to show robust performance over all cases. Such results suggest that the nonparametric sequence embedding space constructed through cosupervision is well-constructed compared to those trained via only NSP.
C.2 Co-supervision encourages high interaction between the two spaces
Correctness degradation in cases where retrieval fails Table 7 shows the degradation of correctness when removing the ones that are correct by the parametric knowledge (categorizing responses associated with incorrect paragraphs as wrong). Models trained with semiparametric tokensequence co-supervision exhibit less degradation compared to those with separate supervision. This suggests that co-supervision promotes interaction between token and sequence embedding space, encouraging the model to share a common space for enhanced performance.
Grounding performance tends to vary by retrieval performance When analyzing the impact of retrieval success on grounding performance, NTP + NSP significantly outperforms in grounding when retrieval is successful compared to when it fails, whereas the model trained only by NTP shows little difference regardless of retrieval outcome. Upon investigating why models trained with NTP + NSP exhibit lower grounding performance upon retrieval failure, it appears to stem from the disconnect caused by attempting to answer the query with the incorrect document. In such cases, the model might fetch information that seems closest to the expected answer from the external knowledge but ends up generating content that also contains knowledge from the given query, leading to less relevance to the retrieved paragraph or fabricating information not present in the paragraph (Examples in Table 11 in Appendix). Conversely, models trained with NTP demonstrate consistent grounding performance, unaffected by the success or failure of retrieval. This suggests that NTP’s grounding capability is more reliant on its generative performance rather than the accuracy of retrieval, leading to similar outcomes irrespective of whether the correct information was retrieved or not. Based on these findings, future work could explore critiquing the success of retrieval based on grounding scores.
Generation performance without the condition of retrieval performance As correctness and grounding performance depend on retrieval performance, we evaluate both NTP + NSP and NTP trained models in a setting where retrieval is always correct or always wrong to see the correctness and grounding performance without the condition of retrieval performance. We evaluate two settings where 1. oracle: retrieval is always correct and 2. failure: retrieval is always wrong. Table 13 shows the average performance of each metric over datasets in the KILT benchmark for oracle setting (all numbers in Table 8) and Table 9 shows performance for failure setting. NTP + NSP trained models shows higher performance in the oracle setup whereas lower performance in the failure setup. Such results correlate with the findings on top where since semiparametric token-sequence cosupervision trained tends to get the answer correct based on the retrieved paragraphs, it shows high correctness in the oracle setup whereas lower correctness in the failure setup. Also as we could see
that semiparametric token-sequence co-supervision trained models tend to show low grounding performance when retrieval fails, we can see that it shows lower grounding performance in the failure setup whereas high grounding performance in the oracle setup.

heading: C.3 Weight parameter λ
text: Table 14 shows the Top20 performance of the KILT dataset with varying values of the weight parameter (λ) in Equation 5. Notably, for λ = 10−3, we could see that the model tends to not converge, resulting in significantly lower performance.

heading: C.4 Effect of Batch Size
text: Results in Figure 6 (Numbers in Table 10) show the average performance of each metric over 8 datasets in KILT with different batch sizes per GPU. Results show the importance of increasing the batch size in NTP + NSP; performance of all metrics tends
to increase with increasing batch size. We hypothesize such a trend largely due to stable training of NSP; as it is widely known that retrieval models tend to show higher performance with larger batch size, which is not always true for generation models (Keskar et al., 2017). As NSP is calculated via in-batch negatives, training with larger batch size in other words increasing the number of negatives consistently improves retrieval performance (Izacard et al., 2021; Karpukhin et al., 2020).
C.5 Performance gap tends to increase as corpus size increases
Figure 7 shows that the performance gap between NTP + NSP and NTP tends to increase as corpus size increases; NTP + NSP shows more stable
and robust performance with different sizes of the corpus.

heading: C.6 Comparison with self-RAG
text: To compare the performance with other models trained on the same dataset, we conducted an evaluation against self-RAG (Asai et al., 2023). From the results in Table 15, three key characteristics emerge. First, Self-RAG appears to be sensitive to thresholding. It requires thresholding to determine whether to retrieve external knowledge at the end of every sentence. In the case of the ALCE benchmark, as it requires to utilization of external knowledge (citation) for every sentence, they keep the threshold as 0, making the model cite for every sentence (self-RAG thres=0). However, for other datasets, as they do not necessitate citations for every sentence, they apply a threshold (self-RAG thres=0.2). When experimenting over both scenarios in the ALCE and the KILT benchmark, we could see that self-RAG is sensitive to threshold; it tends to struggle with retrieval unless it retrieves every sentence for the four datasets. Second, Self-RAG demonstrates comparable or superior performance on in-domain datasets (ASQA, NQ) but exhibits degradation on out-of-domain datasets (ELI5, zsRE). Last, our approach appears to be computationally efficient and performs well, particularly on large corpora. Self-RAG calculates similarity across all documents in a corpus, akin to the cross-encoder architecture, resulting in computational intensity. Additionally, our observations
indicate that the performance of self-RAG tends to decline further with larger corpus sizes, potentially due to distractions from similar documents. We skip the grounding performance as it is difficult to have a fair comparison due to the difference in distinguishing natural part generation from grounded generation.

heading: Future Work
text: Based on these findings, future work could explore critiquing the success of retrieval based on grounding scores. Generation performance without the condition of retrieval performance As correctness and grounding performance depend on retrieval performance, we evaluate both NTP + NSP and NTP trained models in a setting where retrieval is always correct or always wrong to see the correctness and grounding performance without the condition of retrieval performance. We evaluate two settings where 1. oracle: retrieval is always correct and 2. failure: retrieval is always wrong. Table 13 shows the average performance of each metric over datasets in the KILT benchmark for oracle setting (all numbers in Table 8) and Table 9 shows performance for failure setting. NTP + NSP trained models shows higher performance in the oracle setup whereas lower performance in the failure setup. Such results correlate with the findings on top where since semiparametric token-sequence cosupervision trained tends to get the answer correct based on the retrieved paragraphs, it shows high correctness in the oracle setup whereas lower correctness in the failure setup. Also as we could see
that semiparametric token-sequence co-supervision trained models tend to show low grounding performance when retrieval fails, we can see that it shows lower grounding performance in the failure setup whereas high grounding performance in the oracle setup