Show HN: как обучить модель на 125M параметров автодополнять игру на пианино прямо на устройстве

TL;DR: Трансформер на 125 млн параметров научили автодополнять фортепианные исполнения в реальном времени (~108 нот/сек на iPhone 15). Наибольший прирост дали поиск правильного представления MIDI, агрессивная чистка обучающих данных и добавление DPO на этапе постобучения.

Около года назад появилась идея: подключить MIDI-пианино к телефону, сыграть что-то — и дать ИИ автодополнить мелодию. Что-то вроде GitHub Copilot, но для пианино.

Задача оказалась глубже, чем казалось на первый взгляд. После четырнадцати экспериментов результат наконец достиг той стадии, о которой стоит написать.

Приложение RollTab бесплатно доступно здесь — если есть MIDI-клавиатура и iPhone/iPad.¹

Несколько звуковых примеров

Каждая запись начинается с короткого затравочного фрагмента, за которым следует продолжение, сгенерированное моделью.

Pokémon, Pallet Town (затравка из 8 нот)

Final Fantasy VI, Terra's Theme (затравка из 16 нот)

Für Elise (затравка из 16 нот)

Что находится внутри MIDI-файла

MIDI-файл сильно отличается от MP3 или других аудиоформатов. Вместо записанного звука он хранит музыку как последовательность событий: клавиша нажата с определённой высотой тона и силой удара, клавиша отпущена, изменилось состояние педали сустейна и так далее. Есть и другие события — смена инструмента или изменение громкости.

Эти события обычно организованы в несколько дорожек. Поп- или игровой MIDI-файл может содержать мелодию, аккорды, бас, ударные, струнные и несколько синтезаторных партий. Проект сфокусирован на продолжении фортепианной партии, поэтому большая часть материала, похожего на фортепиано, была сохранена, а остальное — удалено или сокращено.

Как токенизировать музыку

Чтобы обучить трансформер на этих исполнениях, сначала нужно было превратить MIDI-события в дискретную последовательность, которую модель может читать и предсказывать. Самый очевидный вариант — сделать токен для каждого MIDI-события:

NOTE_ON_60_80 # {pitch}_{velocity}
NOTE_OFF_60   # {pitch}
TIME_SHIFT_12 # {time step}

Если включить высоту тона и силу удара прямо в токен NOTE_ON, словарь быстро разрастается. Есть 128 значений высоты тона MIDI и 128 значений силы удара, поэтому наивный объединённый словарь note-on может достигать:

128 * 128 + 128 = 16,512

токенов только для note-on и note-off. На практике силу удара, вероятно, стоит группировать по корзинам, но проблема остаётся: многие комбинации редки, и модели приходится извлекать много структуры из разреженных токенов.

Обычное улучшение — разложить представление по грамматике:

[NOTE_ON, PITCH, VELOCITY] | [NOTE_OFF, PITCH] | [TIME_SHIFT, DURATION]

Теперь пространства вывода меньше:

NOTE_ON / NOTE_OFF / TIME_SHIFT
PITCH: 128 значений
VELOCITY: ~16
DURATION: ~100

Грамматику можно навязать прямо во время генерации, маскируя недопустимые следующие токены. После NOTE_ON допустимы только токены высоты тона. После высоты тона — только токены силы удара. Это гарантирует синтаксически корректный вывод.

Представления в стиле note-on/note-off пробовались, но модели имели тенденцию «плыть». Они забывали выдать note-off, оставляли зависшие ноты или теряли отслеживание активного состояния. Особенно плохо это сказывалось на целевой задаче — маленькой модели, работающей близко к реальному времени на ноутбуке или телефоне.

Ещё одно опробованное представление было ближе к следующему:

[NOTE, PITCH, VELOCITY, DURATION] | [TIME_SHIFT, DURATION]

Это избегает дрейфа note-off, потому что длительность ноты задана явно. Токен сдвига времени продвигает воспроизведение вперёд, когда нота не звучит.

Музыкально это работало лучше, но было медленным. Одна музыкальная нота требовала примерно четырёх авторегрессионных шагов трансформера. Также быстро расходовалось окно контекста.

Итоговое представление

В итоге остановились на следующем представлении:

NOTE(pitch, delta_onset, duration, velocity)

В финальной версии нет отдельного события TIME_SHIFT. Тишина представлена через delta_onset следующей ноты — время с момента начала предыдущей ноты.

Например:

NOTE(C4, delta=0,  duration=12, velocity=80)
NOTE(D4, delta=24, duration=12, velocity=80)

означает: сыграть C4, подождать 24 временных шага до начала следующей ноты, затем сыграть D4.

Аккорды представлены как несколько нот с delta_onset = 0, отсортированных по высоте тона²:

NOTE(C4, delta=24, duration=24, velocity=80)
NOTE(E4, delta=0,  duration=24, velocity=78)
NOTE(G4, delta=0,  duration=24, velocity=82)

Это также не плоский поток токенов вроде:

NOTE, PITCH, DELTA, DURATION, VELOCITY

Вместо того чтобы тратить четыре прохода трансформера на генерацию атрибутов одной ноты, трансформер продвигает музыку сразу на одну целую ноту. На практике это позволяет крупной модели достигать около 108 нот в секунду на iPhone — значительно больше, чем нужно человеку для игры вживую.

Внутри каждая нота имеет пять категориальных полей, каждое со своим словарём³, при этом тайминг квантован по фиксированным шагам⁴.

[event_type, pitch_id, delta_id, duration_id, velocity_id]

Каждое поле получает своё эмбеддинг-представление. Токен ноты — это сумма всех эмбеддингов:

note =
    event_type_embedding[NOTE]
  + pitch_embedding[C4]
  + delta_embedding[12]
  + duration_embedding[24]
  + velocity_embedding[80]

У модели есть отдельные выходные головы: высота тона, дельта, длительность и так далее.

Между полями стоит небольшой вложенный декодер, поэтому более поздние поля могут учитывать уже предсказанные ранние поля. Но дорогостоящий трансформерный backbone запускается только один раз на ноту, а не один раз на поле.

Педаль сустейна

Нажатие педали сустейна на пианино заставляет ноты звучать даже после отпускания клавиши. Не хотелось усложнять реализацию добавлением событий педали. Вместо этого сустейн «запекается» в длительность ноты ещё на этапе предобработки.

Если клавиша отпущена при нажатой педали, нота продлевается до момента отпускания педали. Если та же высота тона сыграна снова раньше, предыдущая нота обрывается в момент повторного удара. В результате длительность ноты приближённо соответствует реальной звучащей длительности.

Это теряет явный жест педали, но заметно упрощает задачу моделирования: модели нужно предсказывать только высоту тона, начало, длительность и силу удара.

Датасет

Было просмотрено множество публично доступных датасетов и коллекций, с фокусом в основном на старой классической музыке в общественном достоянии. Качество разнилось очень сильно, поэтому пришлось написать немало скриптов очистки.

Итоговый датасет содержал несколько сотен тысяч MIDI-файлов — примерно 300 миллионов событий-нот.

Итоговый пайплайн:

• отбирал материал, ориентированный на фортепиано
• удалял или сокращал патологические многодорожечные смеси
• фильтровал по плотности и охвату по высоте тона/времени
• дедуплицировал по отпечаткам, игнорирующим глобальную транспозицию и равномерные изменения темпа
• группировал альтернативные версии одного произведения в один и тот же сплит

Была попытка масштабировать датасет примерно в 5 раз в надежде улучшить результат, но полученные модели оказались хуже. Очистка и отбор данных имели большее значение, чем простое увеличение их объёма.

Обучение

Изначально обучение представляло собой просто кросс-энтропию по пяти выходным головам, суммированную вместе:

type_loss
+ pitch_loss
+ delta_loss
+ duration_loss
+ velocity_loss

Это позволяет отдельно отслеживать точность по высоте тона, длительности и силе удара, вместо того чтобы полагаться на единый агрегированный next-token loss.

Тем не менее у этой цели обучения есть важное ограничение: продолжение музыки не имеет единственно правильного ответа. Отложенная песня даёт модели только одну «правильную» следующую ноту, хотя музыкально подошло бы множество разных продолжений. Кросс-энтропия полезна для изучения механики музыки, но плохо отражает то, насколько хорошо звучит полное продолжение.

Аугментация

Аугментация была важна, потому что живой ввод — не идеальный MIDI-файл. Это человек, играющий на пианино не идеально: ноты могут быть чуть раньше, чуть позже, слишком сильными и так далее.

В итоге остановились на следующих аугментациях:

• глобальная транспозиция
• равномерное масштабирование темпа
• джиттер длительности/силы удара
• отбрасывание нот из затравки

Модель

Архитектура — по сути довольно стандартный decoder-only трансформер: RMSNorm, ротационные позиционные эмбеддинги, каузальное self-attention, блоки SwiGLU/MLP и авторегрессионная генерация.

Обучались в основном три размера модели:

small:  около 33M параметров
medium: около 64M параметров
large:  около 125M параметров

Маленькая модель отлично подходила для быстрых экспериментов, но средняя почти всегда её обгоняла. Большая модель показала себя лучше, хотя и не с большим отрывом.

Сейчас ведётся работа над тем, чтобы приблизить качество средней модели к большой — в основном чтобы уменьшить размер и задержку в iOS-приложении.

Scheduled Sampling

Лучшая базовая модель использовала scheduled sampling между полями каждой ноты. Обычно во время обучения предсказания длительности и силы удара «видят» правильную высоту тона. Но во время инференса им приходится работать с той высотой тона, которую модель реально предсказала.

Поэтому во время обучения модели иногда подавалась её же собственная предсказанная высота тона вместо истинной. Начали с 0% на первых нескольких эпохах, затем постепенно увеличивали долю, доведя до 50% в лучшей модели.

Забавно, но это увеличило validation loss, но улучшило качество продолжений.

Оценка

Поначалу оценка сводилась к простому прослушиванию.

Продолжения генерировались из отложенных песен с затравками от 4 до 32 нот, после чего выходы моделей сравнивались вручную. Это было медленно и утомительно, и через некоторое время всё начинало звучать как шум.

Затравки из четырёх нот оказались самыми сложными: музыкального контекста для работы было слишком мало. Восемь нот работали лучше, а затравки в 16–32 ноты давали заметно более надёжный результат, потому что модели хватало структуры, чтобы понять происходящее.

Генерация без затравки — вопрос удачи, но это и не тот сценарий использования, на который нацелен проект.

Также был написан набор автоматических метрик:

• повторяющиеся n-граммы высоты тона
• энтропия высоты тона
• энтропия классов высоты тона
• диапазон высоты тона
• плотность нот
• длинные паузы
• плотность аккордов

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

В итоге для попарной оценки стала использоваться Gemini 3.5 Flash. Просьба выдать единую абсолютную оценку давала непоследовательные результаты. Вопрос вида «Дано A и B, какое продолжение лучше?» работал гораздо лучше — особенно при зеркалировании каждого сравнения для снижения позиционного смещения⁵. Это позволило собрать достаточно большой датасет предпочтений, который затем использовался для DPO.

Изначально Gemini чрезмерно опиралась на то, насколько хорошо продолжение звучит само по себе, а не на то, насколько хорошо оно следует из затравки. Выходы часто звучали лучше по отдельности, но ощущались оторванными от только что сыгранного.

Улучшение промпта помогло, но в итоге оценка была разделена на два критерия: оценка продолжения — насколько хорошо результат следует из затравки, и оценка «звучит хорошо» — насколько хороша музыкальность сама по себе. Оценка продолжения использовалась как основной сигнал для DPO.

DPO: Direct Preference Optimization

DPO дал наибольший прирост после этапа претрейна. Модель перешла от состояния «иногда выдаёт хорошее продолжение» к гораздо более надёжной работе.

Для каждой затравки генерировалось несколько продолжений, а попарная оценка использовалась, чтобы выбрать лучшее и худшее:

prompt -> chosen continuation
prompt -> rejected continuation

DPO обучает модель делать выбранное продолжение более вероятным, чем отклонённое, сохраняя при этом достаточную близость к исходной модели.

После DPO более 69% продолжений в попарной оценке предпочитались базовой модели.

Значение β определяет, насколько сильно DPO штрафует отклонение от базовой модели. В переборе значений β=0.01 и β=0.03 улучшали модель, а β=0.10 давило слишком сильно и ухудшало результат.

Также был опробован «консенсусный» датасет: вместо доверия каждому шумному суждению о предпочтении, сохранялись только те пары, где оценщик был последовательно согласован. Именно этот вариант дал лучший результат в переборе.

Интуитивное ощущение: базовая модель уже усвоила разумную ментальную модель музыки, просто не знала, что делает продолжение хорошим.

Что не сработало

Многое не сработало:

• Note-on/note-off слишком сильно «плыл» для маленьких моделей реального времени.
• Потоки токенов с грамматической маской были корректны, но медленны.
• Расширение датасета ухудшало результат, когда данные были шумными.
• Более крупные модели помогали, но не решали проблему зацикливания магическим образом.
• Mirostat снижал повторы, но часто делал вывод бессвязным.
• Дополнительные локальные вспомогательные функции потерь замедляли обучение без явного выигрыша на слух.
• Абсолютные скалярные оценки Gemini работали хуже попарного сравнения.
• Одного лишь validation loss было недостаточно, чтобы уловить важные различия в качестве результата генерации.
• Born-again networks (переобучение модели на её собственных «мягких» предсказаниях) здесь не улучшили качество.

Упаковка в приложение

Модель PyTorch была экспортирована в Core ML, а веса квантованы до INT8. Первый запуск всё ещё раздражающе медленный, пока рантайм Apple оптимизирует модель под доступное железо.

Модель обучалась только на контекстах до 512 нот, но хотелось поддерживать более длинные сессии. Как только контекст приближается к лимиту, сохраняются последние 384 ноты, из них заново собирается контекст, и генерация продолжается с этой точки. Это означает пересборку KV-кэша, но модель достаточно быстра, чтобы это не стало серьёзной проблемой.

Для позиционного кодирования использовался RoPE, так что теоретически можно было бы сделать что-то более элегантное со сдвинутыми позициями и кольцевым буфером. К сожалению, Core ML не предоставляет прямого доступа к Q, K и V.

Но на этом этапе главным было просто удовлетворение от того, что всё заработало.

Заключение

Проект получился очень увлекательным. Существует немало интересных статей о генерации музыки, но глубокое погружение в них сознательно откладывалось на потом. Хотелось получить удовольствие от самостоятельного решения задачи, а не просто реализовать чужое исследование. Сравнение своего подхода с существующей литературой было сделано уже после основной работы.⁶

Результат всё ещё далёк от идеала. Иногда возникают зацикливания, короткие затравки даются тяжело, и многое хочется улучшить. Что-то вроде GPT-2, но для пианино.

Но наконец достигнута та точка, когда действительно приятно сесть за пианино, сыграть несколько нот и посмотреть, что получится вместе с моделью.