Привет, Хабр! После недавнего релиза Apple VIsion Pro я задумался над тем, какой вклад компания внесла в open-source ML, ведь, очевидно, было много предпосылок, наработок и технологий, о которых массовый пользователь не знает. Но все это гарантировано сыграло роль в появлении тех продуктов на рынке, о которых слышал буквально каждый в мире. Поэтому свою первую статью я решил сделать в формате небольшого обзора того, чем занимались и продолжают заниматься Apple в ML, хотя, казалось бы, что там можно смотреть в продуктах с открытым кодом кроме сотни репозиториев про Swift. Попытаюсь выделить основные тренды за прошедшие года и поделиться тем, что сам нашел интересного (оставив позади вопрос о целесообразности развития ML на яблочных устройствах).

Чем занимается Apple в ML?

На мой взгляд, как и везде, у компании можно выделить два крупных направления развития: research (зачастую в области CV, хотя иногда есть небольшой байес в мультимодальность и LLM) и production, который нацелен на оптимизацию применения ML продуктов на MacOS/IOS системах.

У Apple достаточно много репозиториев, которые маркируются префиксом ml- и, судя по звездочкам, даже пользуются хорошим спросом. При этом существует и отдельная AI лаборатория, занимающаяся в большей степени прикладными задачами и разработкой своих фрэймворков, т.е. аналогов torch/tensorflow, но которые нацелены на инференс на устройствах с M-чипами (напомню, что в том же Vision Pro стоят M3 и R1 чипы, где второй является представителем новой линейки). Давайте по порядку!

Research

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

ml-ferret
Открытый код для свежей статьи (релиз пару месяцев назад) о мультимодальной LLM, которая имеет очень хорошее представление о пространственном расположении на фотографиях. Например, подаешь инпут картинку с приблизительно выделенным объектом и спрашиваешь текстом "что это?" или "как мне использовать это для моей задачи?", а модель сегментирует объект в бокс и выдает тебе контекстно связанный хороший ответ.

Примеры решаемых задач (взяты из train датасета)

Буквально ставишь точку на бутерброд, спрашиваешь, что это и как его сделать, а в ответ получаешь пояснение рецепта и выделение продуктов, которые заодно попали на фото и могут пригодиться. В качестве энкодера картинок взяли CLIP-ViT-L/14 и предложили целый метод по работе с выделенными пользователями областями, назвав его Spatial-aware visual sampler (по факту там последовательность из knn, fusion и pooling). Модель классная и нашла большой отклик в аудитории, о чем уже можно судить по 7 тысячам звездочкам на репозитории.

ml-fastvit
Эта статья вышла почти год назад, но идейно актуальна до сих пор. У Apple много что было об улучшении трансформенных архитектур для CV, ускорении инференса и обучении с уменьшением потребления памяти. Остановился именно на этой работе, т.к. получили хороший буст к результатам, замерив потом все бэнчмарки на iPhone 12 Pro. Команда просто пересмотрела подход RepMixer и провела вдумчивую репараметризацию skip-connection блока, поменяв порядок операций (BN - BatchNorm, DWConv - depthwise convolution, σ - нелинейная функция активации):

Было:
Стало:

На деле там гораздо больше тонкостей (и с positional encodings, и с train процессом, и с ядрами сверток, и т.д. и т.п.). Можно придумывать новые способы оптимизации, а можно провести некоторую вдумчивую реструктуризацию существующей архитектуры, что мне как подход очень нравится. Видимо я так не один думаю, раз есть 2 тысячи звездочек.

ml-hypersim
Не совсем статья, но, оказывается, у Apple даже есть свой открытый датасет на 78k супер фотореалистичных синтетических изображений различных вариантов 500 сцен внутри помещений. Приятный плюс - относительно удобное использование с кучей вариантов команд и мини-гайдов инференса. Очевидно, подобные вещи сыграли не последнюю роль в разработке Vision Pro.

Примеры одной и той же сцены из датасета

ml-neuman
Релиз был еще в 22-ом году. Представили фреймворк для фотореалистичного рендеринга людей в пространстве по одному видео. Обучали NeRF модель как для людей, так и для сцен, со всеми нюансами грубой оценки геометрии человека, что в итоге позволяет из 10-секундного видео сгенерировать новые позы, ракурсы с фонами и т.д. Добавили предобработку видеорядов (например, COLMAP), подобрали разные лоссы для моделей генерации людей и сцен, а потом обучили.

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

Используя приблизительные геометрические оценки положения человека, сделали переходное поле из пространства наблюдения в каноническое пространство. Кажется, что подобные вещи очень актуальны еще со времен Apple Face ID, да и в стезе развития виртуальных аватаров тоже полезно.

Production

Здесь больше упор идет на имплементацию каких-то готовых ML-решений в уже существующие продукты или просто на оптимизацию всего цикла ПО на устройствах с M-чипами, IOS-девайсах и т.д. Все-таки у тех же Mac-ов со своими мощностями будто есть потенциал в "домашнем" ML (лично тюнил gpt от сбера на своем MacBook Pro 2021). На хабре есть хороший обзор ML-инструментов от Apple, где уделяется много времени Create ML со всеми примерами обучения моделей, так что на нем подробно останавливаться не буду.

Существенная часть разработок концентрируется вокруг Core ML экосистемы, которая представляет удобное API для задач Voice/Speech/NLP/CV. Честно отмечу, что выбор готовых моделей достаточно скромный и ограничивается одним вариантом BERT-а для текстов и YOLOv3 для картинок, когда буквально на днях уже вышла YOLOv9. Если хотите больше, то уже приходится использовать более нишевые инструменты.

Есть целый user-friendly инструмент, который на ходу позволяет решать типичные ML задачи

ml-stable-diffusion
Например, когда-то появилась известная всем модель генерации картинок по текстовым промптам, а люди взяли и заставили ее запускаться и работать на своих M-чипах и MacOS Silicon (даже на некоторых IOS можно) с возможностью сжатия весов, добавлением всяких фишек по типу ControlNet и небольшой библиотеки под любимый Swift.

coremltools
Фреймворк на питоне для конвертации ML моделей (есть поддержка PyTorch, TensorFlow, scikit-learn, XGBoost, LibSVM) для их последующего использования в Core ML системах/приложениях на IOS, MacOS и т.д. Полезно для разработчиков на яблоке, как раз помогает расширить изначальный скудный выбор.

ml-cvnets
Сделали небольшой инструмент для тренировки уже готовых CV-архитектур, начиная ImageNet и мобильными имплементациями сеток и заканчивая трансформерами (насчитал штук 20 всего моделей). Опять же все это можно делать и на MacOS устройствах, т.к. все сделано на torch, а в примерах ребята сами предлагают потом использовать свою же конвертацию в CoreML. Все это должно хорошо работать в связке с Xсode, но я лично не проверял, да и обхожу стороной это место...

ml-ane-transformers
Немного старый фреймворк для оптимизации запусков сеток с hugging face на M1 и ранее моделях. Вообще есть отдельная библиотека под tensorflow для M-чипов, но все эти направления сейчас очевидно загнулись, ибо теперь и tensorflow, и torch не без вклада всех этих проектов уже нативно поддерживают бэкенд metal mps, который как раз для M-чипов (хотя и есть некоторые баги или ограничения, вплоть до не релизнутой еще версии torch 2.3+). Поэтому все такие разработки подзабросили пару лет назад, либо в архив кинули, пересмотрев свой подход решения проблемы совместимости Apple-устройств с ML.

Сравнение инференса модели distilbert-base-uncased-finetuned-sst-2-english

mlx
Пожалуй, мой фаворит - открытая ML/DL библиотека на python, как аналог numpy и torch (интерфейс действительно очень похож и интуитивен), все написано на C++. Есть свое API на Swift и C, можно найти большое количество официальных примеров использования, начиная линейными моделями и заканчивая запуском whisper или реализацией bert-а, да и хорошая документация тоже на месте. Разработкой библиотеки занимается отдельный ресечерский центр, который упоминался в самом начале поста.

Официальный пример кода энкодера у VAE для MNIST:

import math
import mlx.core as mx
import mlx.nn as nn

class Encoder(nn.Module):
    def __init__(self, num_latent_dims, image_shape, max_num_filters):
        super().__init__()

        # number of filters in the convolutional layers
        num_filters_1 = max_num_filters // 4
        num_filters_2 = max_num_filters // 2
        num_filters_3 = max_num_filters

        # Output (BHWC):  B x 32 x 32 x num_filters_1
        self.conv1 = nn.Conv2d(image_shape[-1], num_filters_1, 3, stride=2, padding=1)
        # Output (BHWC):  B x 16 x 16 x num_filters_2
        self.conv2 = nn.Conv2d(num_filters_1, num_filters_2, 3, stride=2, padding=1)
        # Output (BHWC):  B x 8 x 8 x num_filters_3
        self.conv3 = nn.Conv2d(num_filters_2, num_filters_3, 3, stride=2, padding=1)

        # Batch Normalization
        self.bn1 = nn.BatchNorm(num_filters_1)
        self.bn2 = nn.BatchNorm(num_filters_2)
        self.bn3 = nn.BatchNorm(num_filters_3)

        # Divide the spatial dimensions by 8 because of the 3 strided convolutions
        output_shape = [num_filters_3] + [dimension // 8 for dimension in image_shape[:-1]]

        flattened_dim = math.prod(output_shape)

        # Linear mappings to mean and standard deviation
        self.proj_mu = nn.Linear(flattened_dim, num_latent_dims)
        self.proj_log_var = nn.Linear(flattened_dim, num_latent_dims)

    def __call__(self, x):
        x = nn.leaky_relu(self.bn1(self.conv1(x)))
        x = nn.leaky_relu(self.bn2(self.conv2(x)))
        x = nn.leaky_relu(self.bn3(self.conv3(x)))
        x = mx.flatten(x, 1)  # flatten all dimensions except batch

        mu = self.proj_mu(x)
        logvar = self.proj_log_var(x)
        # Ensure this is the std deviation, not variance
        sigma = mx.exp(logvar * 0.5)

        # Generate a tensor of random values from a normal distribution
        eps = mx.random.normal(sigma.shape)

        # Reparametrization trick to brackpropagate through sampling.
        z = eps * sigma + mu

        return z, mu, logvar

Радует, что проект активно развивается. Буквально за последних два месяца зимы прошел цикл от 0.1.n до 0.4.n версии. Фреймворк рассчитан на Apple Silicon OS, обладает хорошей нативной поддержкой бэкенда mps, как обычно lazy computation, есть несколько девайсов, да и про dynamic graph construction тоже не забыли. Нет чего-то нового, чего не хватает в том же torch, которым пока можно пользоваться и на Mac-устройствах, но, как мне кажется, за проектом есть хороший потенциал, возможно нужно только подождать.

Итого

По понятным причинам большая ставка делается на CV и мультимодальность вокруг этой области, все-таки релиз Vision Pro, развитие Face ID и лидерство по качеству камер на своих девайсах невольно задают такой тренд. А все технические решения для разработчиков фокусируются вокруг создания удобной среды проектирования своих полноценных ML систем. Все это не означает, что Apple закрывает глаза на другие области и в research разделе можно найти много статьей как про NLP, так и просто про алгоритмы оптимизации, но будто это не главная траектория развития. Каждый месяц от исследователей компании выходит дюжина новых статей, далеко за примером ходить не надо, вот в январе выпустили Acoustic Model Fusion for End-to-end Speech Recognition или же статья с громким названием One Wide Feedforward is All You Need, где разбирается роль FFN в трансформерных архитектурах. Просто у пользователей будто спрос на CV конкретно в этом случае гораздо больше, что как минимум отражается в статистике звездочек, коммитов, форков и реквестов в соответствующих репозиториях компании.

Обидно немного за Siri, хочется какой-нибудь гонки и в LLM стезе, быть может компания и об этом задумается.

Open-source ML от Apple