Сравнение ансамбля алгоритмов машинного обучения и BERT для анализа текстовых описаний КТ головного мозга на предмет наличия внутричерепных кровоизлияний

Цель исследования — обучить и протестировать ансамбль моделей машинного обучения, а также сравнить характеристики его работы с предобученной на медицинских данных языковой моделью BERT в задаче простой бинарной классификации наличия/отсутствия признаков внутричерепного кровоизлияния (ВЧК) в протоколах описаний КТ головного мозга.

Материалы и методы. В качестве моделей, с помощью которых решалась задача бинарной классификации, было выбрано 7 алгоритмов машинного обучения и 3 метода векторизации текста. Обучение моделей проводили на текстовых данных, которые представляли собой протоколы описаний 3980 КТ-исследований головного мозга из 56 медицинских организаций стационарной медицинской помощи Москвы. Эти тексты были векторизованы тремя способами: «мешок слов», TF-IDF и Word2Vec. Далее к ним применяли следующие алгоритмы машинного обучения: дерево решений, случайный лес, логистическая регрессия, метод ближайших соседей, метод опорных векторов, CatBoost и XGBoost. Анализ данных, а также их предварительную обработку осуществляли с использованием библиотеки NLTK (Natural Language Toolkit, версия 3.6.5.) и библиотеки Scikit-learn (версия 0.24.2). В качестве предообученной на медицинских данных модели-трансформера BERT была взята версия MedRuBertTiny2.

Результаты. По результатам обучения и тестирования семи алгоритмов машинного обучения выбраны три алгоритма с наиболее высокими метриками (чувствительность, специфичность): CatBoost, логистическая регрессия и метод ближайших соседей. Самые высокие метрики получены при использовании метода векторизации текста «мешок слов». Эти алгоритмы были собраны в ансамбль методом стекинга (stacking). Показатели чувствительности и специфичности на тестовом наборе данных из исходной выборки составили 0,93 и 0,90 соответственно. Далее ансамбль и модель BERT были обучены на независимом наборе данных, содержащем 9393 текстовых протокола диагностических описаний, разделенных также на обучающую и тестовую выборки. При тестировании на этом наборе данных ансамбля алгоритмов машинного обучения чувствительность и специфичность составила 0,92 и 0,90 соответственно. Тестирование на этих данных модели BERT продемонстрировало чувствительность 0,97 и специфичность 0,90.

Заключение. Обученный ансамбль показал высокие результаты точности работы при анализе текстовых протоколов описаний КТ головного мозга с признаками внутричерепных кровоизлияний, но все равно при его использовании необходимо обеспечить ручной пересмотр результатов для контроля качества. Предобученная модель-трансформер BERT, дополнительно обученная на диагностических текстах, продемонстрировала более высокие метрики точности (p<0,05). Это говорит о перспективности модели в задачах бинарной классификации и для поиска информации по неструктурированным медицинским записям.