Схематическое декодирование: маленькая LM обошла гигантов в генерации MLIR
Исследователи из международной команды представили новый подход к генерации кода на Multi-Level Intermediate Representation (MLIR) — ключевом формате для современных компиляторов машинного обучения. Вместо дообучения под каждый диалект MLIR они предложили использовать информацию из формальных спецификаций операций (ODS) на этапе декодирования.
В рамках работы опубликованы четыре бенчмарка для NL-to-MLIR: MLIR-Spec-150, Linalg-Spec-30, StableHLO-Spec-30 и StableHLO-Held-Out-200 — суммарно 410 пар естественный язык – MLIR, а также тесты на устойчивость и функциональный референс. Все наборы данных распространяются под лицензией Apache-2.0.
Ключевая инновация — трёхуровневый стек ограничений, автоматически извлекаемый из схемы диалекта. Первый уровень задаёт контекстно-свободную грамматику по сигнатурам операций. Второй расщепляет пространство типов на основе решётки типов из ODS. Третий проверяет корректность SSA-областей и запускает до пяти повторных выборок при отбраковке.
Метод не требует нового кода при переносе между диалектами: переход от arith+func+memref+linalg к StableHLO прошёл без изменения слоёв ограничений. На диалектах, где верификатор проверяет в основном структурные ограничения, результаты впечатляют. SmolLM2-1.7B (1,7 млрд параметров) на диалекте linalg достигла 80.0% verify-valid, обойдя CodeLlama-34B, Granite-Code-34B и StarCoder2-15B на 21-44 процентных пункта при непересекающихся доверительных интервалах.
При этом скорость генерации была в 8-25 раз выше по сравнению с крупными моделями. Однако на диалектах, где проверка опирается на значения атрибутов, а не на структуру — например, arith+func и параметризованные StableHLO-Held-Out-200 — большие модели сохраняют преимущество.
Исследователи предоставили код, все поколения для каждого запроса и воспроизводимый образ Docker. Работа подчёркивает потенциал выводных априорных ограничений как альтернативы масштабированию моделей и тонкой настройке в специализированных языковых задачах.
