Flash RNN-T: управление памятью обучения трансдьюсера с помощью блоков и повторного вычисления

При обучении RNN-T память нередко расходуется не кодировщиком, а совместной сетью на решётке всех акустических и текстовых позиций. Flash RNN-T вычисляет её небольшими блоками, сохраняет только нужные вероятности переходов и повторно восстанавливает большие промежуточные тензоры при обратном проходе. В отличие от прореживания выравниваний, это не меняет полную сумму RNN-T.

Метод. Допустимые позиции разных фраз упаковываются без временного дополнения в блоки с ограничением размера. От каждого остаются вероятности blank и очередного правильного символа; состояния и полные словарные оценки удаляются. После вычисления рекурсии вперёд-назад блоки заново рассчитываются и дифференцируются. Сохранённое начальное состояние генератора воспроизводит исходную маску dropout. Специализированные GPU-операции объединяют сложение, нелинейность, dropout и части нормировки; проекция на словарь остаётся матричным умножением PyTorch. Один вызов backward для мини-пакета сохраняется.

Результаты. Основной тест использует H100 80 ГБ, пакет 128, скрытую размерность 640, словарь 1024 и шаг кодировщика 80 мс. Для среднего режима выбранная блочная конфигурация даёт 2839 примеров/с при 1,84 GiB: память примерно в девять раз меньше Apex с сохранением графов, но скорость ниже на 16,3%. Без блочного повторного вычисления Flash повышает скорость на 24–25% и уменьшает память на 49–56% относительно Apex в этом режиме. Для длинных фраз блочный вариант достигает 163,6 примера/с при 4,26 GiB, когда варианты с сохранением всех графов не помещаются даже в 80 ГБ.

Ограничения. Это испытание совместной сети и функции потерь от уже спроецированных состояний, не ускорение полного обучения распознавателя. Замеры исключают компиляцию, настройку и прогрев; динамика форм реального обучения может добавить расходы. Режим освобождения графов по подгруппам позволяет и неблочным решениям обработать длинные фразы, поэтому говорить, что все конкуренты всегда исчерпывают память, неверно. Скорость специализированных операций зависит от GPU; также проверялась RTX 5090. Корректность потерь и градиентов сверялась с эталоном.

Фишка из статьи. Таблица 2 отделяет выигрыш от специальных операций от цены повторного вычисления. Ниже один средний пакет и фиксированный предел блока 200 тысяч позиций; память включает входы и параметры тестируемого участка.

Как это читать: при одинаковых специализированных операциях повторный расчёт сокращает память примерно в десять раз ценой около 30% скорости. Если памяти достаточно, неблочный вариант быстрее. Это настраиваемый обмен вычислений на память, а не одновременное десятикратное улучшение всех характеристик.

Источник. Оригинальное название: Flash RNN-T: Memory-Efficient Joint Network and Transducer Loss Computation. arXiv:2610.23409; полный текст PDF.