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

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

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

Введение

Допустим, вы только что разработали новую классную модель машинного обучения (конечно, с помощью Flux.jl). И теперь хотите начать развёртывать её для своих пользователей. Как вы будете это делать? Наверное, проще всего будет отдать модель пользователям и позволить прогонять её локально на их данных. Но у такого подхода есть недостатки:

1. Модели машинного обучения велики, и пользовательские компьютеры могут не иметь достаточного количества вычислительных или дисковых ресурсов.
2. Модели машинного обучения часто обновляются, и вам может быть неудобно регулярно рассылать по сети большие объёмы данных.
   
3. Разработка моделей занимает много времени и требует большого количества вычислительных ресурсов. И вы можете захотеть компенсацию за это в виде платы за пользование вашей моделью.

Затем обычно вспоминают о том, что модель можно предоставлять в облаке через API. За последние несколько лет появилось много таких сервисов, каждая крупная облачная платформа предлагает подобные услуги корпоративным разработчикам. Но потенциальные пользователи сталкиваются с очевидной дилеммой: теперь их данные обрабатываются на удалённом сервере, который может не заслуживать доверия. У этого есть явные этические и юридические последствия, ограничивающие применение подобных сервисов. В регулируемых законодательством отраслях, особенно в здравоохранении и финансовых услугах, часто нельзя отправлять данные пациентов и клиентов на обработку третьим лицам.

Есть другие варианты?

Оказывается, есть! Недавние открытия в криптографии позволяют проводить вычисления с данными без их расшифровки. Например, пользователь отправляет зашифрованные данные (скажем, изображения) в облачный API, который запускает модель машинного обучения, а затем отправляет зашифрованный ответ. Ни на одном из этапов данные не расшифровываются, облачный провайдер не получает доступа к исходным изображениям и не может расшифровать вычисленный прогноз. Как такое возможно? Давайте выясним это на примере создания сервиса для распознавания рукописного текста на зашифрованных изображениях из датасета MNIST.

О гомоморфном шифровании

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

• pub_key, eval_key, priv_key = keygen()
  
• encrypted = encrypt(pub_key, plaintext)
  
• decrypted = decrypt(priv_key, encrypted)
  
• encrypted′ = eval(eval_key, f, encrypted)
  

Первые три операции просты и знакомы всем, кто уже использовал любые алгоритмы асимметричного шифрования (например, если вы подключились через TLS). Вся магия происходит в последней операции. При шифровании она оценивает функцию f и возвращает другое зашифрованное значение, вычисленное в соответствии с результатом оценки f на зашифрованном значении. Эта особенность и дала такому подходу его название. Оценка связана с операцией шифрования:

f(decrypt(priv_key, encrypted)) == decrypt(priv_key, eval(eval_key, f, encrypted))

Аналогично, можно с помощью зашифрованного значения оценить произвольные гомоморфизмы f.

Какие поддерживаются функции f, зависит от криптографических схем и поддерживаемых операций. Если поддерживается лишь одна f (например, f = +), то схема называется «частично гомоморфной». Если f может быть любым полным набором шлюзов, на основе которого можно создать произвольные схемы, то при ограниченном размере схемы это называется другой разновидностью частично гомоморфного вычисления — «в некотором роде гомоморфное», а при неограниченном размере — «полностью гомоморфным» вычислением. Можно превратить «в некотором роде» в полностью гомоморфное шифрование с помощью методики bootstrapping, но это выходит за рамки нашей статьи. Полностью гомоморфное шифрование — относительно недавнее открытие, первая рабочая схема (хотя и непрактичная) была опубликована Крейгом Джентри в 2009-м. Есть и ряд более поздних (и практичных) полностью гомоморфных схем. Также есть программные пакеты, которые качественно реализуют эти схемы. Чаще всего используют Microsoft SEAL и PALISADE. Кроме того, недавно я открыл код реализации этих алгоритмов Pure Julia. Для этой статьи мы будем использовать реализованное в ней шифрование CKKS.

Обзор CKKS

CKKS (по именам авторов научной работы Cheon-Kim-Kim-Song, предложивших алгоритм в 2016-м) — это схема гомоморфного шифрования, позволяющая гомоморфно оценивать следующие примитивные операции:

• Поэлементное сложение длин n векторов комплексных чисел.
  
• Поэлементное умножение длин n комплексных векторов.
  
• Поворачивание (в контексте circshift) элементов в векторе.
  
• Комплексное сопряжение векторных элементов.
  

Параметр n зависит от желаемого уровня безопасности и точности, и обычно он довольно высок. В нашем примере он будет равен 4096 (более высокое значение повышает безопасность, но и тяжелее в вычислениях, масштабируется примерно как n log n).

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

С другой стороны, подобные ограничения не являются чем-то необычным для разработчиков пакетов для машинного обучения. Специальные ускорители вроде GPU тоже обычно оперируют векторами чисел. Кроме того, для многих разработчиков числа с плавающей запятой иногда выглядят шумными из-за влияния алгоритмов выбора, многопоточности и так далее. Хочу подчеркнуть, что ключевым отличием здесь является то, что арифметические вычисления с числами с плавающей запятой являются изначально детерминистскими, даже если это не очевидно из-за сложности реализации, хотя примитивы CKKS действительно являются шумными. Но, быть может, это позволяет пользователям понять, что шумность не так страшна, как может показаться.

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

julia> using ToyFHE
# Let's play with 8 element vectors
julia> N = 8;
# Choose some parameters - we'll talk about it later
julia> ℛ = NegacyclicRing(2N, (40, 40, 40))
ℤ₁₃₂₉₂₂₇₉₉₇₅₆₈₀₈₁₄₅₇₄₀₂₇₀₁₂₀₇₁₀₄₂₄₈₂₅₇/(x¹⁶ + 1)
# We'll use CKKS
julia> params = CKKSParams(ℛ)
CKKS parameters
# We need to pick a scaling factor for a numbers - again we'll talk about that later
julia> Tscale = FixedRational{2^40}
FixedRational{1099511627776,T} where T
# Let's start with a plain Vector of zeros
julia> plain = CKKSEncoding{Tscale}(zero(ℛ))
8-element CKKSEncoding{FixedRational{1099511627776,T} where T} with indices 0:7:
 0.0 + 0.0im
 0.0 + 0.0im
 0.0 + 0.0im
 0.0 + 0.0im
 0.0 + 0.0im
 0.0 + 0.0im
 0.0 + 0.0im
 0.0 + 0.0im
# Ok, we're ready to get started, but first we'll need some keys
julia> kp = keygen(params)
CKKS key pair
julia> kp.priv
CKKS private key
julia> kp.pub
CKKS public key
# Alright, let's encrypt some things:
julia> foreach(i->plain[i] = i+1, 0:7); plain
8-element CKKSEncoding{FixedRational{1099511627776,T} where T} with indices 0:7:
 1.0 + 0.0im
 2.0 + 0.0im
 3.0 + 0.0im
 4.0 + 0.0im
 5.0 + 0.0im
 6.0 + 0.0im
 7.0 + 0.0im
 8.0 + 0.0im
julia> c = encrypt(kp.pub, plain)
CKKS ciphertext (length 2, encoding CKKSEncoding{FixedRational{1099511627776,T} where T})
# And decrypt it again
julia> decrypt(kp.priv, c)
8-element CKKSEncoding{FixedRational{1099511627776,T} where T} with indices 0:7:
 0.9999999999995506 - 2.7335193113350057e-16im
 1.9999999999989408 - 3.885780586188048e-16im
  3.000000000000205 + 1.6772825551165524e-16im
  4.000000000000538 - 3.885780586188048e-16im
  4.999999999998865 + 8.382500573679615e-17im
  6.000000000000185 + 4.996003610813204e-16im
  7.000000000001043 - 2.0024593503998215e-16im
  8.000000000000673 + 4.996003610813204e-16im
# Note that we had some noise. Let's go through all the primitive operations we'll need:
julia> decrypt(kp.priv, c+c)
8-element CKKSEncoding{FixedRational{1099511627776,T} where T} with indices 0:7:
 1.9999999999991012 - 5.467038622670011e-16im
 3.9999999999978817 - 7.771561172376096e-16im
   6.00000000000041 + 3.354565110233105e-16im
  8.000000000001076 - 7.771561172376096e-16im
   9.99999999999773 + 1.676500114735923e-16im
  12.00000000000037 + 9.992007221626409e-16im
 14.000000000002085 - 4.004918700799643e-16im
 16.000000000001346 + 9.992007221626409e-16im
julia> csq = c*c
CKKS ciphertext (length 3, encoding CKKSEncoding{FixedRational{1208925819614629174706176,T} where T})
julia> decrypt(kp.priv, csq)
8-element CKKSEncoding{FixedRational{1208925819614629174706176,T} where T} with indices 0:7:
 0.9999999999991012 - 2.350516767363621e-15im
 3.9999999999957616 - 5.773159728050814e-15im
  9.000000000001226 - 2.534464540987068e-15im
 16.000000000004306 - 2.220446049250313e-15im
  24.99999999998865 + 2.0903753311370056e-15im
  36.00000000000222 + 4.884981308350689e-15im
 49.000000000014595 + 1.0182491378134327e-15im
  64.00000000001077 + 4.884981308350689e-15im

Вот так просто! Внимательный читатель мог заметить, что CSQ немного отличается от предыдущего шифротекста. В частности, у шифротекста «длина 3» и масштаб куда больше. Объяснение, что это и для чего нужно, выходит за рамки статьи. Достаточно сказать, что нам нужно снизить значения, прежде чем продолжать вычисления, иначе в шифротексте кончится «место». К счастью, мы можем уменьшить каждое из двух выросших значений:

# To get back down to length 2, we need to `keyswitch` (aka
# relinerarize), which requires an evaluation key. Generating
# this requires the private key. In a real application we would
# have generated this up front and sent it along with the encrypted
# data, but since we have the private key, we can just do it now.
julia> ek = keygen(EvalMultKey, kp.priv)
CKKS multiplication key
julia> csq_length2 = keyswitch(ek, csq)
CKKS ciphertext (length 2, encoding CKKSEncoding{FixedRational{1208925819614629174706176,T} where T})
# Getting the scale back down is done using modswitching.
julia> csq_smaller = modswitch(csq_length2)
CKKS ciphertext (length 2, encoding CKKSEncoding{FixedRational{1.099511626783e12,T} where T})
# And it still decrypts correctly (though note we've lost some precision)
julia> decrypt(kp.priv, csq_smaller)
8-element CKKSEncoding{FixedRational{1.099511626783e12,T} where T} with indices 0:7:
 0.9999999999802469 - 5.005163520332181e-11im
 3.9999999999957723 - 1.0468514951188039e-11im
  8.999999999998249 - 4.7588542623100616e-12im
 16.000000000023014 - 1.0413447889166631e-11im
 24.999999999955193 - 6.187833723406491e-12im
 36.000000000002345 + 1.860733715346631e-13im
  49.00000000001647 - 1.442396043149794e-12im
 63.999999999988695 - 1.0722489563648028e-10im

Кроме того, modswitching (сокращение от modulus switching, переключение модуля) уменьшает размер модуля шифротекста, так что мы не можем продолжать делать это бесконечно долго (мы используем схему somewhat-гомоморфного шифрования):

julia> ℛ # Remember the ring we initially created
ℤ₁₃₂₉₂₂₇₉₉₇₅₆₈₀₈₁₄₅₇₄₀₂₇₀₁₂₀₇₁₀₄₂₄₈₂₅₇/(x¹⁶ + 1)
julia> ToyFHE.ring(csq_smaller) # It shrunk!
ℤ₁₂₀₈₉₂₅₈₂₀₁₄₄₅₉₃₇₇₉₃₃₁₅₅₃/(x¹⁶ + 1)</code>
Осталось выполнить последнюю операцию — повороты (rotations). Как и в случае с keyswitch, для этого нам потребуется ключ оценки (evaluation key, также известный как ключ Галуа):

<source lang="julia">julia> gk = keygen(GaloisKey, kp.priv; steps=2)
CKKS galois key (element 25)
julia> decrypt(circshift(c, gk))
decrypt(kp, circshift(c, gk))
8-element CKKSEncoding{FixedRational{1099511627776,T} where T} with indices 0:7:
  7.000000000001042 + 5.68459112632516e-16im
  8.000000000000673 + 5.551115123125783e-17im
  0.999999999999551 - 2.308655353580721e-16im
 1.9999999999989408 + 2.7755575615628914e-16im
  3.000000000000205 - 6.009767921608429e-16im
  4.000000000000538 + 5.551115123125783e-17im
  4.999999999998865 + 4.133860996136768e-17im
  6.000000000000185 - 1.6653345369377348e-16im
# And let's compare to doing the same on the plaintext
julia> circshift(plain, 2)
8-element OffsetArray(::Array{Complex{Float64},1}, 0:7) with eltype Complex{Float64} with indices 0:7:
 7.0 + 0.0im
 8.0 + 0.0im
 1.0 + 0.0im
 2.0 + 0.0im
 3.0 + 0.0im
 4.0 + 0.0im
 5.0 + 0.0im
 6.0 + 0.0im

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

Модель машинного обучения

Если вы не знакомы с машинным обучением или библиотекой Flux.jl, то рекомендую кратко пробежаться по документации Flux.jl или посмотреть бесплатное введение в машинное обучение, потому что мы будем обсуждать только изменения в применении модели к шифрованным данным.

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

function reshape_and_vcat(x)
    let y=reshape(x, 64, 4, size(x, 4))
        vcat((y[:,i,:] for i=axes(y,2))...)
    end
end
model = Chain(
    # First convolution, operating upon a 28x28 image
    Conv((7, 7), 1=>4, stride=(3,3), x->x.^2),
    reshape_and_vcat,
    Dense(256, 64, x->x.^2),
    Dense(64, 10),
)

Это та же модель, что и в работе «Secure Outsourced Matrix Computation and Application to Neural Networks», которая использует ту же криптографическую схему с двумя отличиями: 1) ради простоты мы не стали шифровать саму модель, и 2) после каждого слоя у нас используются байесовские векторы (во Flux это делается по умолчанию), не уверен, что это было в упомянутой работе. Быть может, из-за второго пункта точность на тестовом наборе у нашей модели оказалась чуть выше (98,6 % против 98,1 %), но причиной могут быть и гиперпараметрические различия.

Необычным (для тех, у кого есть опыт в машинном обучении) является x.^2-активация функций. Чаще всего в подобных случаях используют tanh, relu или что-нибудь более причудливое. Но хотя эти функции (особенно relu) легко вычисляются для обычных текстовых значений, однако для их оценки в шифрованном виде может потребоваться много вычислительных ресурсов (обычно мы оцениваем полиномиальное приближение). К счастью, в данном случае x.^2 прекрасно работает.

В остальном цикл обучения остался таким же. Мы удалили из модели softmax ради loss-функции logitcrossentropy (её можно было оставить и оценить softmax после расшифровки на клиенте). Полный код для обучения модели лежит на GitHub, он выполняется за несколько минут на любой свежей видеокарте.

Эффективное выполнение операций

Теперь мы знаем, какие операции нам нужно выполнить:

• Свёртывание.
  
• Поэлементное возведение в квадрат.
  
• Матричное умножение.
  

С возведением в квадрат всё просто, мы уже рассмотрели его выше, так что рассмотрим две другие операции. Будем считать, что длина пакета данных равна 64 (вы могли заметить, что параметры модели и размер пакета выбраны так, чтобы использовать преимущества 4096-элементного вектора, который мы получили в результате реалистического выбора параметров).

Свёртывание

Вспомним, как работает свёртывание. Возьмём окно (в нашем случае 7х7) исходного входного массива, и каждый элемент окна умножим на элемент свёрточной маски. Затем передвинем окно на какой-то шаг (в нашем случае шаг равен 3, то есть двигаем на 3 элемента) и повторяем процесс (с той же свёрточной маской). Ниже показан анимация процесса (источник) для свёртки 3x3 с шагом (2, 2) (синий массив — входной, зелёный — выходной):

Кроме того, мы выполняем свёртки в четырёх разных «каналах» (то есть повторяем свёртывание ещё 3 раза с разными масками).

Теперь мы знаем, что нужно делать, осталось понять, как. Нам повезло, что свёртка — первая операция в нашей модели. В результате, чтобы сэкономить ресурсы, мы можем предварительно обработать данные на клиенте, а потом шифровать их (без использования весов). Сделаем вот что:

• Предварительно вычислим каждое свёрточное окно (то есть, выборку 7x7 из исходных изображений), что дает нам 64 матрицы 7x7 на каждое входное изображение. Обратите внимание, что для окна 7x7 с шагом 2 будут свёрточные окна 8x8 для оценки входного изображения 28x28.
  
• Соберём в один вектор одинаковые позиции в каждом окне. То есть, для каждого изображения у нас будет 64-элементный вектор, либо вектор из 64х64 элементов для пакета размером 64 (всего 49 матриц 64х64).
  
• Зашифруем.
  

Затем свертывание просто превращается в скалярное умножение всей матрицы с соответствующим элементом маски. А просуммировав позднее все 49 элементов, мы получим результат свёртывания. Вот как может выглядеть реализация этой стратегии (в виде обычного текста):

function public_preprocess(batch)
    ka = OffsetArray(0:7, 0:7)
    # Create feature extracted matrix
    I = [[batch[i′*3 .+ (1:7), j′*3 .+ (1:7), 1, k] for i′=ka, j′=ka] for k = 1:64]
    # Reshape into the ciphertext
    Iᵢⱼ = [[I[k][l...][i,j] for k=1:64, l=product(ka, ka)] for i=1:7, j=1:7]
end
Iᵢⱼ = public_preprocess(batch)
# Evaluate the convolution
weights = model.layers[1].weight
conv_weights = reverse(reverse(weights, dims=1), dims=2)
conved = [sum(Iᵢⱼ[i,j]*conv_weights[i,j,1,channel] for i=1:7, j=1:7) for channel = 1:4]
conved = map(((x,b),)->x .+ b, zip(conved, model.layers[1].bias))

Это (модуль изменения размерности) (по модулю — изменение порядка размеров) дает такой же ответ, как и операции model.layers[1](batch).

Добавим операции шифрования:

Iᵢⱼ = public_preprocess(batch)
C_Iᵢⱼ = map(Iᵢⱼ) do Iij
    plain = CKKSEncoding{Tscale}(zero(plaintext_space(ckks_params)))
    plain .= OffsetArray(vec(Iij), 0:(N÷2-1))
    encrypt(kp, plain)
end
weights = model.layers[1].weight
conv_weights = reverse(reverse(weights, dims=1), dims=2)
conved3 = [sum(C_Iᵢⱼ[i,j]*conv_weights[i,j,1,channel] for i=1:7, j=1:7) for channel = 1:4]
conved2 = map(((x,b),)->x .+ b, zip(conved3, model.layers[1].bias))
conved1 = map(ToyFHE.modswitch, conved2)

Обратите внимание, что здесь не требуется keyswitch, потому что веса являются публичными. Так что мы не увеличиваем длину шифротекста.

Матричное умножение

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

function matmul_square_reordered(weights, x)
    sum(1:size(weights, 1)) do k
        # We rotate the columns of the LHS and take the diagonal
        weight_diag = diag(circshift(weights, (0,(k-1))))
        # We rotate the rows of the RHS
        x_rotated = circshift(x, (k-1,0))
        # We do an elementwise, broadcast multiply
        weight_diag .* x_rotated
    end
end
function matmul_reorderd(weights, x)
    sum(partition(1:256, 64)) do range
        matmul_square_reordered(weights[:, range], x[range, :])
    end
end
fc1_weights = model.layers[3].W
x = rand(Float64, 256, 64)
@assert (fc1_weights*x) ≈ matmul_reorderd(fc1_weights, x)

Конечно, для общего матричного умножения потребуется что-то посложнее, но пока этого достаточно.

Улучшаем методику

Теперь все компоненты нашей методики работают. Вот код целиком (за исключением настройки параметров выбора и подобных вещей):

ek = keygen(EvalMultKey, kp.priv)
gk = keygen(GaloisKey, kp.priv; steps=64)
Iᵢⱼ = public_preprocess(batch)
C_Iᵢⱼ = map(Iᵢⱼ) do Iij
    plain = CKKSEncoding{Tscale}(zero(plaintext_space(ckks_params)))
    plain .= OffsetArray(vec(Iij), 0:(N÷2-1))
    encrypt(kp, plain)
end
weights = model.layers[1].weight
conv_weights = reverse(reverse(weights, dims=1), dims=2)
conved3 = [sum(C_Iᵢⱼ[i,j]*conv_weights[i,j,1,channel] for i=1:7, j=1:7) for channel = 1:4]
conved2 = map(((x,b),)->x .+ b, zip(conved3, model.layers[1].bias))
conved1 = map(ToyFHE.modswitch, conved2)
Csqed1 = map(x->x*x, conved1)
Csqed1 = map(x->keyswitch(ek, x), Csqed1)
Csqed1 = map(ToyFHE.modswitch, Csqed1)
function encrypted_matmul(gk, weights, x::ToyFHE.CipherText)
    result = repeat(diag(weights), inner=64).*x
    rotated = x
    for k = 2:64
        rotated = ToyFHE.rotate(gk, rotated)
        result += repeat(diag(circshift(weights, (0,(k-1)))), inner=64) .* rotated
    end
    result
end
fq1_weights = model.layers[3].W
Cfq1 = sum(enumerate(partition(1:256, 64))) do (i,range)
    encrypted_matmul(gk, fq1_weights[:, range], Csqed1[i])
end
Cfq1 = Cfq1 .+ OffsetArray(repeat(model.layers[3].b, inner=64), 0:4095)
Cfq1 = modswitch(Cfq1)
Csqed2 = Cfq1*Cfq1
Csqed2 = keyswitch(ek, Csqed2)
Csqed2 = modswitch(Csqed2)
function naive_rectangular_matmul(gk, weights, x)
    @assert size(weights, 1) < size(weights, 2)
    weights = vcat(weights, zeros(eltype(weights), size(weights, 2)-size(weights, 1), size(weights, 2)))
    encrypted_matmul(gk, weights, x)
end
fq2_weights = model.layers[4].W
Cresult = naive_rectangular_matmul(gk, fq2_weights, Csqed2)
Cresult = Cresult .+ OffsetArray(repeat(vcat(model.layers[4].b, zeros(54)), inner=64), 0:4095)

Выглядит не слишком аккуратно, но если ты всё это сделал, то должен понимать каждый шаг.
Теперь давайте подумаем, какие абстракции могли бы упростить нам жизнь. Мы покидаем область картографии и машинного обучения и переходим к архитектуре языка программирования, поэтому давайте воспользуемся тем, что Julia позволяет использовать и создавать мощные абстракции. Например, можно инкапсулировать весь процесс извлечения свёрток в свой тип массива:

using BlockArrays
"""
    ExplodedConvArray{T, Dims, Storage} <: AbstractArray{T, 4}
Represents a an `nxmx1xb` array of images, but rearranged into a
series of convolution windows. Evaluating a convolution compatible
with `Dims` on this array is achievable through a sequence of
scalar multiplications and sums on the underling storage.
"""
struct ExplodedConvArray{T, Dims, Storage} <: AbstractArray{T, 4}
    # sx*sy matrix of b*(dx*dy) matrices of extracted elements
    # where (sx, sy) = kernel_size(Dims)
    #       (dx, dy) = output_size(DenseConvDims(...))
    cdims::Dims
    x::Matrix{Storage}
    function ExplodedConvArray{T, Dims, Storage}(cdims::Dims, storage::Matrix{Storage}) where {T, Dims, Storage}
        @assert all(==(size(storage[1])), size.(storage))
        new{T, Dims, Storage}(cdims, storage)
    end
end
Base.size(ex::ExplodedConvArray) = (NNlib.input_size(ex.cdims)..., 1, size(ex.x[1], 1))
function ExplodedConvArray{T}(cdims, batch::AbstractArray{T, 4}) where {T}
    x, y = NNlib.output_size(cdims)
    kx, ky = NNlib.kernel_size(cdims)
    stridex, stridey = NNlib.stride(cdims)
    kax = OffsetArray(0:x-1, 0:x-1)
    kay = OffsetArray(0:x-1, 0:x-1)
    I = [[batch[i′*stridex .+ (1:kx), j′*stridey .+ (1:ky), 1, k] for i′=kax, j′=kay] for k = 1:size(batch, 4)]
    Iᵢⱼ = [[I[k][l...][i,j] for k=1:size(batch, 4), l=product(kax, kay)] for (i,j) in product(1:kx, 1:ky)]
    ExplodedConvArray{T, typeof(cdims), eltype(Iᵢⱼ)}(cdims, Iᵢⱼ)
end
function NNlib.conv(x::ExplodedConvArray{<:Any, Dims}, weights::AbstractArray{<:Any, 4}, cdims::Dims) where {Dims<:ConvDims}
    blocks = reshape([ Base.ReshapedArray(sum(x.x[i,j]*weights[i,j,1,channel] for i=1:7, j=1:7), (NNlib.output_size(cdims)...,1,size(x, 4)), ()) for channel = 1:4 ],(1,1,4,1))
    BlockArrays._BlockArray(blocks, BlockArrays.BlockSizes([8], [8], [1,1,1,1], [64]))
end

Здесь мы снова использовали BlockArrays для представления массива 8x8x4x64 в виде четырёх массивов 8x8x1x64 как в исходном коде. Теперь представление первого этапа стало гораздо красивее, хотя бы при незашифрованных массивах:

julia> cdims = DenseConvDims(batch, model.layers[1].weight; stride=(3,3), padding=(0,0,0,0), dilation=(1,1))
DenseConvDims: (28, 28, 1) * (7, 7) -> (8, 8, 4), stride: (3, 3) pad: (0, 0, 0, 0), dil: (1, 1), flip: false
julia> a = ExplodedConvArray{eltype(batch)}(cdims, batch);
julia> model(a)
10×64 Array{Float32,2}:
[snip]

Теперь как нам это соединить с шифрованием? Для этого нужно:

1. Зашифровать структуру (ExplodedConvArray) так, чтобы мы получали шифротекст для каждого поля. Операции с такой зашифрованной структурой будут поверять, что сделала бы функция с исходной структурой, и делать то же самое гомоморфно.
   
2. Перехватывать определённые операции, чтобы по-другому выполнять их в зашифрованном контексте.

К счастью, Julia предоставляет нам для этого абстракцию: плагин для компилятора, который использует механизм Cassette.jl. Что это такое и как работает я рассказывать не буду, вкратце скажу, что он может определять контекст, например, Encrypted, а затем определяет правила, как должны работать операции в этом контексте. Например, так можно написать для второго требования:

# Define Matrix multiplication between an array and an encrypted block array
function (*::Encrypted{typeof(*)})(a::Array{T, 2}, b::Encrypted{<:BlockArray{T, 2}}) where {T}
    sum(a*b for (i,range) in enumerate(partition(1:size(a, 2), size(b.blocks[1], 1))))
end
# Define Matrix multiplication between an array and an encrypted array
function (*::Encrypted{typeof(*)})(a::Array{T, 2}, b::Encrypted{Array{T, 2}}) where {T}
    result = repeat(diag(a), inner=size(a, 1)).*x
    rotated = b
    for k = 2:size(a, 2)
        rotated = ToyFHE.rotate(GaloisKey(*), rotated)
        result += repeat(diag(circshift(a, (0,(k-1)))), inner=size(a, 1)) .* rotated
    end
    result
end

В результате пользователь сможет писать всё вышеприведённое с минимальным количеством ручной работы:

kp = keygen(ckks_params)
ek = keygen(EvalMultKey, kp.priv)
gk = keygen(GaloisKey, kp.priv; steps=64)
# Create evaluation context
ctx = Encrypted(ek, gk)
# Do public preprocessing
batch = ExplodedConvArray{eltype(batch)}(cdims, batch);
# Run on encrypted data under the encryption context
Cresult = ctx(model)(encrypt(kp.pub, batch))
# Decrypt the answer
decrypt(kp, Cresult)

Конечно, даже этого может быть недостаточно. Параметры криптосистемы (то есть кольцевое ℛ, когда применять modswitch, keyswitch и т.д.) отражают компромисс между точностью ответа и производительностью, и сильно зависят от исполняемого кода. Хотелось бы, чтобы компилятор проанализировал код, который он собирается запустить, предложил параметры для заданного уровня безопасности и желаемой точности, а затем сгенерировал код с минимальным ручным вмешательством.

Заключение

Достижение мечты об автоматическом и безопасном выполнении произвольных вычислений — это сложная задача для любой системы. Но возможности метапрограммирования в Julia и дружелюбный синтаксис делают этот инструмент подходящей платформой для разработки. Коллаборация RAMPARTS (paper, JuliaCon talk) уже сделал первые шаги в этом направлении: простой Julia-код компилируется в полно-гомогенной библиотеке PALISADE. Julia Computing сотрудничает с экспертами RAMPARTS в работе над Verona, недавно анонсированной версией этой системы следующего поколения. В последний год производительность гомоморфных систем шифрования позволяет оценивать интересные вычисления со скоростью, которая близка уровню практической применимости. Врата открыты. С новыми разработками в алгоритмах, ПО и оборудовании, гомоморфное шифрование однозначно станет основной технологией защиты приватности миллионов пользователей.

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