Key-Value Query-Response associative memory model

[SergeyGubanov]

Предлагаю участникам форума покритиковать следующую модель ассоциативной памяти. Она возникла у меня в голове в результате творческого переосмысления "современной Хопфилдовской" модели Дмитрия Кротова.

Пусть есть набор $N$ пар векторов ключ-значение $\vec{K}_{n}$, $\vec{V}_{n}$ (Key-Value). Размерность векторного пространства ключей и размерность векторного пространства значений может быть разной:

$$ \dim{K} \neq \dim{V}. $$

Задача ассоциативной памяти состоит в том, чтобы по заданному вектору запроса $\vec{Q}$ (Query) [принадлежащего векторному пространству ключей] вычислить вектор ответа $\vec{R}$ (Response) [принадлежащего векторному пространству значений], такой что если вектор запроса близок к какому-то вектору ключа $\vec{Q} \approx \vec{K}_{n}$, то вектор ответа должен быть близок к соответствующему вектору значения $\vec{R} \approx \vec{V}_{n}$.

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

$$ \mathcal{L}(Q, R) = \frac{1}{2} \vec{Q}^2 + \frac{1}{2} \vec{R}^2 - \log \sum_{n = 1}^{N} \exp \left( \left(\vec{K}_{n} \vec{Q}\right) + \left(\vec{V}_{n} \vec{R}\right) \right). $$

Тогда для запроса $\vec{Q}$ ответ $\vec{R}$ минимизирующий функцию $\mathcal{L}(Q, R) $ должен быть вычислен градиентным спуском:

$$ \frac{d \vec{R}}{d t} = - \frac{\partial \mathcal{L}}{\partial \vec{R}}. $$

В явном виде:

$$ \frac{d \vec{R}}{d t} + \vec{R} = \frac{\sum_{n = 1}^{N} \vec{V}_{n} \exp \left( \left(\vec{K}_{n} \vec{Q}\right) + \left(\vec{V}_{n} \vec{R}\right) \right)}{\sum_{n = 1}^{N} \exp \left( \left(\vec{K}_{n} \vec{Q}\right) + \left(\vec{V}_{n} \vec{R}\right) \right)}. $$

Начальное условие: $\vec{R} = 0$ при $t = 0$. В конечном состоянии $\frac{d \vec{R}}{d t} \to 0$ при $t \to \infty$

Плюсы:
1) Процесса обучения такой модели вообще нет, так как набор $N$ пар ключ-значение задан по условию.

Минусы:
1) Во время инференса надо численно решать дифференциальное уравнение.
2) Если $N$ слишком велико, то процедура инференса становится слишком дорогой.
3) В связи с (2), модель годится только для случаев плавной зависимости $\vec{V}$ от $\vec{K}$, чем плавнее, тем лучше.

[SergeyGubanov]

А теперь добавим скрытый слой (Hidden layer).

Давайте теперь входящий вектор запроса $\vec{Q}$ сначала преобразуем во внутренний (скрытый, hidden) вектор $\vec{H}$. А затем, этот внутренний вектор $\vec{H}$ преобразуем в вектор ответа $\vec{R}$:

$$ \vec{Q} \to \vec{H} \to \vec{R}. $$

Функция потерь преобразования $\vec{Q} \to \vec{H}$ (энкодер):

$$ \mathcal{L}_{1}(Q, H) = \frac{1}{2} \vec{H}^2 - \log \sum_{n = 1}^{N_{1}} \exp \left( \left(\vec{K}^{(1)}_{n} \vec{Q}\right) + \left(\vec{V}^{(1)}_{n} \vec{H}\right) \right), $$

Функция потерь преобразования $\vec{H} \to \vec{R}$ (декодер):

$$ \mathcal{L}_{2}(H, R) = \frac{1}{2} \vec{R}^2 - \log \sum_{n = 1}^{N_{2}} \exp \left( \left(\vec{K}^{(2)}_{n} \vec{H}\right) + \left(\vec{V}^{(2)}_{n} \vec{R}\right) \right). $$

Размерности рассматриваемых векторных пространств могут не совпадать:

$$ \dim K^{(1)} \neq \dim V^{(1)}, \quad \dim K^{(2)} \neq \dim V^{(2)}, $$

но, разумеется, должны быть одинаковы размерности векторных пространств $K^{(2)}$ и $V^{(1)}$:

$$ \dim K^{(2)} = \dim V^{(1)}. $$

Уравнение энкодера $\vec{Q} \to \vec{H}$:

$$ \frac{d \vec{H}}{d t} = - \frac{\partial \mathcal{L}_{1}}{\partial \vec{H}} \qquad \to \qquad \frac{d \vec{H}}{d t} + \vec{H} = \frac{\sum_{n = 1}^{N_{1}} \vec{V}^{(1)}_{n} \exp \left( \left(\vec{K}^{(1)}_{n} \vec{Q}\right) + \left(\vec{V}^{(1)}_{n} \vec{H}\right) \right)}{\sum_{n = 1}^{N_{1}} \exp \left( \left(\vec{K}^{(1)}_{n} \vec{Q}\right) + \left(\vec{V}^{(1)}_{n} \vec{H}\right) \right)}, $$

Уравнение декодера $\vec{H} \to \vec{R}$:

$$ \frac{d \vec{R}}{d t} = - \frac{\partial \mathcal{L}_{2}}{\partial \vec{R}} \qquad \to \qquad \frac{d \vec{R}}{d t} + \vec{R} = \frac{\sum_{n = 1}^{N_{2}} \vec{V}^{(2)}_{n} \exp \left( \left(\vec{K}^{(2)}_{n} \vec{H}\right) + \left(\vec{V}^{(2)}_{n} \vec{R}\right) \right)}{\sum_{n = 1}^{N_{2}} \exp \left( \left(\vec{K}^{(2)}_{n} \vec{H}\right) + \left(\vec{V}^{(2)}_{n} \vec{R}\right) \right)}. $$

Введение скрытого слоя целесообразно если можно подобрать параметры так, что $N_{1} \ll N$ и $N_{2} \ll N$, то есть если можно уменьшить количество параметров по сравнению со случаем без скрытого слоя.

Разумеется число скрытых слоёв можно увеличить:

$$ \vec{Q} \to \vec{H}^{(1)} \to \vec{H}^{(2)} \to \ldots \to \vec{H}^{(K)} \to \vec{R}. $$

[SergeyGubanov]

В первом сообщении была рассмотрена задача как по заданным $N$ парам ключей $\vec{K}_{n} \in K$ и значений $\vec{V}_{n} \in V$ интерполировать ответ $\vec{R} \in V$ на вопрос $\vec{Q} \in K$. Там же было сказано, что если $N$ слишком велико, то процедура такой интерполяции становится слишком дорогой. Нам выгодно чтобы $N$ было как можно меньше.

Теперь рассмотрим обратную задачу. Пусть теперь нам известно $M$ пар запросов $\vec{Q}_{a}$ и ответов $\vec{R}_{a}$ и по ним надо отыскать (выучить) $N$ пар ключей и значений. Причём, мы хотим, чтобы $N$ было как можно меньше: $N \ll M$. Другими словами, давайте выведем формулы для регрессии $M$ пар запросов/ответов к $N$ парам ключей/значений.

Довольно понятно, что функция подлежащая минимизации могла бы иметь следующий вид:

$$ \mathcal{L} = \frac{1}{2} \sum_{a=1}^{M} \left( \vec{Q}^{2}_{a} + \vec{R}^{2}_{a} \right) + \frac{1}{2} \sum_{n=1}^{N} \left( \vec{K}^{2}_{n} + \vec{V}^{2}_{n} \right) - \sum_{a=1}^{M} \log \sum_{n = 1}^{N} \exp \left( \left(\vec{K}_{n} \vec{Q}_{a} \right) + \left(\vec{V}_{n} \vec{R}_{a} \right) \right). $$

Уравнения для нахожения $\vec{K}_{i}$ и $\vec{V}_{i}$:

$$ \frac{d \vec{K}_{i}}{d t} = - \frac{\partial \mathcal{L}}{\partial \vec{K}_{i}}, \qquad \frac{d \vec{V}_{i}}{d t} = - \frac{\partial \mathcal{L}}{\partial \vec{V}_{i}}. $$

В явном виде:

$$ \frac{d \vec{K}_{i}}{d t} + \vec{K}_{i} = \sum_{a = 1}^{M} \frac{ \vec{Q}_{a} \exp \left( \left(\vec{K}_{i} \vec{Q}_{a} \right) + \left(\vec{V}_{i} \vec{R}_{a} \right) \right) }{ \sum_{n = 1}^{N} \exp \left( \left(\vec{K}_{n} \vec{Q}_{a} \right) + \left(\vec{V}_{n} \vec{R}_{a}\right) \right)}, $$

$$ \frac{d \vec{V}_{i}}{d t} + \vec{V}_{i} = \sum_{a = 1}^{M} \frac{ \vec{R}_{a} \exp \left( \left(\vec{K}_{i} \vec{Q}_{a} \right) + \left(\vec{V}_{i} \vec{R}_{a} \right) \right) }{ \sum_{n = 1}^{N} \exp \left( \left(\vec{K}_{n} \vec{Q}_{a} \right) + \left(\vec{V}_{n} \vec{R}_{a}\right) \right)}. $$

Однако, уравнения получились с подвохом. Надо ж как-то специально следить чтобы все ключи были попарно различны $\vec{K}_{i} \ne \vec{K}_{j}$. Чтобы это произошло автоматически, пожалуй следует добавить специальное отталкивающее взаимодействие между ключами. Например, вот такая добавка

$$ \mathcal{L}_{\text{regularized}} = \mathcal{L} + \frac{\lambda}{2} \sum_{i=1}^{N} \sum_{j \neq i}^{N} \left( \vec{K}_{i} \vec{K}_{j} \right)^2 $$

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

$$ \frac{d \vec{K}_{i}}{d t} + \vec{K}_{i} + \lambda \sum_{j \neq i}^{N} \vec{K}_{j} \left( \vec{K}_{i} \vec{K}_{j} \right) = \sum_{a = 1}^{M} \frac{ \vec{Q}_{a} \exp \left( \left(\vec{K}_{i} \vec{Q}_{a} \right) + \left(\vec{V}_{i} \vec{R}_{a} \right) \right) }{ \sum_{n = 1}^{N} \exp \left( \left(\vec{K}_{n} \vec{Q}_{a} \right) + \left(\vec{V}_{n} \vec{R}_{a}\right) \right)}. $$

[SergeyGubanov]

Замечание про нормировку векторов $\vec{K}_{i}$ и $\vec{V}_{i}$.

Для того чтобы в сумме

$$\sum_{n = 1}^{N} \exp \left( \left(\vec{K}_{n} \vec{Q}\right) + \left(\vec{V}_{n} \vec{R}\right) \right)$$

слагаемое с $\vec{Q} \approx \vec{K}_{i}$ и $\vec{R} \approx \vec{V}_{i}$ заметно выделялось на фоне всех остальных слагаемых нужно чтобы векторы ключей и значений были нормированы примерно так:

$$ \left| \vec{K}_{n} \right|^2 > \log N, \qquad \left| \vec{V}_{n} \right|^2 > \log N. $$

В противном случае интересное слагаемое не будет сильно выделяться на фоне всех остальных слагаемых и ответ $\vec{R}$ на вопрос $\vec{Q}$ будет сильно сглажен по всем $\vec{V}_{n}$.

[SergeyGubanov]

SergeyGubanov в сообщении #1707482 писал(а):

Теперь рассмотрим обратную задачу.

Есть решение получше, называется дистилляция знаний.

Пусть есть $N$ пар ключей $\vec{K}_{n}$ и значений $\vec{V}_{n}$, причём $N$ на столько большое, что доставляет вычислительные неприятности. Нам хватило бы менее точной, но более компактной модели. Всё что нам нужно, это просто сэмплировать $M$ "популярных" векторов вопросов $\vec{Q}_{a}$ и вычислить для них точные ответы $\vec{R}_{a}$ при помощи большой модели ($N \gg M$). Затем эти $M$ пар векторов $\vec{Q}_{a}$ и $\vec{R}_{a}$ можно использовать в качестве векторов ключей $\vec{Q}_{a} \to \vec{K}_{a}$ и значений $\vec{R}_{a} \to \vec{V}_{a}$ в маленькой модели.