KAN модель. Пробуем на задаче восстановления табличных данных

В этом году в научном сообществе прогремела статья MIT про архитектуру Колмогорова-Арнольда. С тех пор выходили различные теоретические разборы идеи. В ИТМО также попробовали применить модель KAN к своим вычислениям и пришли к довольно оптимистичным выводам, по крайней мере в ракурсе задач восстановления табличных данных. Архитектура KAN выглядит перспективной в задачах с большим количеством параметров. На фоне перцептрона KAN нужно меньше данных для обучения, а оптимизированные пакеты для расчета позволяют выполнять вычисления со скоростью, сравнимой с MLP. Рассказываем подробнее, что у нас получилось.

Что такое KAN и каковы перспективы этой архитектуры?

Для начала вспомним, что собой представляет KAN с точки зрения математики. Эту часть мы попросили прокомментировать руководителя лаборатории композитного искусственного интеллекта ИТМО, кандидата физико-математических наук Александра Хватова.

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

В основе большинства существующих архитектур лежит многослойный перцептрон, предложенный еще в 50-60-х годах прошлого века ― во время первой весны искусственного интеллекта. Эта сущность хорошо изучена, уже сложилась определенная практика эффективных компьютерных вычислений. За прошедшие 50 лет перцептрон претерпевал некоторые эволюционные изменения ― мы перешли на более сложные модели, вроде LLM, но никакой революции с ним не происходило.

Авторы из MIT предложили отойти от этой классической архитектуры ― от многослойного перцептрона ― заменив его на KAN ― сети Колмогорова-Арнольда. Статья с этой идеей была опубликована в Arxiv.org. Ее влияние на научный мир оказалось действительно масштабным, по крайней мере, в части того, что множество научных групп взялись попробовать отойти от классического многослойного перцептрона. В первой статье от MIT есть своеобразная затравочка ― пример, как применить эту архитектуру для решения дифференциальных уравнений. Позже появились публикации о том, как KAN справляется с распознаванием изображений, прогнозированием временных рядов и так далее.

В научном мире KAN наделали много шума. Однако каждое столь масштабное изменение должно давать прирост в каких-то задачах. Пока KAN не дает заметного прироста. Авторы первоначальной статьи, к сожалению, не проводили подробного исследования, где эта разработка могла бы быть наиболее полезной. Такие исследования появились позднее (например, вот и вот.

Многослойными перцептронами в ИИ занимаются более полувека, а значит у нас есть самые разнообразные инструменты и методы для работы с ними, в том числе очень быстрые. Для KAN такие методы еще не сформировались. Да, есть готовые реализации, но пока это своего рода базовый вариант. Будут ли формироваться более эффективные методы ― большой вопрос, потому что для этого нужно сначала на каком-то классе задач получить прирост в скорости или точности решения, а потом уже думать, как все ускорить. Пока прироста нет. В среднем KAN работают медленнее. Хотя уже появляются реализации, сравнимые  по числу операций ― то есть по скорости обучения ― с перцептроном, например, вот и вот. 

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

Не так давно ― 3-4 года назад ― уже предлагали альтернативный фундамент ― гиперболические нейронные сети. Это довольно сложная идея, в рамках которой мы переходим от Евклидового пространства к геометрии Лобачевского, где арифметические операции с матрицами выполняются по-другому. И эта архитектура нашла свое применение в рекомендательных системах, дав там неожиданный прирост. Возможно, какую-то аналогичную нишу со временем найдут и для сетей Колмогорова-Арнольда.

Думаю, что на хайпе KAN мы увидим некоторый пласт работ, где исследователи будут пытаться применить сети Колмогорова-Арнольда в своих частных задачах. Первую волну работ стоит ждать уже к зимним конференциям. Возможно, в части из них результаты будут чуть-чуть приукрашены. К первым восторженным заголовкам всегда стоит относиться очень аккуратно, проверяя результаты самостоятельно или выискивая так называемые синхронные источники и складывая картину из результатов разных научных групп. Будем анализировать и искать не случайный, а закономерный прирост.

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

Практический пример: восстановление табличных данных

Мы в ИТМО еще несколько месяцев назад попытались применить эту модель к восстановлению табличных данных одного из экспериментов Центра ИИ в Химии. Суть эксперимента для данного примера не имеет значения. Важно, что результат ― это таблица, в нашем случае содержащая 3 млн. записей.

Чтобы проверить возможности архитектуры, заменим около 30% датасета на специфическое значение, которое сильно отличается от корректного результата (чтобы пропуски проще было найти). После этого используем KAN для восстановления таблицы, вычисляя для каждого функцию потерь.

Ниже представлен простой автоэнкодер, который в своей основе имеет только KAN (никаких других архитектур).

from efficient_kan.kan import KAN import torch from dataclasses import dataclass   @dataclass class KANAEConfig:    struct: list[int]    grid: int    k: int   class EfEncoder(torch.nn.Module):    def __init__(self, config: KANAEConfig):        super().__init__()         self.kan = KAN(layers_hidden=config.struct,                       grid_size=config.grid,                       spline_order=config.k,                       )     def forward(self, x):        x = self.kan(x)        return x   class EfDecoder(torch.nn.Module):    def __init__(self, config: KANAEConfig):        super().__init__()        self.str = config.struct[::-1]        self.kan = KAN(layers_hidden=self.str,                       grid_size=config.grid,                       spline_order=config.k,                       )     def forward(self, x):        x = self.kan(x)        return x   class EfKANAE(torch.nn.Module):    def __init__(self, config):        super().__init__()        self.encoder = EfEncoder(config)        self.decoder = EfDecoder(config)     def forward(self, x):        x = self.encoder(x)        x = self.decoder(x)        return x     def encode(self, x):        return self.encoder(x)      def decode(self, x):        return self.decoder(x)

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

Если оценить результат по MSE на данных размерности (35357, 86), получаются следующие цифры:

loss: 0.007484327536076307: 100%|██████████| 25/25

loss: 0.006712096743285656: 100%|██████████| 25/25

loss: 0.004756155889481306: 100%|██████████| 25/25

loss: 0.006610468495637178: 100%|██████████| 25/25

loss: 0.004374181851744652: 100%|██████████| 25/25

loss: 0.0044711739756166935: 100%|██████████| 25/25

loss: 0.003979865461587906: 100%|██████████| 25/25

loss: 0.0039029153995215893: 100%|██████████| 25/25

loss: 0.00714501366019249: 100%|██████████| 25/25

loss: 0.0038636275567114353: 100%|██████████| 25/25

Для расчетов использовался пакет Efficient KAN. Это оптимизированная версия оригинального KAN, скорость работы которого сравнима с MLP по времени.

А что дальше?

По итогам экспериментов архитектура KAN показалась коллегам из Центра ИИ в Химии довольно перспективной. В многоуровневом перцептроне фактически есть большое количество нейронов, которые проходят через одну функцию активации. По сути функция активации определяет справа или слева находится нейрон по значению (т.е. относится он к одному или к другому классу). В KAN на каждый нейрон предусмотрено столько функций активации, сколько нейронов в следующем слое. За счет этого нейронной сети нужен меньший объем данных для обучения, и она сама имеет меньшие размеры. Правда, при этом процесс обучения занимает больше времени. Таким образом, KAN лучше всего подойдет к задачам, где мало данных. И это не обязательно должно быть восстановление таблицы. Уже когда готовился этот эксперимент, вышла публикация, сравнивающая использование архитектур MLP и KAN для ряда других задач. 

Ждем более удачных реализаций KAN.

Автор эксперимента: Алексей Дмитренко, ML-инженер, Центр ИИ в Химии.