语言学处理：k2中的FST原理、源码与代码实现

一、引言

在语音识别系统中，语言学知识的编码是一个核心问题。如何将词典（lexicon）、音素（phone）序列、语言模型（LM）等语言学约束高效地整合到声学模型的解码过程中？k2生态给出的答案是：有限状态转录机（Finite State Transducer, FST）。FST提供了一种统一的数学框架，使得我们可以将不同层面的语言学知识表示为图结构，并通过图组合（composition）操作将它们无缝融合。

本文从原理、源码和代码实现三个维度，深入剖析k2生态中语言学处理的核心机制。

二、原理篇：FST与语言学知识的编码

2.1 有限状态转录机的基本概念

FST是一种带有输入和输出标签的有限状态自动机。与仅接受或拒绝符号序列的有限状态自动机不同，FST能够将一个符号序列转录为另一个符号序列，这一特性使其天然适合表示从音素序列到词序列的映射。

一个FST由以下几个要素构成：

• • 状态：自动机所处的节点，通常有一个起始状态和若干终止状态。
• • 弧：连接状态的有向边，每条弧携带四个关键信息：源状态、目标状态、输入标签、输出标签、权重（通常表示对数概率）。
• • 符号表：将符号（如音素、词）映射为整数ID的映射表，是实现符号与数值互转的核心数据结构。

FST的核心运算包括：组合将两个FST合并为一个，使其输入输出关系等于两个FST的串联，这是构建解码图的核心操作；确定化将非确定性的FST转换为等价的确定性FST，消除歧义路径；最短路径在加权FST中寻找权重最小的路径，对应解码时的最优识别结果。

2.2 语言学知识在FST中的编码

在k2生态中，语言学知识通过三类FST进行编码：

L.fst（词典FST）实现从音素序列到词序列的映射。例如，对于词”hello”，其音素序列为”h e l l o”，L.fst中会有一条路径，输入为”h e l l o”，输出为”hello”。由于一词多音现象，一个词可能对应多条音素序列，L.fst中会为每条发音变体设置不同的弧，每条弧携带一个权重表示该发音的可能性。L.fst的构建过程涉及消歧符号的添加：当多个词的音素序列存在前缀关系或完全相同时，解码时会产生歧义，系统会在音素序列末尾添加”#1“、”#2“等消歧符号，确保每个词在FST中对应的路径具有唯一性。

**G.fst（语言模型FST）**是一个接受词序列并输出相同词序列的加权接受器，其权重表示语言模型概率。在n-gram语言模型中，G.fst的状态对应n-gram历史，弧上的权重表示在当前历史条件下生成下一个词的对数概率。G.fst的构建通常从ARPA格式的语言模型文件开始，通过工具将其转换为FST格式。

**H.fst（上下文相关音素FST）**用于将上下文相关音素映射为上下文无关音素。在端到端系统中，H.fst通常简化为CTC拓扑或RNN-T拓扑，其作用是约束声学模型输出的对齐方式。

2.3 解码图的组合：HLG的构建

有了上述三类FST，解码图的构建通过组合操作完成，其核心表达式为HLG = H ∘ L ∘ G。组合的含义是：将多个FST串联，使前一个的输出成为后一个的输入，从而形成一个从声学输出直接到词序列的复合FST。

具体来说：G接受词序列并输出相同词序列（附加语言模型权重）；L接受词序列并输出音素序列（实现词到音的映射）；H接受音素序列并输出声学单元序列。组合后的HLG是一个巨大的加权FST，它编码了所有可能的词序列、音素序列与声学单元序列之间的映射关系及其概率权重。在解码时，系统将声学模型的输出与HLG进行交集，然后寻找最短路径，即得到最优的词序列识别结果。

2.4 可微FST：k2的核心创新

传统的FST框架是离散的，无法与神经网络联合训练。k2的核心创新在于实现了可微分的FST，即支持在FST运算上进行反向传播。

这一能力的实现基于一个关键洞察：在FST的前向计算中，如果记录每个输出弧与输入弧之间的对应关系，就可以在后向传播时通过链式法则将梯度回传至输入权重。具体实现时，k2采用”自顶向下”的自动微分策略：当寻找最佳路径时，系统会记录每个输出弧对应哪些输入弧；当计算损失函数时，梯度可以沿着这些对应关系传播回神经网络的输出层。这种设计使得我们可以将语言学约束（以FST形式编码）直接融入神经网络的训练过程中。

三、源码篇：k2的底层实现机制

3.1 RaggedTensor：GPU上FST的数据结构基础

k2在GPU上高效实现FST算法的核心秘密在于RaggedTensor数据结构。

FST的弧表示本质上是一个不规则的数据结构：不同状态的出弧数量可能不同，不同样本的路径长度也可能不同。传统的张量表示无法有效处理这种不规则性。k2的RaggedTensor采用了一种类似TensorFlow RaggedTensor的设计，将数据存储为连续的内存块，同时使用索引数组标记每个序列的边界。

在源码层面，RaggedTensor是一个模板化的C++类，其设计使得所有弧的属性（源状态、目标状态、输入标签、输出标签、权重）存储为多个并行的数组，这些数组通过索引偏移量与状态关联。这种设计实现了三点关键优势：内存连续，所有数据存储在一维数组中，减少内存碎片；并行友好，在GPU上不同状态或不同样本的处理可以独立并行；算法简洁，许多FST操作可以表达为对RaggedTensor的规约操作。

k2的CUDA实现大量使用NVIDIA cub库提供的并行原语，如exclusive prefix sum（前缀和操作），用于高效地计算路径概率的累积和。代码中C和CUDA混合编写，CUDA内核通过模板实例化，C11 lambda表达式直接在算法实现中操作数据指针，使得最耗时的计算部分实现”高度可并行化”。

3.2 源码目录结构解析

k2的源码主要分布在以下目录中：

• • k2/csrc/：核心C++和CUDA实现，包含GPU导向的算法
• • k2/csrc/host/：早期CPU版本的算法实现
• • k2/python/：pybind11封装的Python绑定

3.3 k2.Fsa的核心实现

k2.Fsa对象在底层被实现为RaggedTensor，模板参数是一个表示弧的结构体。该结构体包含弧的源状态、目标状态、输入标签、输出标签和权重。Fsa对象可以表示单个FSA或一批FSA（FsaVec），这通过RaggedTensor的多轴结构实现：第一轴区分不同的FSA，第二轴是每个FSA内的状态，第三轴是每个状态的弧列表。

3.4 自动微分的实现机制

k2源码中并没有显式的autograd引用，而是采用了一种与PyTorch不同的设计哲学：自顶向下实现可微分性。以最佳路径搜索为例，当执行k2.shortest_path操作时，底层C++代码在计算最短路径的同时，会构建一个映射表，记录输出路径上的每个弧对应于输入FST中的哪些弧。在Python层面，这个操作被封装为torch.autograd.Function的子类，其backward方法利用这个映射表将梯度稀疏地传播回输入权重。

这种设计的优势在于：它避免了将复杂FST操作分解为许多细粒度的可微分原子操作，从而获得更好的数值精度和更高的计算效率。

四、代码实现篇：从lang目录到解码图

4.1 lang目录的构建流程

在Icefall和SpeechBrain的k2集成中，lang目录是语言学知识的集中存放地。一个典型的lang目录包含以下文件：tokens.txt（音素符号表）、words.txt（词符号表）、L.pt（词典FST）、L_disambig.pt（含消歧符号的词典FST）。

以下代码展示了构建lang目录的核心流程，来自Icefall的prepare_lang_bpe.py：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
def lexicon_to_fst_no_sil(    lexicon: Lexicon,    token2id: Dict[str, int],    word2id: Dict[str, int],    need_self_loops: bool = False,) -> k2.Fsa:    """将词典转换为k2格式的FST"""    loop_state = 0      # 起始状态，词从这里进入和离开    next_state = 1      # 下一个未分配的状态    arcs = []    eps = 0   # epsilon符号的ID    for word, pieces in lexicon:        cur_state = loop_state        word_id = word2id[word]        piece_ids = [token2id[i] for i in pieces]        # 为音素序列中的每个音素创建弧        for i in range(len(piece_ids) - 1):            w = word_id if i == 0 else eps            arcs.append([cur_state, next_state, piece_ids[i], w, 0])            cur_state = next_state            next_state += 1        # 最后一个音素，输出词ID        i = len(piece_ids) - 1        w = word_id if i == 0 else eps        arcs.append([cur_state, loop_state, piece_ids[i], w, 0])    # 添加终止状态    final_state = next_state    arcs.append([loop_state, final_state, -1, -1, 0])    arcs.append([final_state])    # 构建FSA    fsa = k2.Fsa.from_str(arcs_str, acceptor=False)    return fsa

4.2 GraphCompiler的实现

SpeechBrain中的GraphCompiler类负责将语言学FST编译为可供解码使用的复合图。其核心方法包括：

compile_HL()方法：编译用于解码的HL图（不含语言模型），将CTC拓扑与词典FST组合。

compile_HLG()方法：编译包含语言模型的完整解码图：

1
2
3
4
5
6
7
8
9
10
11
def compile_HLG(self, G: k2.Fsa, cache_dir: Optional[Path] = None):    """编译HLG解码图"""    # L与G组合得到LG    LG = k2.compose(L_inv, G)    # H与LG组合得到HLG    HLG = k2.compose(H, LG)    # 确定化    HLG = k2.determinize(HLG)    # 最短路径剪枝    HLG = k2.arc_sort(HLG)    return HLG

4.3 ARPA语言模型到FST的转换

语言模型的FST转换通过arpa_to_fst函数完成：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
def arpa_to_fst(    words_txt: Path,    in_arpa: Path,    out_fst: Path,    ngram_order: int = 3,):    """将ARPA格式的语言模型转换为FST"""    # 读取词符号表    word_table = k2.SymbolTable.from_file(words_txt)    # 读取ARPA语言模型    with open(in_arpa, 'r') as f:        arpa_content = f.read()    # 转换为FST格式（内部实现处理backoff弧）    fst = _arpa_to_fst(arpa_content, word_table, ngram_order)    # 保存    k2.Fsa.save(fst, out_fst)

4.4 解码时的图操作

解码过程通过get_lattice函数实现：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
@torch.no_grad()def get_lattice(    log_probs_nnet_output: torch.Tensor,  # (B, T, num_tokens)    input_lens: torch.Tensor,              # (B,)    decoder: k2.Fsa,                        # 解码图    search_beam: int = 5,    output_beam: int = 5,    min_active_states: int = 300,    max_active_states: int = 1000,    ac_scale: float = 1.0,) -> k2.Fsa:    """从解码图和神经网络输出获取解码格子"""    # 1. 将神经网络输出转换为dense FSA    dense_fsa = k2.DenseFsaVec(        log_probs_nnet_output * ac_scale,        input_lens,        subsampling_factor    )    # 2. 与解码图进行剪枝组合    lattice = k2.intersect_dense_pruned(        decoder,        dense_fsa,        search_beam=search_beam,        output_beam=output_beam,        min_active_states=min_active_states,        max_active_states=max_active_states,    )    return lattice

4.5 最佳路径解码

获得格子后，通过one_best_decoding提取最优路径：

1
2
3
4
5
6
7
def one_best_decoding(lattice: k2.Fsa, use_double_scores: bool = True) -> k2.Fsa:    """从格子中获取最佳路径"""    # 使用最短路径算法    best_path = k2.shortest_path(lattice, use_double_scores=use_double_scores)    # 去除epsilon弧    best_path = k2.remove_epsilon(best_path)    return best_path

4.6 完整解码流程示例

以下是完整的解码流程代码示例：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
# 1. 创建解码图decode = get_decoding(    {        "compose_HL_with_G": True,      # 使用语言模型        "decoding_method": "onebest",   # 最优路径解码        "lang_dir": "/path/to/lang",        "lm_dir": "/path/to/lm",        "G_arpa": "3gram.arpa",    },    graphCompiler)# 2. 获取格子lattice = get_lattice(    log_probs,           # 神经网络输出    input_lens,          # 输入长度    decode["decoding_graph"],  # 解码图    search_beam=5,    output_beam=5)# 3. 解码得到最优路径best_path = decode["decoding_method"](lattice)# 4. 转换为文本text = lattice_paths_to_text(best_path, lexicon.word_table)

4.7 二次重打分

对于更高精度的解码，可以使用rescore_with_whole_lattice进行二次重打分：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
def rescore_with_whole_lattice(    lattice: k2.Fsa,    G_with_epsilon_loops: k2.Fsa,   # 更大的语言模型（如4-gram）    lm_scale_list: List[float],) -> Dict[str, k2.Fsa]:    """使用更大的语言模型对格子进行重打分"""    results = {}    for lm_scale in lm_scale_list:        # 将格子与语言模型FST相交        rescored = k2.intersect(lattice, G_with_epsilon_loops)        # 寻找最短路径        best_path = k2.shortest_path(rescored)        results[f"lm_scale_{lm_scale}"] = best_path    return results

五、总结

从原理到源码再到代码实现，k2生态中的语言学处理呈现出一个完整的技术栈：

1. 1. 原理层：将词典、语言模型等语言学知识编码为FST，通过H∘L∘G的组合操作形成统一解码图，核心创新在于可微FST的设计使语言学约束能够参与端到端训练。
2. 2. 源码层：通过RaggedTensor实现GPU上的高效FST操作，使用C++/CUDA混合编程和cub库实现高度并行化，采用自顶向下的自动微分策略记录弧对应关系。
3. 3. 代码实现层：通过prepare_lang_bpe.py构建lang目录，GraphCompiler编译解码图，get_lattice和one_best_decoding完成解码推理。

这一技术栈的设计核心在于：将语言学约束作为可微分组件无缝集成到神经网络的训练和推理流程中。这不仅实现了传统FST框架的表达能力，更赋予了其端到端优化的能力——这正是k2作为新一代Kaldi核心组件的本质价值所在。