Seq2Seq + Prefix Tree: A New Paradigm for Retrieval Tasks (Taking KgCLUE as an Example)

Task Introduction

KgCLUE is a Chinese knowledge graph QA evaluation task recently launched by the CLUE organization. The data is standardized and suitable for scientific research testing. Specifically, it uses a knowledge base of approximately 20 million triples, where each triple is in the format (S, P, O) (Subject-Predicate-Object), as shown below:

It also provides a set of annotated corpora for training. Each sample is in a simple question-and-answer format. Essentially, questions can be abstracted as "What is the P of S?", and the answer is the corresponding triple (S, P, O), as shown below:

In principle, as long as (S, P) is determined, the corresponding O can be extracted from the knowledge base. Therefore, our main task is to parse the correct S and P from the question.

General Approach

A simple way to complete this task is in two steps: first, train a labeling model to extract S from the question, then find all triples corresponding to that S in the knowledge base, combine S with each P, and calculate the similarity with the question one by one. Thus, the second step is a similarity model.

However, it is important to note that there may be many homonymous entities in the knowledge base. To distinguish them, the knowledge base includes the concept of "Meaning (M)" to specify the exact referent. In the KgCLUE knowledge base, meanings are added in parentheses after the Subject, such as "Cowherd and Weaver Girl (Chinese folk tale)." When extracting from a question, usually only the part outside the parentheses is extracted. Therefore, in practice, we often treat each piece of knowledge as a quadruple (S, M, P, O) rather than a triple.

Determining which Meaning a Subject belongs to based on a given question is the "Entity Linking" problem, a necessary step in KGQA. However, within the current KgCLUE task, since the corpus is not very large, for a two-step model, we can consider the entity linking step to be integrated into the similarity model rather than performing it separately.

Proposed Solution

If we use "Seq2Seq + Prefix Tree," the model training becomes very simple.

Specifically, we only need a Seq2Seq model. We treat the question as the Seq2Seq input and use (S, P, M) connected by [SEP] as the target for standard Seq2Seq training. One trick here is that concatenating in the order (S, P, M) yields much better results (8–10 percentage points higher) than the order (S, M, P). This is because predicting S and P from the question is easier than predicting M; we should predict the easier parts first to reduce the number of candidate answers.

In the reference code, the Seq2Seq model used is the RoFormer-Sim-FT model introduced in "SimBERTv2 is here! RoFormer-Sim Model Integrating Retrieval and Generation" and "Enhancing RoFormer-Sim with Open Source Manual Annotations". It is a similar-question generation model pre-trained using the UniLM framework. Comparisons show that using RoFormer-Sim-FT improves performance by at least 2 percentage points over using RoFormer directly. This also indicates that similar-question generation is an effective pre-training method for this scheme.

Error Analysis

During decoding, we first build a prefix tree for all (S, P, M) and then decode according to the prefix tree. This ensures the decoding result falls into a triple in the knowledge base. The detailed prefix decoding steps were introduced earlier.

However, when observing bad cases, I found that the model occasionally makes very "simple" mistakes. For example, for the question "How far is the Shangri-La Hotel Pudong Shanghai from the railway station?", the correct (S, P) should be "(Shangri-La Hotel Pudong Shanghai, Distance to railway station)," but the model generated "(Shangri-La Hotel Pudong Shanghai, Hotel star rating)." Sometimes for "What does XXX hate?", the model generates "(XXX, Like)." Sometimes for "What courses does XXX mainly teach?", the correct answer is "(XXX, Main courses)," but the model generates "(XXX, Main achievements)." That is, the model seems to fail on questions that look very simple literally (generating the wrong P).

After reflection, I believe these bad cases are essentially caused by inherent weaknesses of Seq2Seq, mainly involving two aspects: 1. Exposure Bias during training; 2. The greedy nature of Beam Search during decoding. First, because Seq2Seq knows the previous true label during training, it weakens the training difficulty, leading to a lack of "global vision." Second, even with Beam Search, decoding is essentially greedy, making it difficult to predict the current token by considering several subsequent tokens.

Theoretically, applying strategies to mitigate Exposure Bias in Seq2Seq, such as those in "Analysis and Countermeasures of Exposure Bias in Seq2Seq" or "TeaForN: Making Teacher Forcing More ’Farsighted’", should help. These methods are numerous and vary in complexity; I leave them for readers to try.

"Look-ahead" Strategy

Here, I propose another "look-ahead" strategy that can mitigate this problem.

In the "Seq2Seq + Prefix Tree" scheme, our Seq2Seq is not truly generating arbitrary text but performing what is essentially a retrieval operation under the constraint of a prefix tree. Therefore, after generating S, we know the allowed subsequent Ps. Generally, there aren’t too many allowed Ps, so we can directly enumerate them and adjust the current token’s prediction based on the coverage of the question by P.

The specific steps are: assume the current question Q has already decoded S and the first t-1 characters of P, denoted as P_{< t}. Now, to predict the t-th character P_t, we retrieve all possible P^{(1)}, P^{(2)}, \dots, P^{(n)} based on S and P_{< t}. Each P^{(k)} can be represented as [P_{< t}, P_t^{(k)}, P_{> t}^{(k)}]. We then use a coverage function—here, we directly use the Longest Common Subsequence (LCS) length—to calculate the coverage gain brought by each candidate P: \Delta^{(k)} = \text{LCS}(P^{(k)}, Q) - \text{LCS}(P_{< t}, Q) This gain is regarded as the "potential benefit" of predicting P_t as P_t^{(k)}. If multiple P_t^{(k)} correspond to the same character, we take the maximum.

Thus, for all candidate values k of P_t, we have calculated a "potential benefit" \Delta^{(k)}. We can adjust the Seq2Seq prediction probability to strengthen tokens with higher \Delta^{(k)}. The reinforcement rule I used is: p_k \leftarrow p_k^{1/(\Delta^{(k)} + 1)} That is, if the potential benefit is \Delta^{(k)}, the corresponding probability is raised to the power of 1/(\Delta^{(k)} + 1), and then re-normalized. This strategy brought an improvement of about 4 percentage points.

Effectiveness Analysis

The code for this article is shared at:

Github: https://github.com/bojone/KgCLUE-bert4keras

The final results are:

It currently ranks second on the leaderboard:

The debugging journey was roughly:

1. Initially used RoFormer + UniLM predicting in (S, M, P) order; validation EM was around 70+.

2. After changing to (S, P, M) order, validation EM reached 82.

3. After switching the pre-trained model to RoFormer-Sim-FT, it improved to 84–85.

4. Finally, adding the "look-ahead" strategy brought it to the current 89.