前缀树约束大语言模型解码

随着计算资源和大规模数据集的发展，大语言模型因其庞大的参数量和训练数据，在各项任务中表现极其优异，但是模型在面临代码生成、结构化数据输出等场景时会存在非法输出，这就需要对模型解码过程进行有效干预。

1.背景

随着计算资源和大规模数据集的发展，大语言模型因其庞大的参数量和训练数据，在各项任务中表现极其优异，但是模型在面临代码生成、结构化数据输出等场景时会存在非法输出，这就需要对模型解码过程进行有效干预。

2.前缀树

前缀树（字典树）是一种树结构，每一个节点代表一个字符串前缀，每个节点可以连接多个节点（N叉树形式），子节点代表字符串是由当前节点和通往该节点路径的所有字符组成的。
那么前缀树的构建、插入、查找就成为了较为重要的前置基础，因此下文将会对前缀树的各项操作进行说明。

注：为了巩固练习，可以做下leetcode的这道题目：前缀树练习.

2.1前缀树的构建

对于树的初始化，首先需要给定一个根节点，前缀树的根可以用一个空的哈希表来初始化。

def __init__(self):self.root = {}

2.2前缀树的插入

对于前缀树的插入，需要从根节点开始遍历，若找到前缀树中存在对应字符，则继续在当前节点的孩子节点进行查找，直到找不到对应字符；若前缀树中未出现对应字符，则在当前节点挂载该字符，并未其创建一个空哈希表，已被挂载其孩子节点。

def insert(self, word):cur = self.rootfor w in word:if w not in cur:cur[w] = {}cur = cur[w]cur['#'] = True

2.3前缀树的查找

对于前缀树的查找，本质与插入的思想一致，从根节点开始遍历，若未找到相应字符返回False，若找到相应字符则在其孩子节点中继续查找剩余字符。

def startsWith(self, prefix):cur = self.rootfor w in prefix:if w not in cur:return Falsecur = cur[w]return True

3.前缀树约束解码

有了前缀树的基础之后，那么我们可以来学习如何用前缀树来约束LLM解码。
LLM在解码最后阶段会计算词表中所有词的logit，然后进行softmax之后选择分数最高的词作为下一个token的输出（贪心解码策略），因此我们可以拿到模型计算的logit进行修改，将输出范围外的token设置为很大的负数，再返回给模型进行解码。

3.1约束解码逻辑

①设置last_token，从这个词之后要对解码过程进行约束，并找到这个词在序列中的索引。
②判断last_token之后的词是否符合前缀树中的路径。
③若合法，则获取最后一个词下面挂载的孩子节点，对孩子节点以外的字符进行屏蔽；若不合法，则终止序列生成。
④将logit传回LLM。

import itertools
import torch
from transformers import LogitsProcessorclass TrieLogitsProcessor(LogitsProcessor):def __init__(self, trie, tokenizer, num_beams, last_token=':'):self.trie = trieself.num_beams = num_beams # 候选序列的个数self.last_token = tokenizer(last_token)['input_ids'][1:][0] # 从这个词之后，要对解码进行约束self.eos_token_id = tokenizer.eos_token_id # 结束标志def __call__(self, input_ids, scores):mask = torch.ones_like(scores, dtype=torch.bool)for i, input_id in enumerate(input_ids):sequence = input_id.tolist()# 找到last_token在序列中的索引if self.last_token in sequence:index_from_end = len(sequence) - 1 - sequence[::-1].index(self.last_token)sub_sequence = sequence[index_from_end + 1:]cur = self.trie# 判断序列是否合法for s in sub_sequence:if s in cur.children:cur = cur.children[s]else:cur = Nonebreakif cur:next_states = cur.childrenelse: next_states = {self.eos_token_id}scores[i, list(next_states)] = -1e10valid_token_ids = list(next_states)# 对不合法的token进行屏蔽mask[i, valid_token_ids] = Falsescores[i] = scores[i].masked_fill(mask[i], -float(1e12))if len(next_states) < self.num_beams:complementary_ids = list(i for i in range(self.num_beams - len(next_states)))scores[i, list(complementary_ids)] = -1e10return scores

3.2约束解码实现

最后在model.generate中可以接收一个参数logits_processor，可以将上述约束解码策略传入。

outputs = model.generate(input_ids,logits_processor=logits_processor,max_length=50,num_beams=num_beams,num_return_sequences=num_beams,output_scores=True,return_dict_in_generate=True,early_stopping=True)

4.总结

通过构建一个包含所有有效输出模式的前缀树，我们在解码阶段动态地屏蔽那些不符合预定规则的token，从而显著提升模型在垂直领域或特定任务中的输出准确性和可靠性，减少不合法输出。