Одиннаковая модель для преобразования кода в числа

Исследователи из Корнелла и Google представили Regression Language Model (RLM), которая читает исходный текст программ, Triton-ядра и ONNX-графы и напрямую предсказывает числовые метрики: задержку GPU-ядра, пиковое потребление памяти и точность моделей вместе с аппаратно-зависимой задержкой. Подход устраняет необходимость в ручной инженерии признаков, графовых энкодерах или специальных прокси.

Что предсказывает модель

RLM обучается выдавать несколько числовых целей из одного текстового представления. Примеры целей:

• Пиковое потребление памяти для программ на Python, C, C++ и других языках.
• Задержка для Triton GPU-ядр (замеры на устройствах вроде RTX A6000).
• Точность нейросетей и их задержка на конкретных устройствах для архитектур, экспортированных в ONNX.

Выходы генерируются как цифровые токены (знак, экспонента, мантисса) с ограниченным декодированием, что обеспечивает корректность форматов и позволяет моделировать неопределённость через сэмплинг.

Основные результаты

Авторы приводят высокие ранговые корреляции по разным задачам:

• APPS LeetCode, память: Spearman ρ > 0.9.
• Triton latency (A6000): ρ ≈ 0.52.
• CodeNet по 17 языкам, память: среднее ρ > 0.5, для C/C++ ≈ 0.74–0.75.
• Ранжирование NAS: средний Kendall τ ≈ 0.46.

Эти показатели конкурентоспособны с графовыми предикторами и в ряде задач их превосходят.

Архитектура и обучение

Бекбоном служит encoder–decoder, инициализированный T5-Gemma (примерно 300M параметров). На вход подаются сырьевые строки (код, Triton IR, ONNX). Декодер автогрессивно генерирует числовые токены при ограниченном декодировании, чтобы гарантировать валидные числа. Абляции показывают:

• Языковая предобученность ускоряет сходимость и улучшает предсказания задержки Triton.
• Автогрессивная генерация чисел превосходит классические MSE-регрессоры даже после нормализации целей.
• Обучаемые токенайзеры, адаптированные под ONNX-операторы, увеличивают эффективный контекст.
• Более длинные контексты и масштабирование энкодера дают дополнительные улучшения при аккуратной настройке.

Библиотека regress-lm реализует утилиты для обучения, ограниченного декодирования и multi-task предобучения/дообучения.

Почему это важно

Традиционные пайплайны предсказания производительности завязаны на хрупких рукодельных признаках или GNN, привязанных к определённым IR/операторам. Представляя регрессию как генерацию текста-на-число, RLM унифицирует стек: представить артефакты как простой текст, дообучать модель на новых метриках или железе и декодировать откалиброванные числа по цифре. Это снижает поддерживающие затраты и облегчает переносимость между языками и аппаратами.

Датасеты и инструменты

Code-Regression объединяет APPS/LeetCode по памяти, KernelBook-производные измерения задержек Triton и CodeNet по памяти. NAS/ONNX набор экспортирует архитектуры из NASBench, FBNet, Once-for-All и других в ONNX-текст для предсказания точности и задержки по устройствам. regress-lm даёт рецепты и реализации ограниченного декодинга.

Практические выводы

Одна модель (~300M параметров) с инициализацией T5-Gemma способна предсказывать память, задержки Triton и точность+задержку моделей прямо из текстовых представлений без ручной инженерии признаков. Качество ранговых предсказаний делает метод применимым для эвристик компиляторов, отбора/прореживания ядер и многокритериального NAS, при этом снижая порог входа для адаптации под новое железо и языки.