question about the first training stage

[colinzhaoxp]

hello authors, thanks for your open-source code, and I have learned a lot.

I have a question about the first stage of training pipeline:

GLEN/src/tevatron/modeling/glen_phase1.py

Lines 88 to 108 in d165ef8

	for i in range(max_output_len_in_batch):
	decoder_attention_mask_cur = decoder_attention_mask[:, : i + 1]
	cur_out = self.hf_model(
	input_ids=input_ids,
	attention_mask=attention_mask,
	decoder_inputs_embeds=decoder_inputs_embeds,
	decoder_attention_mask=decoder_attention_mask_cur,
	encoder_outputs=encoder_outputs,
	past_key_values=past_key_values,
	output_hidden_states=True,
	use_cache=True,
	return_dict=True,
	)
	if encoder_outputs is None:
	encoder_outputs = cur_out.encoder_outputs
	past_key_values = cur_out.past_key_values
	decoder_inputs_embeds = cur_out.decoder_hidden_states[-1][
	:, -1:, :
	] # [B, 1, H]
	lm_logits.append(cur_out.logits[:, -1, :])
	sequence_output.append(cur_out.decoder_hidden_states[-1][:, -1:, :])

According formula 4 in the paper, the sequence-to-sequence cross-entropy loss should use the teacher-forcing style training, but the code shown above demonstrate the model generate docid token by token and compute loss (without teacher-forcing).

Are there any considerations for this type of training?

[skleee]

Hi, thank you for your interest in our work!

For a quick answer, both phases 1 and 2 use the document representation as input, token by token, during decoding without teacher forcing. This means we run a for-loop over each token in the sequence during generation.

As mentioned in Section 3.2 in the paper, we feed the document representation d_t, not z_t, into the decoder at each step.

This design choice was made to keep consistent the training procedures of phases 1 and 2, especially considering that the document identifiers may fluctuate in phase 2 and the model needs to remain stable under such conditions.

Please feel free to reach out if you have any further questions!

[colinzhaoxp]

thanks for your response!

I have the question about dynamic docid.
In the first stage, the docid is generated by the key word extractor (BM25), then the model will learn the relationship between doc (and query) to the docid. So the docids are fixed, such as Figure 1, the docid is olymic-games-list.
In the second stage, I guess the docid of the same document in Figure 1 is also olymic-games-list, so where is the dynamic docid ? they are same and not changed.

So my question is whether the docid will change?

the $$z^+_t$$ indicates the identifier predicted from the positive document $$d^+$$, and $$z^+_t$$ is a t-th token of $$z_t$$.

Is the $$z_t$$ equal to the one in the first stage (got by BM25)?

[skleee]

@colinzhaoxp Apologies for the delayed response. Thank you for asking such an insightful question! It gets to a very interesting part of our work.

1. Will the docid change?
To answer your question directly: Yes, the DocID does change.

You are correct that Figure 1 in the paper doesn't fully illustrate this due to space constraints. However, once the model completes phase 2 of training, it learns to adjust the initial IDs. This process is what creates the dynamic IDs.

We actually have a figure in our poster that visualizes this concept more clearly as below:

Furthermore, we provide checkpoints trained on NQ320k, MS MARCO. If you run inference using these checkpoints, you will notice that the outputs are different from those of a naive T5-base. Those are direct results of the learned dynamic document identifiers.

2. Is zt equal to the one in the first stage (got by BM25)?
There is no guarantee that they will be the same. The identifier z_t is what the model learns and predicts, so it can and will change from the initial BM25-generated ID based on the training.

I hope this clears things up. Thanks again for the great questions, and please let me know if you have any more!