用规则约束替代微调，让小模型高效生成跨方言MLIR代码
Cross-Dialect Generalization Without Retraining: Benchmarks and Evaluation of Schema-Derived Constrained Decoding for MLIR
- 基于操作定义规范自动生成语法约束，实现无微调的跨方言生成
- 小模型在特定方言上验证通过率超340亿参数大模型，速度提升8-25倍
- 适合编译器研发者和低资源语言建模场景，支持快速部署
多级中间表示（MLIR）是现代机器学习编译器基础设施（如TensorFlow、JAX/StableHLO、PyTorch Inductor、IREE）的核心，但在代码大模型预训练语料中出现极少。且MLIR设计可扩展：新方言按应用领域发布，因此为每种方言微调模型不可持续。本文探索是否能在推理时通过机械方式从各方言的操作定义规范（ODS）生成先验约束，替代基于梯度的适配。首先，我们发布了四个自然语言到MLIR的基准数据集，覆盖三种方言——MLIR-Spec-150、Linalg-Spec-30、StableHLO-Spec-30与StableHLO-Held-Out-200，总计410个有效对，并附带25个超出语法的压力量测集及人工撰写的30个功能参考集，均以Apache-2.0许可发布，包含Gebru数据表与Croissant 1.0元数据。其次，构建三层由模式推导的约束栈：基于操作签名的上下文无关文法（C1）、从ODS提取的类型格中划分的类型域（C2），以及驱动五次重试拒绝采样的SSA作用域验证器（C3）。从arith+func+memref+linalg迁移到StableHLO无需新增约束层代码。在验证语义主要依赖结构约束的方言上，基于模式的先验使SmolLM2-1.7B在生成效率上达到15B-34B代码大模型的水平，速度提升8-25倍：在linalg任务上，SmolLM2平均验证通过率达80.0%（三种子实验，n=125），优于CodeLlama-34B、Granite-Code-34B和StarCoder2-15B达21-44个百分点，置信区间不重叠。在arith+func及模板化参数化的StableHLO-Held-Out-200上，因验证依赖属性值而非结构，相同基线模型表现持平或更优，被界定为非胜出区域。所有基准、解码器、每提示生成结果及可复现容器镜像均已公开。
原文摘要 · Abstract (English)
Multi-Level Intermediate Representation (MLIR) underlies modern ML compiler infrastructure (TensorFlow, JAX/StableHLO, PyTorch Inductor, IREE), yet appears only in trace amounts in code-LM pretraining corpora. MLIR is also extensible by design: new dialects ship per application domain, so a fine-tuned model per dialect does not scale. We ask whether inference-time priors derived mechanically from each dialect's Operation Definition Specification (ODS) can substitute for gradient-based adaptation. First, we release four natural-language-to-MLIR benchmarks across three dialects - MLIR-Spec-150, Linalg-Spec-30, StableHLO-Spec-30, and StableHLO-Held-Out-200 - totaling 410 in-scope NL-to-MLIR pairs, plus a 25-program out-of-grammar stress set and a hand-authored n=30 functional reference set, shipped under Apache-2.0 with Gebru datasheets and Croissant 1.0 metadata. Second, we build a three-layer schema-derived constraint stack: a CFG over op signatures(C1), type-domain splits from an ODS-extracted type lattice (C2), and an SSA-scope validator driving five-retry rejection sampling (C3). Porting from arith+func+memref+linalg to StableHLO required no new constraint-layer code. On dialects whose verifier semantics are dominated by structural constraints, schema-derived priors let SmolLM2-1.7B match or exceed 15B-34B code LMs at 8-25x the per-generation speed: on linalg, SmolLM2 reaches 80.0% verify-valid (three-seed mean, n=125), beating CodeLlama-34B, Granite-Code-34B, and StarCoder2-15B by 21-44 percentage points with non-overlapping CIs. On arith+func and on the templated parametric StableHLO-Held-Out-200, where verifier semantics turn on attribute values rather than structure, the same baselines match or beat the SLM; we scope these as non-win cells. We release benchmarks, decoder, all per-prompt generations, and a reproducibility Docker image.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。