Seq2Seq + Prefix Tree: A New Paradigm for Retrieval Tasks (Using KgCLUE as an Example)
Two years ago, in 《万能的seq2seq：基于seq2seq的阅读理解问答》 and 《"非自回归"也不差：基于MLM的阅读理解问答》, we experimented with "Seq2Seq + prefix tree" and "MLM + prefix tree" approaches for extractive reading comprehension tasks and obtained pretty good results. Then, at ICLR 2021, Facebook's paper 《Autoregressive Entity Retrieval》 similarly used the "Seq2Seq + prefix tree" combination to achieve a "win-win" of effectiveness and efficiency in entity linking and document retrieval.
In fact, the "Seq2Seq + prefix tree" combination can in principle be applied to any retrieval-style task, making it something of a "new paradigm" for retrieval. This post revisits the "Seq2Seq + prefix tree" idea and uses it to build a baseline for the recently released KgCLUE knowledge-graph QA benchmark.
Retrieval Tasks
Retrieval tasks are surely familiar to everyone. Besides ordinary similar-question retrieval, many other NLP tasks can also be framed as retrieval, such as extractive reading comprehension, entity linking, and even knowledge-graph-based question answering.
Taking similar-question retrieval as an example, the typical retrieval pipeline goes as follows:
1. Train a sentence encoding model, which usually involves a complex negative-sample construction process — the quality of the negatives directly affects the final performance;
2. Encode every sentence into a vector and store it in a vector index such as Faiss, a step that usually consumes considerable storage;
3. Encode the query sentence into a vector, run retrieval, and return the top-k results with their similarity scores.
Now imagine someone told you there's a new retrieval scheme that requires no effort in picking negative samples, no huge-memory vector index, and yet achieves comparable accuracy and speed to the old scheme — wouldn't you be eager to try it out? Indeed, that's the protagonist of this post: "Seq2Seq + prefix tree."
Seq2Seq
Let's set aside all the fiddly details and think about what retrieval tasks are actually doing. For similar-question retrieval, we input a query and want to output the most similar sentence in the database; for entity linking, we input a sentence containing an entity and want to output the entity name or ID in the knowledge base that correctly refers to it; and so on.
Ignoring certain constraints, we can see that these tasks essentially boil down to the same abstraction:
Input one sentence, output another sentence.
Isn't that exactly what Seq2Seq is best at? So isn't it natural to use Seq2Seq for this? Of course, some readers might object: but the sentence I want to output already exists in the database — what if Seq2Seq "overflows" (decodes a sentence that isn't in the database)? This is where the prefix tree comes in: we can use it to constrain the decoding process so that the generated result is guaranteed to be in the database. We'll get into the details shortly.
Once we set aside the worry about "overflow," we find that the Seq2Seq approach really does combine many advantages:
1. Training a Seq2Seq model only requires inputs and targets, meaning we only need positive samples — no more headaches over constructing negatives. Or, put another way, every other sentence is implicitly treated as negative;
2. Seq2Seq directly decodes the target sentence, eliminating the need to store and search sentence vectors, and hence tools like Faiss;
3. Seq2Seq involves token-level interaction between the target sentence and the input sentence, which in principle allows for finer-grained comparison than inner-product-based vector retrieval, yielding better retrieval quality.
Prefix-Constrained Decoding
Now let's discuss in detail how a prefix tree is used to constrain decoding, and how it guarantees that the output lies in the database. Suppose our database contains the following sentences:
明月几时有
明天会更好
明天下雨
明天下午开会
明天下午放假
明年见
今夕是何年
今天去哪里玩
We would store these sentences using the following prefix tree:
Illustration of a prefix tree: essentially a compressed representation of the sequences
The construction simply merges identical tokens at identical positions, scanning left to right. Every complete path in the tree (starting with [BOS] and ending with [EOS]) corresponds to one sentence in the database. It's called a "prefix tree" because this tree structure lets us quickly look up all characters/sentences beginning with a given prefix. For instance, from the figure above we can see that the first character can only be "明" or "今"; "明" can only be followed by "月," "天," or "年"; "明天" can only be followed by "会" or "下"; and so on.
Once we have the prefix tree, constraining Seq2Seq decoding is straightforward. For example, since the first character can only be "明" or "今," when predicting the first character we can zero out the probability of every other character, so that the model can only choose between these two. Once the first character is fixed — say "明" — then when predicting the second character, we can likewise zero out the probability of every character other than "月," "天," or "年," forcing the model to choose among these three, giving "明月," "明天," or "明年." And so on: by zeroing out candidate tokens that are not on the prefix tree, we ensure that decoding only follows branches of the tree, and must follow them all the way to the end — guaranteeing that the decoded result is necessarily a sentence already present in the database.
Compared with a conventional vector-retrieval scheme, the "Seq2Seq + prefix tree" approach replaces the stored retrieval vectors with a prefix tree, and the prefix tree is essentially a "compressed representation" of the original sentences — so it's not hard to imagine that the storage required by a prefix tree is far less than that of dense retrieval vectors. In Python, a convenient way to implement a prefix tree is to use nested dictionaries; see the KgCLUE code below for a concrete example.
KgCLUE
Practice is the sole criterion for testing truth, so let's now implement a KgCLUE baseline using the "Seq2Seq + prefix tree" scheme, to check how well it works.
Task Overview
KgCLUE is a Chinese knowledge-graph QA benchmark recently released by the CLUE organization. Its data is well-formatted and suitable for research experiments. Specifically, it uses a knowledge base of about 20 million triples, each triple in the $(S, P, O)$ format (Subject-Predicate-Object), as shown below:
Screenshot of the KgCLUE knowledge base
It also provides a batch of annotated corpora for training. Each sample is a simple question-answer pair, where the question can essentially always be abstracted as "What is the $P$ of $S$," and the answer is the corresponding triple $(S,P,O)$, as shown below:
Screenshot of the KgCLUE annotated corpus
In principle, once $(S,P)$ is determined, we can retrieve the corresponding $O$ from the knowledge base. So our main task is essentially to parse out the correct $S$ and $P$ from the question.
The Conventional Approach
The relatively straightforward way to tackle this task is in two steps: first train a tagging model to extract $S$ from the question, then find all triples in the knowledge base corresponding to that $S$, combine $S$ with each $P$ in turn, and compute a similarity score against the question for each — so the second step is a similarity model.
One thing worth noting, though, is that the knowledge base may contain many entities that share the same name — for example, "牛郎织女" (Cowherd and Weaver Girl) could refer to the folk tale, a book, or a song, among others. To distinguish among these, the knowledge base also has a notion of "Meaning," which specifies exactly what the term refers to. In the KgCLUE knowledge base, the meaning is appended after the Subject in parentheses, e.g., "牛郎织女（中国著名民间故事）" ("Cowherd and Weaver Girl (a famous Chinese folk tale)"), "牛郎织女（2015年东方出版社出版的图书）" ("Cowherd and Weaver Girl (a book published by Oriental Press in 2015)"), and so on. But when extracting directly from a question, we typically can only recover the part outside the meaning annotation, i.e., "牛郎织女." So in practice we usually separate these out, treating each piece of knowledge as a quadruple $(S, M, P, O)$ rather than a triple.
Determining which Meaning a given Subject belongs to, based on the question, is the "entity linking" problem — a necessary step in knowledge-graph QA. But within the scope of the current KgCLUE task, since the corpus itself isn't very large, we can treat the two-step model as effectively folding entity linking into the similarity model, rather than treating it as a separate task.
This Post's Approach
If we use "Seq2Seq + prefix tree," the training side of things becomes remarkably simple.
Specifically, all we need is a single Seq2Seq model: feed the question in as the Seq2Seq input, join $(S,P,M)$ with [SEP] as the target, and train it as an ordinary Seq2Seq model. One trick here is that concatenating in the order $(S,P,M)$ works noticeably better (by 8–10 percentage points) than concatenating in the order $(S,M,P)$, because predicting $S,P$ from the question is easier than predicting $M$ — we should predict the easier part first, to reduce the number of candidate answers.
Illustration of this post's baseline model
In the reference code, the Seq2Seq model we use is the RoFormer-Sim-FT model introduced in 《SimBERTv2来了！融合检索和生成的RoFormer-Sim模型》 and 《用开源的人工标注数据来增强RoFormer-Sim》, a similar-question generation model pretrained with UniLM. In our comparisons, using RoFormer-Sim-FT rather than plain RoFormer improved accuracy by at least 2 percentage points. This suggests that similar-question generation is an effective pretraining strategy for this scheme.
Error Analysis
For decoding, we first build all the $(S,P,M)$ into a prefix tree, then decode following the prefix tree, guaranteeing that the decoded result lands on some triple in the knowledge base, producing a reasonable output. We've already covered the details of prefix-constrained decoding above, so we won't repeat them here.
However, when examining bad cases, we found the model sometimes makes surprisingly "simple" mistakes. For example, given "How far is the Pudong Shangri-La Hotel Shanghai from the train station?" the correct $(S,P)$ should be "(Pudong Shangri-La Hotel Shanghai, distance to train station)," but the model generated "(Pudong Shangri-La Hotel Shanghai, hotel star rating)." Sometimes when asking "What does XXX dislike," the model would instead generate "(XXX, likes)." Sometimes when asking "What course does XXX mainly teach," the correct answer should be "(XXX, courses taught)," but the model generated "(XXX, main achievements)." In other words, the model seems to make mistakes (generating the wrong $P$) on questions that look extremely simple on the surface.
After some thought, I believe this kind of bad case is fundamentally caused by inherent shortcomings of Seq2Seq itself, mainly in two respects: 1) the exposure bias problem during training; 2) the greediness of beam search during decoding. First, because Seq2Seq is trained with knowledge of the true previous label, this weakens the training difficulty and results in the model lacking a sufficiently "global view." Second, decoding — even with beam search — is fundamentally greedy, and it's hard for the model to take several subsequent tokens into account when predicting the current one. Take the "What course does XXX mainly teach" example: when the model generates $P$, it greedily generates the two characters "主要" ("mainly/main") first; then, constrained by the prefix tree, the only thing that can follow "主要" is "成就" ("achievements") — since "主讲课程" ("courses taught") begins with "主讲" instead — so it ends up producing "主要成就" ("main achievements").
In principle, applying some of the strategies designed to alleviate the exposure bias problem in Seq2Seq — such as those in 《Seq2Seq中Exposure Bias现象的浅析与对策》 or 《TeaForN：让Teacher Forcing更有"远见"一些》 — should help here. There are quite a few such methods, with varying complexity, so I'll leave it to readers to try them out on their own.
The "Look-Ahead" Strategy
Here I propose another strategy that can alleviate this problem — a "look-ahead" strategy.
In the "Seq2Seq + prefix tree" scheme, our Seq2Seq model isn't really generating arbitrary free text — it's performing what is essentially a retrieval operation constrained by the prefix tree. So after generating $S$, we already know the set of permissible continuations $P$, and in general there won't be too many permissible $P$ candidates. This lets us directly enumerate the permissible $P$ one by one, and adjust the prediction of the current token according to how well each $P$ covers the question.
Concretely, suppose that for the current question $Q$, we have already decoded $S$, together with the first $t-1$ characters $P_{< t}$ of $P$, and now we want to predict the $t$-th character $P_t$ of $P$. We then look up, based on $S,P_{< t}$, all possible $P^{(1)},P^{(2)},\cdots,P^{(n)}$, where each $P^{(k)}$ can be written in the form $[P_{< t},P_t^{(k)},P_{> t}^{(k)}]$. Using a coverage function — here simply the longest-common-subsequence length $\text{LCS}$ — we compute the coverage gain that each candidate $P$ would bring:
\begin{equation}\Delta^{(k)} = \text{LCS}(P^{(k)},Q) - \text{LCS}(P_{< t},Q)\end{equation}
This gain is treated 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.
In this way, for every candidate value $k$ of $P_t$, we obtain a "potential benefit" $\Delta^{(k)}$, and we can adjust the Seq2Seq prediction probabilities to boost the tokens $\Delta^{(k)}$. The boosting rule I used is:
\begin{equation}p_k \leftarrow p_k^{1/(\Delta^{(k)} + 1)}\end{equation}
That is, if the potential benefit is $\Delta^{(k)}$, we raise the corresponding probability to the power $\Delta^{(k)}+1$ and re-normalize. Since probabilities are less than 1, taking a root has an amplifying effect. This strategy brings roughly a 4-point improvement. Other boosting rules could also be tried; since this part is fairly subjective, I won't enumerate them all here.
Results
The code for this post is shared at:
GitHub: https://github.com/bojone/KgCLUE-bert4keras
The final result is:
$$\begin{array}{c|ccc} \hline & \text{F1} & \text{EM} & \text{average} \\ \hline \text{valid} & 89.20 & 91.04 & 90.12\\ \text{test} & 90.25 & 92.48 & 91.37\\ \text{online} & 86.03 & 88.45 & 87.24\\ \hline \end{array}$$
Currently ranked second on the leaderboard:
Current screenshot of the KgCLUE leaderboard
Roughly, the tuning journey went as follows:
1. Starting with RoFormer + UniLM, predicting in the order $(S,M,P)$, the validation EM was around 70-something;
2. Switching to the order $(S,P,M)$, validation EM reached 82;
3. Switching the pretrained model to RoFormer-Sim-FT bumped it up to 84–85;
4. Finally, adding the "look-ahead" strategy brought it to the current 89.
Some Shortcomings
We've now covered "Seq2Seq + prefix tree" in reasonable detail and shared a baseline built on KgCLUE as an example. Overall, "Seq2Seq + prefix tree" has clear advantages and can achieve competitive results on retrieval tasks, but there are still a few issues worth thinking about.
The most typical issue is the inherent limitation of Seq2Seq itself, which we've already discussed above. Although the "look-ahead" strategy brings a decent improvement, the underlying problem isn't fully resolved. How to address it more naturally, or design more natural decoding rules, remains an open question with no standard answer. Also, I haven't yet tried the general-purpose exposure-bias mitigation strategies mentioned earlier, so I'm not sure how well they would actually perform here.
Another issue is that, in the "Seq2Seq + prefix tree" scheme, results are "generated" by the model, so whenever a rare character gets mapped to [UNK], the generation will most likely fail — especially if the very first character of S is [UNK], in which case failure is almost guaranteed. How to better handle the [UNK] problem is worth investigating. One could also try a traditional copy mechanism — that's a matter of taste.
Lastly, while "Seq2Seq + prefix tree" may achieve good results on evaluation metrics, it has a property that isn't ideal from an engineering standpoint: fixing bad cases becomes fairly difficult. With traditional methods, fixing a bad case usually just means adding more training samples; with "Seq2Seq + prefix tree," you'd need to modify the decoding process itself, which is typically much harder.
Summary
This post introduced a new scheme for retrieval models — "Seq2Seq + prefix tree" — and presented a concrete baseline built on KgCLUE. The "Seq2Seq + prefix tree" scheme has advantages such as simple training and low storage footprint, along with some shortcomings; overall, it counts as a simple yet competitive approach.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.