论文

无需重训的跨方言泛化:MLIR模式派生受限解码的基准与评估

Cross-Dialect Generalization Without Retraining: Benchmarks and Evaluation of Schema-Derived Constrained Decoding for MLIR

模型推理解码与生成控制

摘要

多级中间表示(MLIR)是现代ML编译器基础设施(TensorFlow、JAX/StableHLO、PyTorch Inductor、IREE)的基础,但在代码LM预训练语料中只有痕量存在。MLIR还按设计可扩展:新方言随应用领域发布,因此每方言微调一个模型不可扩展。我们要问:从每个方言的操作定义规范(ODS)机械派生的推理时先验能否替代基于梯度的适配。第一,我们发布四个自然语言到MLIR的基准——MLIR-Spec-150、Linalg-Spec-30、StableHLO-Spec-30与StableHLO-Held-Out-200,跨三个方言共410个范围内的NL到MLIR对,外加25程序的语法外语料压力集与手工编写的n=30功能参考集,以Apache-2.0发布并附Gebru数据表与Croissant 1.0元数据。第二,我们构建三层模式派生约束栈:操作签名上的CFG(C1)、从ODS抽取的类型格导出的类型域拆分(C2)、驱动五次重试拒绝采样的SSA作用域验证器(C3)。从arith+func+memref+linalg移植到StableHLO无需新约束层代码。在验证器语义由结构约束主导的方言上,模式派生先验让SmolLM2-1.7B以8–25倍的每代生成速度匹敌或超过15B–34B代码LM:在linalg上,SmolLM2达80.0%验证通过(三种子均值,n=125),以21–44个百分点、不重叠置信区间胜过CodeLlama-34B、Granite-Code-34B与StarCoder2-15B。在arith+func与模板化参数化的StableHLO-Held-Out-200上——其验证器语义取决于属性值而非结构——同样的基线匹敌或超过SLM;我们把这些划为非赢格。基准、解码器、全部逐提示生成与可复现Docker镜像均已发布。

无需重训的跨方言泛化:MLIR模式派生受限解码的基准与评估:论文配图
图 1:C3C_{3} 范围验证器混淆矩阵与 200 SmolLM2 + C1C_{1} 代上的 mlir-opt 的比较。 SmolLM2 输出上的(范围 FAIL,验证 PASS)单元为零。在 34B 模型分布 (CodeLlama) 上,相同的单元格不为零,请参阅附录 A“零错误拒绝是特定于分布的”以及表 3 中的 −22.2-22.2pp CodeLlama C3C_{3} 下降。