Привет, Хабр!

Меня зовут Дима Архипов. Я учусь на четвертом курсе института на направлении прикладной математики в НИТУ МИСИС. В марте 2024 года мне удалось попасть на стажировку в центр медицины Sber AI Lab, где я занимался классификацией ЭЭГ сигнала в режиме реального времени. Эта тема крайне важна, поскольку анализ электроэнцефалограммы (ЭЭГ) может помочь в диагностике различных неврологических расстройств, таких как эпилепсия, нарушения сна и другие. Использование искусственного интеллекта для классификации ЭЭГ сигнала позволяет повысить точность и скорость диагностики, что, в свою очередь, способствует улучшению качества медицинской помощи.

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

0. Что такое ЭЭГ-сигнал?

Согласно Википедии электроэнцефалография (ЭЭГ) — неинвазивный метод исследования функционального состояния головного мозга путём регистрации его биоэлектрической активности. Простыми словами, электроэнцефалография помогает определять активность мозга, будь то состояние бодрствования, сна или эпилептический приступ, не требуя для этого проникновения внутрь организма.

Снятие параметров сигнала происходит при помощи такой шапочки:

При записи получаются сигналы такого вида:

Стоит отметить, что электроэнцефалография – это не новинка в медицинской области. В этом году этот метод исследования отмечает столетие, ещё в 1924 году Ханс Бергер получил первую запись ЭЭГ человека. Однако для области  машинного обучения работа с электроэнцефалографией пока не является чем-то классическим или обыденным, а это значит, что в ней есть интересные вопросы для исследований и разработки. 

1. Классификация ЭЭГ сигнала в реальном времени

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

Для обучения ML моделей необходимы данные, поэтому был взят открытый датасет TUH EEG Seizure Corpus v2.0.0. Датасет содержит train (4610 записей), dev (1831 записей) и eval (865 записей) выборки. Для обучения мы взяли train выборку, для валидации — dev, а для теста — eval, как и предполагается авторами датасета. Наш выбор остановился именно на нём по нескольким причинам:

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

• датасет поддерживается сообществом исследователей, регулярно обновляется и дополняется;

• Де-факто это стандарт в области анализа ЭЭГ –  удобно сравниваться с результатами других научных работ.

Первоначально необходимо было воспроизвести работу, описанной в статье Real-Time Seizure Detection using EEG: A Comprehensive Comparison of Recent Approaches under a Realistic Setting, которая исследует ровно ту же самую задачу. Основная идея этого подхода заключалась в классификации последних 4 секунд ЭЭГ сигнала каждую секунду с помощью LSTM с энкодером из простой CNN, состоящей из одномерных сверточных слоев. Однако, при работе с кодом и подходом, описанным в статье вскрылось множество проблем, речь о которых пойдет далее.

Pipeline обучения

В своей статье авторы преобразуют исходный сигнал ЭЭГ в биполярную схему с 20 каналами и частотой дискретизации 200 Гц. Это означает, что данные сэмплируются 200 раз в секунду для каждого из 20 каналов, что создает 20 одномерных временных рядов (по одному для каждого канала) или, иначе говоря, один многомерный временной ряд, где каждый канал представляет собой отдельное измерение.

Для обработки данных авторы используют подход, в котором каждый одномерный временной ряд обрабатывается независимо от других. Они применяют к каждому каналу энкодер (CNN), чтобы извлечь признаки из временных рядов. После этого признаки, полученные из каждого временного ряда, усредняются, чтобы получить признаки уже многомерного временного ряда, которые в свою очередь подаются на вход LSTM. Схематически модель выглядит так:

Для обучения модели данные из одной записи ЭЭГ разбиваются на фрагменты длиной 30 секунд. Затем каждый из этих фрагментов делится на 4-секундные окна с перекрытием в 3 секунды между последовательными окнами. В результате каждый 30-секундный сегмент разбивается на 27 перекрывающихся 4-секундных окон.

Модель в процессе обучения получает массивы данных размером , где:

• 27 – количество 4-секундных окон, необходимых для полного покрытия 30-секундного фрагмента;

• 20 – количество каналов ЭЭГ-сигнала;

• 800 – количество точек отсчета в каждом 4-секундном сегменте при частоте в 200 Гц.

При изучении датасета видно, что большую часть времени записей занимают отрезки, где у пациентов нет приступов. Таким образом, модель будет будет видеть слишком мало примеров, включающих в себя участки с приступами, что в свою очередь снижает качество предсказаний. Для борьбы с таким эффектом авторы делят всю обучающую выборку на 6 типов, классов и сэмплируют примеры при обучении так, чтобы модель “видела” равномерное количество примеров каждого типа. В статье предлагается следующее разделение на классы:

• Non-Ictal (Patients) — отсутствие приступа в сегменте, содержащимся в записи, имеющей отрезки приступов;

• Non-Ictal (Patients, Normal Control) — отсутствие приступа в сегменте, содержащимся в записи, не имеющей отрезки приступов;

• Ictal — сегмент полностью состоящий из приступа;

• Ictal Onset — сегмент, начало которого не является приступом, а конец является;

• Ictal Offset — сегмент, начало которого является приступом, а конец не является;

• Alternated — сегмент, содержащий минимум 1 полный отрезок приступа или нормального состояния.

Схематически отрезки записей каждого класса выглядят так:

Важно рассказать о том, как “учится” модель в официальной имплементации статьи. Дело в том, что обучение построено нестандартно и не очевидно, о чем не указано ни слова в статье. Код, используемый в имплементации, схематически представлен ниже.

# объявление модели и лосса
model = CNN1D_LSTM(...)
criterion = nn.CrossEntropyLoss()

# итерация по батчам
for batch in dataloader:
  x, y = batch

  # итерация по окнам
  for i in range(n_windows):

    # получение i-ого окна и соответствующего таргета
    slices, targets = x[:, [i], ...], y[:, i]

    # получение предсказания
    output = model.forward(slices)

    # обновление весов
    loss = criterion(output, targets)
    loss.backward(retain_graph=False)
    optimizer.step()
    optimizer.zero_grad()
    scheduler.step()

Для пояснения: тензор x содержит входные данные и имеет размерность (B,  N,  C,  T) где:

• B — размер батча,

• N — количество окон (27),

• C — количество каналов (20),

• T — количество отсчетов в одном окне (800).

Тензор y представляет собой метки классов, указывающие на наличие приступов в каждом окне, и имеет размерность (B, N).

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

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

Невоспроизводимость результатов

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

Перед тем как сделать выводы о сравнении, стоит сказать, что за все время работы единственной метрикой, вызывающей проблемы, была Precision-Recall кривая, поэтому и выводы о том, какой подход лучше, мы делали в основном основываясь на том, какая модель имеет выше значение Precision-Recall AUC. Как видно результаты даже близко не повторяют тех, что указаны авторами. В попытке достигнуть метрик авторов, мы стали искать решение проблемы. Однако необходимо сделать важное замечание. При рассмотрении кода авторов можно заметить, как они смешивают валидационную и тестовую выборки. Это проявляется в том, что в валидацию и тест попадают куски одних и тех же записей, что является утечкой. Поэтому строго говоря результат авторов невалиден, так как он был получен не на оригинальной тестовой выборке из датасета. Об этом, конечно же, авторы в статье не упомянули.

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

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

Тест подходов к обучению RNN

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

После получения знания о подходе авторов к обучению RNN возникает закономерный вопрос: работает ли это лучше в сравнении с классическим подходом? Для выяснения этого мы провели прямое сравнение подхода авторов с привычным обучением RNN при полностью идентичной модели, данных, гиперпараметров и т.п. Сравнение подхода к обучению авторов и классического обучения будет происходить через сравнение метрик на тестовой выборке:

Учитывая наш акцент на метрике PR-AUC, мы посчитали, что подход авторов действительно имеет смысл и нет нужды переходить к классическому обучению RNN.

Дополнительные улучшения

В попытках улучшить результат мы попробовали 2 дополнительные техники машинного обучения, которые не применяли авторы, а именно:

• negative mining;

• аугментации.

Negative Mining

Как известно ЭЭГ-сигнал крайне шумный по своей природе, подвержен искажениям из-за различных артефактов таких как моргание, движение мышц и т.п. Мы решили рассмотреть гипотезу, согласно которой наша обученная модель может путать отрезки, содержащие артефакты, и отрезки с приступами. Поскольку в используемом датасете разметка не содержит данных об артефактах, то просто так мы не могли проверить эту теорию. Для этого мы решили попробовать применить технику negative mining и посмотреть, станут ли результаты лучше.

Negative mining — это метод, применяемый в машинном обучении для улучшения качества классификации. Он заключается в том, что модель чаще видит негативные примеры (те, которые не относятся к целевому классу), которые она ошибочно классифицирует как положительные. Эти "трудные" негативные примеры получают больший вес при обучении, что заставляет модель более тщательно их анализировать и различать от положительных примеров.

Взяв дополнительный датасет TUH EEG Artifact Corpus v3.0.1, состоящий из 159 записей с с артефактами, от тех же авторов, мы добавили данные в обучающую выборку и изменили классы для сэмплирования. Мы убрали класс Non-Ictal (Patients, Normal Control) (отсутствие приступа в сегменте, содержащимся в записи, не имеющей отрезки приступов) и добавили класс Artefact. Сравнительная таблица показана ниже:

Как видно наша гипотеза не подтвердилась — никаких улучшений использование записей с артефактами при обучении не дало.

Аугментации

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

Список опробованных аугментаций:

• Time Mask, вероятность — 0.4: Маскирование случайного временного интервала по всем каналам в каждом окне в сэмпле, в результате чего выбранная часть временного ряда заменяется нулями, чтобы создать эффект пропуска данных в этом интервале.

• Frequency Shift, вероятность — 0.33: Изменение частотных компонентов сигнала путем добавления сдвига к спектру. Это приводит к изменению основных частотных характеристик сигнала.

• Channel Dropout, вероятность — 0.4: Случайное зануление до 4 каналов ЭЭГ-сигнала.

• Discrete Cosine Transform Convert, вероятность — 0.2: Преобразование временного сигнала одним алгоритмом конвертации в спектральное представление с использованием дискретного косинусного преобразования (DCT) и обратно, используя другой алгоритм конвертации.

• Sign Flip, вероятность — 0.2: Инвертирование знака значений сигнала (умножение всех значений на -1) в каждом канале, что отражает сигнал относительно оси времени.

• Gaussian Noise, вероятность — 0.5: Добавление шума из нормального распределения со стандартным отклонением в 13 от стандартного отклонения в сэмпле.

• Cut And Wise, вероятность — 0.2: Обрезание первых 200 отсчетов у всех окон в сэмле и линейное интерполирование обратно до 800 отсчетов.

• Time Reverse, вероятность — 0.2: Инверсия временной оси сигнала в каждом окне отдельно, при которой временные ряды сигналов воспроизводятся в обратном порядке.

• Band Filter, вероятность — 0.2: Применение полосового фильтра, который фильтрует частоты от 10 до 90 Гц.

• Amplitude Scale, вероятность — 0.4: Умножение всех значений сигнала на 0.66 для изменения амплитуды сигнала.

Внизу представлена таблица с метриками, показывающими влияние аугментации на результат:

Как видно ни одна аугментация не сделала модели хуже, однако много аугментаций дали хороший прирост в “главной” метрике PR-AUC, а именно: Time Mask, Gaussian Noise, Cut And Wise, Time Reverse, Band Filter, Amplitude Scale.

2. Построение собственного решения

Не получив удовлетворительных результатов, используя подход авторов, речь о котором шла ранее, мы перешли к построению собственного решения задачи классификации ЭЭГ-сигнала. Мы решили воспользоваться знаниями из области цифровой обработки сигналов и использовать вейвлет преобразования.

Подробно о вейвлет преобразованиях на Хабре уже писали, поэтому здесь мы опустим математическую составляющую и опишем лишь суть того, что мы получаем. Если кратко, то вейвлет-преобразование позволяет получить детальное представление сигнала как во временной, так и в частотной области. Это означает, что мы можем одновременно анализировать, какие частоты и фазы присутствуют в сигнале и в какие моменты времени они проявляются. Похожим на вейвлет-преобразование является популярное кратковременное преобразование Фурье (STFT), однако последнее не имеет тех преимуществ, что есть у вейвлет. Так можно выделить следующие преимущества:

• локализация: анализ сигнала на разных частотах с разным разрешением с сохранением качества;

• адаптивность: возможность выбора вейвлета под конкретную задачу для лучшей локализации.

Примеры того, как могут выглядеть различные участки ЭЭГ-сигнала после применения вейвлет-преобразования, показаны на изображении ниже:

Для применения преобразования необходимо выбрать функцию, на основе которой будет строится само преобразование. Выбор происходит из ограниченного класса функций, называемых вейвлетами. В литературе, связанной с анализом ЭЭГ-сигнала и эпилептических припадков, основной для использования выбирают Комплексную вейвлет Морле . Определившись с используемой функцией, мы стали искать лучшие диапазон частот и параметры функции: центральную частоту (далее С) и ширину полосы (далее В).

Эксперименты

Для проведения экспериментов использовался все тот же датасет TUH EEG Seizure Corpus v2.0.0 и то же равномерное сэмплирование 6 классов при обучении, что и в статье, бывшей ориентиром в первой части статьи. Перед тем, как рассказать о результатах, стоит рассказать об предобработке данных.

Обработка данных

На выходе из вейвлет-преобразования мы получаем 20 матриц, так как ЭЭГ-сигнал имеет 20 каналов. Каждая матрица состоит из комплексных чисел (обусловленно нашей комплекснозначной вейвлет). Пример того, как может выглядеть матрица на изображении ниже:

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

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

Также из-за разной величины отклика (чем больше частота, тем меньше амплитуда отклика в матрице), то мы могли применить стандартизацию по временной оси матрицы или логарифм к амплитуде.

Также необходимо уточнение про изменение частоты выходного сигнала. На вход вейвлет-преобразованию подаются сигнала с частотой в 200 Гц. Однако из-за большого веса получаемых после преобразования данных мы ограничены в использовании сигнала высокой частоты. Поэтому выходной сигнал из вейвлет-преобразования, имеющий ту же частоту в 200 Гц, мы преобразовали в сигнал частотой в 10 Гц. Преобразование происходило усреднением отрезков длиной в 20 значений.

Тесты

Объяснив, как формируются вектора, подаваемых на вход модели, можно перейти к результатам. За все время тестов на все искомые параметры были наложены ограничения:

• Параметр С мог быть и ;

• Параметр В был выбран 2;

• Нижний диапазон частот был выбран 1 Гц;

• Верхний диапазон частот перебирался от 10 Гц до 100 Гц с шагом 10 Гц;

• Разрешение по частоте было выбрано в 32;

• Частоты были равномерно расположены на всем диапазоне.

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

Вначале в качестве модели мы взяли простую модель, которая состоит из последовательных блоков: Batch Norm, LSTM, Linear. В модель подавались только амплитуды комплексных чисел, к которым были применена стандартизация. Лучший результат:

Далее пришла идея заменить стандартизацию на логарифм. Лучший результат:

Далее мы добавили фазу, но при подсчете мы не применяли arctg. Лучший результат:

Как видно, “главная” метрика PR-AUC планомерно растет, однако она все еще принципиально ниже той, что мы достигали при старом подходе. Поэтому мы решили попробовать другие RNN: S4 и новую xLSTM.

S4 — это не новая модель, однако она представляет новый тренд в DL. S4 принадлежит семейству линейных RNN, самым производительным и набравшим популярность представителем которой является Mamba, и именно она начала тренд на это семейство RNN. S4 дала лучший прирост, однако наверняка можно еще лучше, так как мы использовали неофициальную реализацию и могли неоптимального подобрать гиперпараметры. Лучший результат:

xLSTM — это самая свежая и самая большая из опробованных нами модель. Модель развивает идею оригинальной LSTM и привносит в архитектуру блоки на манер трансформеров. С помощью этой модели удалось получить схожие метрики, и вероятно, что это тоже не предел, так как гиперпараметры выбраны не самым оптимальные. Лучший результат:

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

3. Заключение

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