Материал подготовлен автоматически по первоисточникам: ссылки на них — в конце статьи.
Naveen Mysore представил FlashKAN — заявленную как совместимая замену стандартным слоям B-spline KAN; это препринт, не прошедший рецензирование. Вместо последовательного расчёта базисных функций автор сводит его к единому объединённому GPU-ядру. Работа важна командам, которые уже используют KAN с равномерной сеткой и упираются в стоимость вычисления сплайнов.
Что сделали
В обычной нейросети фиксированная функция активации применяется в узле. В KAN обучаемые функции размещаются на связях между узлами, а каждая такая функция собирается из B-сплайнов — кусочно-полиномиальных базисных функций с локальной областью влияния.
Стандартные реализации вычисляют B-сплайны рекурсией Cox-de Boor. Каждый проход зависит от результата предыдущего, поэтому GPU не может выполнить их параллельно. Дополнительные операции поиска нужного участка сетки и выборки данных также создают нерегулярный доступ к памяти.
FlashKAN заменяет рекурсию замкнутой формулой через усечённые степенные функции: сдвинутые выражения вида max(0, x)^3 складываются с постоянными коэффициентами. Для равномерной сетки, где узлы расположены на одинаковом расстоянии, формула не требует поиска участка и содержит только поэлементные операции. По утверждению автора, torch.compile объединяет их в одно GPU-ядро.
Проблема такого представления — потеря точности при вычитании близких больших величин. Она возникает, когда вход далеко выходит за локальную область сплайна, хотя правильный результат там равен нулю. FlashKAN ограничивает нормализованную координату границами этой области до вычисления степеней. Автор утверждает, что это не меняет значение функции внутри области и ограничивает промежуточные величины снаружи.
Меняется только вычисление базиса. Обучаемые коэффициенты, остаточная ветвь, прохождение градиентов и интерфейс слоя должны оставаться прежними. Код оформлен как пакет под лицензией MIT, который автор предлагает подключать вместо существующей реализации KAN без изменения остальной модели.
Что показали
Главный количественный результат работы относится не к FlashKAN, а к исходной реализации. В профиле одного слоя вычисление B-сплайнов заняло 91% прямого прохода: 1,44 мс из общих 1,58 мс. Это показывает, что рекурсия действительно была основным локальным узким местом в выбранной конфигурации.
Обоснование численной устойчивости преимущественно аналитическое. По расчёту автора, в float16 куб нормализованной координаты может переполниться при значениях примерно выше 40. После ограничения координаты максимальное промежуточное кубическое значение составляет 64, поэтому ошибка больше не должна расти по мере удаления входа от области сплайна.
Включённые в пакет модульные тесты проверяют совпадение базисных функций, локальность, сумму базисов, гладкость и прохождение градиента. Однако в тексте нет прямого замера того, насколько быстрее выполняется сам FlashKAN.
Ограничения
Замкнутая формула применима к равномерным сеткам. Неравномерные узлы нужны, когда сетку адаптируют к распределению входов или уточняют только отдельные области. Для них потребуются коэффициенты по участкам и зависимый от данных выбор нужного участка, то есть часть устранённых расходов вернётся. Расширение на такие сетки автор оставляет для будущей работы.
Профилирование проведено для одного слоя формы 256×784→64 на GPU через MPS. В работе не приведены замеры полного обучения или вывода модели, расхода памяти, поведения на нескольких GPU и длительной устойчивости в смешанной точности. Также нет прикладного датасета или задачи, на которой можно было бы сравнить качество и сходимость.
Авторы обсуждают Efficient-KAN и FastKAN, но не приводят экспериментального сравнения с ними по скорости, памяти или качеству. Нет и прямой таблицы производительности FlashKAN против исходной рекурсии на одинаковом оборудовании. Поэтому работа показывает устройство оптимизации и наличие узкого места, но не устанавливает величину итогового ускорения.
Что это значит
Для команды, уже строящей продукт на равномерных кубических B-spline KAN в PyTorch, работа может изменить ближайший технический план: вместо переписывания архитектуры имеет смысл проверить FlashKAN как локальную замену вычислительного слоя. Совпадение интерфейса и сохранение базисных свойств, если они подтвердятся в собственных тестах, снижают объём миграции.
Планирование производительности и инфраструктуры по этому препринту менять рано. Перед внедрением нужен локальный замер полного сценария: обучение и вывод, целевые размеры пакета данных, используемый GPU, требуемая точность и реальные диапазоны скрытых активаций. Особенно важно проверить не только время базиса, но и долю этого базиса после компиляции всей модели: устранение локального узкого места может сделать ограничением умножение на веса или обмен данными.
Для команд, которым нужны адаптивные неравномерные сетки, FlashKAN в описанном виде планы не меняет. Работа также не даёт оснований выбирать KAN вместо обычной многослойной сети: она оптимизирует конкретную реализацию уже выбранной архитектуры, но не сравнивает архитектуры по качеству, стоимости обучения или эксплуатации.
Источники
Похоже на вашу задачу?
Расскажите, что собираете. За полчаса разложим на этапы и назовём сроки — это бесплатно и ни к чему не обязывает.
