Переосмысление архитектуры облачного федеративного обучения на платформе AWS

Преобразование архитектуры облачного федеративного обучения на платформе AWS

Машинное обучение (ML), особенно глубокое обучение, требует большого количества данных для повышения производительности модели. Заказчики часто должны обучать модель с данными из разных регионов, организаций или учетных записей AWS. Централизация таких данных для ML является сложной задачей из-за требований к конфиденциальности, высоких затрат на передачу данных или операционной сложности.

Федеративное обучение (FL) – это распределенный подход к ML, который обучает модели ML на распределенных наборах данных. Целью FL является повышение точности моделей ML путем использования большего количества данных, при этом сохраняя конфиденциальность и локализацию распределенных наборов данных. FL увеличивает количество данных, доступных для обучения моделей ML, особенно данных, связанных с редкими и новыми событиями, что приводит к более общей модели ML. Существующие партнерские открытые решения FL на AWS включают FedML и NVIDIA FLARE. Эти пакеты с открытым исходным кодом развертываются в облаке, работая на виртуальных машинах, без использования облачных сервисов, доступных на AWS.

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

Сначала мы рассмотрим различные подходы и проблемы FL. Затем мы покажем, как создать облачную архитектуру FL на AWS. Пример кода для создания этой архитектуры доступен на GitHub. Мы используем AWS Cloud Development Kit (AWS CDK) для развертывания архитектуры с одним щелчком развертывания. Пример кода демонстрирует сценарий, когда сервер и все клиенты принадлежат одной и той же организации (такому же учетной записи AWS), но их наборы данных не могут быть централизованы из-за требований к локализации данных. Пример кода поддерживает горизонтальное и синхронное FL для обучения моделей нейронных сетей. Используемый фреймворк ML на клиентах FL – TensorFlow.

• Шептущие модели для автоматического распознавания речи теперь доступны в Amazon SageMaker JumpStart.
• Новое – возможности генеративного искусственного интеллекта без кодирования теперь доступны в Amazon SageMaker Canvas
• Генерировать музыку с помощью MusicGen от Meta на Colab

Обзор федеративного обучения

FL обычно включает центральный сервер FL и группу клиентов. Клиенты – это вычислительные узлы, которые выполняют локальное обучение. Во время раунда обучения FL центральный сервер сначала отправляет общую глобальную модель группе клиентов. Клиенты обучают глобальную модель с использованием локальных данных, затем предоставляют локальные модели обратно серверу. Сервер агрегирует локальные модели в новую глобальную модель, затем начинает новый раунд обучения. Могут быть десятки раундов обучения до того, как глобальная модель сойдется или до достижения порогового количества раундов обучения. Поэтому FL обменивается моделями ML между центральным сервером FL и клиентами, не перемещая тренировочные данные в центральное расположение.

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

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

FL применяется в различных отраслях, таких как финансы, здравоохранение, медицина и телекоммуникации, где сохранение конфиденциальности является критическим или требуется локализация данных. FL использовалось для обучения глобальной модели для обнаружения финансовых преступлений среди нескольких финансовых учреждений. Глобальная модель превосходит модели, обученные только на локальных наборах данных, на 20%. В здравоохранении FL использовалось для прогнозирования смертности госпитализированных пациентов на основе электронных медицинских записей из нескольких больниц. Глобальная модель, предсказывающая смертность, превосходит локальные модели всех участвующих больниц. FL также использовалось для сегментации мозговых опухолей. Глобальные модели для сегментации мозговых опухолей выполняются так же, как модель, обученная собиранием распределенных наборов данных в центральном месте. В телекоммуникациях FL может быть применено для краевых вычислений, управления беспроводным спектром и ядра 5G сетей.

Существует множество других способов классификации FL:

• Горизонтальное или вертикальное – В зависимости от разделения особенностей в распределенных наборах данных, FL может быть классифицирован как горизонтальное или вертикальное. В горизонтальном FL все распределенные наборы данных имеют одинаковый набор особенностей. В вертикальном FL наборы данных имеют разные группы особенностей, требуется дополнительная коммуникация для выравнивания образцов на основе перекрывающихся особенностей.
• Синхронное или асинхронное – В зависимости от стратегии агрегации на сервере FL, FL можно классифицировать как синхронное или асинхронное. Синхронный сервер FL агрегирует локальные модели выбранного набора клиентов в глобальную модель. Асинхронный сервер FL сразу обновляет глобальную модель после получения локальной модели от клиента, тем самым сокращая время ожидания и повышая эффективность обучения.
• Централизованная или peer-to-peer – Типичная топология FL – это централизованная, где центральный сервер FL координирует набор клиентов. Другая топология FL – это peer-to-peer без централизованного сервера FL, где клиенты FL агрегируют информацию от соседних клиентов для обучения модели.

Проблемы в FL

Вы можете решить следующие проблемы с помощью алгоритмов, работающих на серверах и клиентах FL в общей архитектуре FL:

• Гетерогенность данных – Локальные данные клиентов FL могут различаться (т.е. иметь гетерогенность данных) из-за конкретного географического расположения, организаций или временных периодов. Гетерогенность данных влияет на точность глобальной модели, что приводит к большему количеству итераций обучения и длительному времени обучения. Было предложено множество решений для смягчения влияния гетерогенности данных, таких как алгоритмы оптимизации, частичное совместное использование данных между клиентами и адаптация к домену.
• Сохранение конфиденциальности – Локальные и глобальные модели могут утечь конфиденциальную информацию при атаке злоумышленника. Для FL было предложено множество подходов к сохранению конфиденциальности. Подход с безопасной агрегацией может использоваться для сохранения конфиденциальности локальных моделей, обмениваемых между серверами и клиентами FL. Локальные и глобальные подходы к дифференциальной конфиденциальности позволяют ограничить потерю конфиденциальности, добавляя шум к локальным или глобальным моделям, что обеспечивает контролируемый баланс между конфиденциальностью и точностью модели. В зависимости от требований к конфиденциальности могут использоваться комбинации различных подходов к сохранению конфиденциальности.
• Федеративная аналитика – Федеративная аналитика обеспечивает статистические измерения распределенных наборов данных без нарушения требований к конфиденциальности. Федеративная аналитика важна не только для анализа данных по распределенным наборам данных перед обучением, но и для мониторинга моделей при выполнении.

Несмотря на эти проблемы алгоритмов FL, критически важно построить безопасную архитектуру, предоставляющую операции FL от начала до конца. Одной из важных проблем при построении такой архитектуры является обеспечение простоты развертывания. Архитектура должна координировать серверы и клиенты FL для создания, обучения и развертывания моделей FL, включая непрерывную интеграцию и разработку (CI/CD) среди клиентов, отслеживаемость, аутентификацию и контроль доступа для серверов и клиентов FL. Эти функции аналогичны централизованным операциям машинного обучения (ML Ops), но их реализация более сложна из-за участия большего числа сторон. Архитектура также должна быть гибкой для реализации различных топологий FL и синхронной или асинхронной агрегации.

Обзор решения

Мы предлагаем облачную архитектуру FL на AWS, как показано на следующей диаграмме. Архитектура включает центральный сервер FL и два клиента FL. При наличии клиентов, принадлежащих разным силам, количество клиентов FL может достигать сотен. Сервер FL должен быть расположен на облачной платформе AWS, так как он состоит из набора микросервисов, предлагаемых на облаке. Клиенты FL могут быть на AWS или на стороне клиента. Клиенты FL хранят свои собственные локальные наборы данных и имеют свою собственную ИТ и систему машинного обучения для обучения моделей машинного обучения.

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

Рисунок 1 – Облачная федеративная архитектура для тренировки моделей между сервером федеративного обучения и клиентами.

Предпосылки

Для реализации этого решения вам понадобится учетная запись AWS для запуска служб центрального сервера федеративного обучения и двух клиентов. Клиентам на месте необходимо установить AWS Command Line Interface (AWS CLI), который предоставляет доступ к службам AWS на сервере федеративного обучения, включая Amazon Simple Queue Service (Amazon SQS), Amazon Simple Storage Service (Amazon S3) и Amazon DynamoDB.

Шаги федеративного обучения

В этом разделе мы рассмотрим предложенную архитектуру на Рисунке 1. На сервере федеративного обучения запускается статическая машина состояний AWS Step Functions, как показано на Рисунке 2, которая выполняет Шаги 0, 1 и 5 с Рисунка 1. Статическая машина инициирует службы AWS на сервере (Шаг 0) и выполняет итерации раундов федеративного обучения. Для каждого раунда обучения статическая машина отправляет уведомление Amazon Simple Notification Service (Amazon SNS) на тему global_model_ready вместе с задачным токеном (Шаг 1). Затем статическая машина приостанавливается и ожидает обратного вызова с задачным токеном. Есть очереди SQS, подписанные на тему global_model_ready. Каждая очередь SQS соответствует клиенту федеративного обучения и очередям уведомлений, отправленных с сервера на клиент.

Рисунок 2 – Рабочий процесс статической машины Step Functions.

Каждый клиент федеративного обучения извлекает сообщения из своей назначенной очереди SQS. При получении уведомления global_model_ready клиент загружает глобальную модель из Amazon S3 (Шаг 2) и запускает локальное обучение (Шаг 3). Локальное обучение генерирует локальную модель. Затем клиент загружает локальную модель в Amazon S3 и записывает информацию о локальной модели в таблицу DynamoDB вместе с полученным задачным токеном (Шаг 4).

Мы реализуем реестр моделей федеративного обучения с использованием Amazon S3 и DynamoDB. Мы используем Amazon S3 для хранения глобальных и локальных моделей. Мы используем таблицу DynamoDB для хранения информации о локальной модели, поскольку информация о локальной модели может отличаться в разных алгоритмах федеративного обучения и требует гибкой схемы, поддерживаемой таблицей DynamoDB.

Мы также включаем поток DynamoDB для вызова функции Lambda, так что каждый раз, когда запись записывается в таблицу DynamoDB (когда получается новая локальная модель), функция Lambda вызывается для проверки, собраны ли необходимые локальные модели (Шаг 5). Если это так, функция Lambda запускает функцию агрегации для объединения локальных моделей в глобальные модели. Полученная глобальная модель записывается в Amazon S3. Функция также отправляет обратный вызов вместе с задачным токеном, полученным из таблицы DynamoDB, в статическую машину Step Functions. Затем статическая машина определяет, должно ли федеративное обучение продолжаться с новым раундом обучения или быть остановлено на основе условия, например, достижение порога количества раундов обучения.

Каждый клиент федеративного обучения использует следующий образец кода для взаимодействия с сервером федеративного обучения. Если вы хотите настроить локальное обучение на своих клиентах федеративного обучения, функцию localTraining() можно изменить, пока возвращаемые значения будут local_model_name и local_model_info для загрузки на сервер федеративного обучения. Вы можете выбрать любую ML-платформу для обучения локальных моделей на клиентах федеративного обучения, при условии, что все клиенты используют одну и ту же ML-платформу.

 # Шаг 2: получить уведомления и имя файла модели из его очереди SQS
client.receiveNotificationsFromServer(sqs_region, client_queue_name)
# Шаг 3: загрузить глобальную модель и тренироваться локально
local_model_name, local_model_info = client.localTraining(global_model_name, s3_fl_model_registry)
# Шаг 4: загрузить локальную модель и информацию о локальной модели на сервер FL
client.uploadToFLServer(s3_fl_model_registry, local_model_name, dynamodb_table_model_info, local_model_info)

Функция Lambda для выполнения агрегационной функции на сервере имеет следующий образец кода. Если вы хотите настроить алгоритм агрегации, необходимо изменить функцию fedAvg() и вывод.

# Шаг 5: агрегировать локальные модели в функции Lambda
def lambda_handler(event, context):
    # получение task_name из события, вызванного потоком DynamoDB
    task_name = event['Records'][0]['dynamodb']['Keys']['taskName']['S']
    # извлечение транзакций из таблицы DynamoDB
    transactions = readFromFLServerTaskTable(os.environ['TASKS_TABLE_NAME'], task_name)
    # чтение информации о локальной модели от необходимых клиентов
    # token - это токен обратного вызова из машины состояний Step Functions
    local_model_info, round_id, token = receiveUpdatedModelsFromClients(transactions, task_name)
    # функция fedAvg агрегирует локальные модели в глобальную модель и сохраняет глобальную модель в S3
    global_model_name, avg_train_acc, avg_test_acc, avg_train_loss, avg_test_loss = fedAvg(local_model_info, round_id)
    # вывод отправляется в машину состояний Step Functions
    output = {'taskName': task_name, 'roundId': str(round_id), 'trainAcc': str(avg_train_acc), 'testAcc': str(avg_test_acc), 'trainLoss': str(avg_train_loss), 'testLoss': str(avg_test_loss), 'weightsFile': str(global_model_name)}
    # отправка обратного вызова в машину состояний Step Functions для сообщения о успешном завершении задачи, идентифицированной токеном
    step_client = boto3.client('stepfunctions')
    out_str = json.dumps(output)
    step_client.send_task_success(taskToken=token, output=out_str)

Эта архитектура имеет два инновационных решения. Во-первых, сервер FL использует облачные сервисы, такие как Step Functions и Lambda. Поэтому для сервера FL не требуется постоянно работающий вычислитель, что минимизирует затраты на вычисления. Во-вторых, клиенты FL извлекают сообщения из своих назначенных очередей SQS и загружают или загружают модели и информацию на сервер FL. Такой подход позволяет избежать непосредственного доступа сервера FL к ресурсам клиентов, что является важным условием для обеспечения конфиденциальности и гибкой IT и ML среды (на месте или на облаке AWS) клиентам FL.

Преимущества облачной архитектуры

Эта архитектура является облачной и обеспечивает прозрачность от начала до конца с использованием служб AWS с подтвержденной безопасностью и операционным превосходством. Например, у вас может быть клиенты с кросс-аккаунтами, которые могут предпринимать действия для доступа к ресурсам на сервере FL. Для клиентов на месте AWS CLI и AWS SDK для Python (Boto3) автоматически обеспечивают защищенное сетевое соединение между сервером FL и клиентами. Для клиентов на облачной платформе AWS вы можете использовать AWS PrivateLink и службы AWS с шифрованием данных в пути и в покое. Вы также можете использовать Amazon Cognito и AWS Identity and Access Management (IAM) для аутентификации и контроля доступа к серверам и клиентам FL. Для развертывания обученной глобальной модели вы можете использовать возможности ML Ops в Amazon SageMaker.

Облачная архитектура также позволяет интегрировать собственные фреймворки машинного обучения и алгоритмы и протоколы федеративного обучения. Например, вы можете выбрать фреймворк машинного обучения для тренировки локальных моделей на клиентах FL и настраивать различные алгоритмы агрегации как сценарии, выполняемые в функциях Lambda на сервере. Также вы можете изменять рабочие процессы в Step Functions для адаптации к различным протоколам коммуникации между сервером и клиентами.

Еще одним преимуществом облачной архитектуры является простота развертывания с помощью инструментов IaC, предлагаемых для облачных решений. Вы можете использовать AWS Cloud Development Kit (AWS CDK) и AWS CloudFormation для развертывания в один клик.

Заключение

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

Вы можете начать экспериментировать и создавать облачные архитектуры для федеративного обучения для своих случаев использования. Вы можете настроить архитектуру для поддержки различных фреймворков машинного обучения, таких как TensorFlow или PyTorch. Вы также можете настроить ее для поддержки различных алгоритмов федеративного обучения, включая асинхронное федеративное обучение, алгоритмы агрегации и алгоритмы дифференциальной конфиденциальности. Вы можете активировать эту архитектуру с помощью операций FL Ops, используя возможности ML Ops в Amazon SageMaker.