Улучшение нейросетевого прогноза температуры воздуха

Всем привет! В этой статье расскажу про свои ошибки и новшества во время работы над проектом, а точнее его улучшением. Эта статья является продолжением серии статей про данный проект.

Речь пойдет о моем проекте edge-weather-forecast — лёгкой нейросетевой модели прогнозирования температуры, которую можно запускать прямо на метеостанции или на простом CPU-устройстве вроде Raspberry Pi.

Почему я вообще решил продолжать работу над проектом

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

Использование Gatv2Conv вместо GCNConv

GCNConv агрегирует соседей с фиксированными весами, которые зависят напрямую от структуры графа. Да, сейчас на 14 станциях этого может быть достаточно. Но для более сложных зависимостей и для расширения на весь мир требуются доработки. Кстати, стало даже лучше. GATv2Conv вычисляет веса динамически - в зависимости от признаков узлов. Для метеостанций это важно, потому что соседняя станция в горах и станция на равнине должны влиять по-разному.

В итоге я сделал комбинацию, теперь мой GNN выглядит так: GCNConv -> GATv2Conv. Кстати, GATv2 исправляет баг оригинального GAT (статичное ранжирование внимания). Логика простая: GCNConv дешево делает грубую агрегацию соседей и формирует базовые представления узлов, обогащенные контекстом соседства, а затем GATv2Conv уже работает с этими представлениями и вычисляет веса внимания точнее.

Почему не два GATv2? Дорого и избыточно. Внимание на первом слое по сырым признакам вычисляется "вслепую", сигнал еще шумный, и механизм внимания не дает реального выигрыша.

Добавление нормализации в слои (почему не BatchNorm)

Раньше нормализации в TCN слоях не было. На первый взгляд кажется, что это упрощение - меньше параметров, быстрее обучение. Однако это приводит к нестабильности градиентов в модели, а масштабы активаций в разных каналах могли расходиться во время обучения.

Казалось бы, добавить BatchNorm как универсальное и частоиспользуемое решение - и дело сделано. И батч размером 112 достаточен для стабильных статистик. Однако проблема не в размере, а в составе. В одном бачте одновременно могут оказаться станции из совершенно разных климатических зон (сейчас это не так заметно, но при расширении очень даже). Например, если в одном батче будем нормализовывать сибирь и тропики - мы внесем шум, а не порядок.

GroupNorm же нормализует каждый пример независимо, по группам каналов. Каждая станция нормализуется сама по себе, не происходит смешение ее контекста с другими станциями. Поведение одинаково на train и inference, что важно для стабильного продакшн-прогноза.

BiasHead теперь на весь горизонт

Ранее BiasHead возвращал одно скалярное значение - единую поправку на весь горизонт. Это грубое допущение: систематическое смещение на +6 часов м на +72 часа разные по природе ошибки.

Теперь BiasHead предсказывает отдельную поправку на каждый шаг горизонта. Модель может выучить, что на коротких горизонтах смещение незначительное, а на длинных нарастает и имеет определенную форму (например регрессия к климатологической норме).

Attention Pool вместо последнего состояния

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

AttentionPool взвешивает все временные шаги и формирует итоговое представление как взвешенную сумму. Модель сама учится, какие моменты временны важны - это некая дополнительная степень свободы практически без прироста числа параметров.

FiLM наряду с обычной конкатенацией

Feature-wise Linear Modulation - это механизм, при котором внешние условия, например эмбеддинги станций не просто конкатенируются к признакам, а модулируются через scale и shift. Теоретически в добавок к обычной конкатенации добавляет мощности.

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

Skip из энкодера в декодер

Идея заимствована из U-Net: декодер на каждом уровне получает не только сжатое представление, но и "сырые" признаки соответствующего уровня энкодера через прямое соединение. Это позволяет не терять краткосрочные паттерны при сжатии и улучшить градиентный поток.

На малых датасетах skip-connections практически не дают прироста качества - модель просто не успевает выучить, как ими правильно пользоваться. Это решение ориентировано на будущее: при расширении датасета (больше станций, длиннее история) skip-соединения раскроют свой потенциал. Сейчас это больше инвестиция в масштабируемость архитектуры.

Настраиваемые слои в EncoderDecoder

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

Новый спектральный слой

Периодические паттерны, например суточный цикл, и сезонность - это то, с чем классические TCN работаю неэффективно. Им нужно большое рецептивное поле, чтобы увидеть период, тем более не всегда это у них хорошо получается. Спектральный слой явно работает в частотном домене, в нем строго зашита дневная и недельная цикличность. При этом маски - мягкие и обучаемые, модель сама подбирает насколько широко считать "суточным" паттерном. Также здесь каждая станция получает индивидуальный спектральный фильтр, как бы фильтруя саму себя. Это позволяет улучшить градиентный поток в GNN.

Данный слой ставится строго после последнего слоя только в энкодере.

Облегчение модели в несколько раз

Одним из самых неожиданных улучшений стало... уменьшение модели. Анализ кривых обучения показал признаки переобучения: train loss уверенно снижался, validation loss нет. Причина в избыточной ёмкости: модель запоминала особенности конкретных станций в обучающей выборке, а не общие физические закономерности.

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

Да, метрики слегка упали на малом горизонте, но модель стала более смелой и точной на дальнем горизонте.

Изменения в стратегии обучения Teacher и Student

Общие изменения: я теперь использую более агрессивные lr и wd. А также в качестве планировщика используется CosineAnnealingLR с линейным warmup.

Начнем с Teacher. Напомню, это модель которая обучает GNN и имеет слив статиков станций напрямую в декодер с большим drop. Важно, что Teacher обучается без BiasHead и AttentionPool. BiasHead может портить обучение GNN, так как может брать на себя много ответственности за станцию. AttentionPool аналогично. Если включить его сразу, он начинает выбирать "важные" временны́е шаги ещё до того, как энкодер научился делать эти шаги содержательными. Более того, AttentionPool стоит после энкодера, а энкодер использует эмбеддинг из GNN. Значит AttentionPool косвенно влияет на то, какой градиент получает GNN. Пока GNN не устоялся, этот сигнал шумный и может мешать. Количество слоев рассчитаны так, чтобы RF покрывал весь входной горизонт. В энкодере используется gated блоки, а в декодере кобминация: 3 gated, 4 dw_sep. Gated лучше раскрывают данные для обучения GNN. При обучении Teacher мы не используем никакие добавки в loss, используется обычный стандартный MAE, ведь важно, чтобы GNN учил климатологию станции.

Важно после обучения GNN проверить, насколько модель его использует (MAE при использовании GNN и MAE при подачи вместо эмбеддинга 0-вектора). Например при переобучении и недообучении будет маленькое влияние. Сейчас у меня получилась разница в ~ 1.7 градуса, что говорит о том, что эмбеддинг реально используется.

Переходим к Student. Это модель которая использует уже готовые эмбеддинги каждой станции и ориентирована на обучение уже самой EncoderDecoder архитектуры. Здесь мы используем ранжированную MAE (ближний горизонт дает больше ошибку), lam diff и std_penalty. Все реализации можно посмотреть в репозитории. Также используем меньше слоев в энкодере (так как спектральный слой позволяет нам делать RF меньше, чем размер истории) и меньше каналов в декодере. При этом используем и BiasHead и AttentionPool, они здесь не мешают. А в качестве слоев используются только dw_sep из-за своей вычислительной легкости.

Все остальное осталось без изменений.

Что в итоге получилось, что стало лучше, а что хуже

1. Модель стала легче.

2. Модель имеет базу для расширения (увеличить гиперпараметры и все).

3. Слегка ухудшились метрики на коротком горизонте, но увеличились на длинном.

Для наглядности прикладываю результаты теста baselines:

• climatology - усреднение за несколько лет.

• persistence - прогноз сейчас как в предыдущем часу.

• gru - обычный gru + linear head, статики станций передаются через линейный слой.

• tcn - обычный tcn + linear head, статики станций передаются через линейный слой.

• tcn_emb - обычный embedding слой для станций + обычный tcn + linear head

Вот метрики моей модели.

Что будет дальше?

Мои планы:

1. Обучить модель на весь мир (не отрицаю, что придется чуть редактировать модель снова).

2. Сравнить модель с другими известными time-series архитектурами (для начала возьму модели из pytorch forecasting).

По результатам этих планов будет продолжение серии постов по этому проекту.

Код и материалы

Кстати, мой проект переехал в новое пространство feyra labs 👉 edge-weather-forecast

Буду рад критике, вопросам и идеям.