Key-Value Query-Response associative memory model

[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}$ в маленькой модели.

[SergeyGubanov]

Доказательство концепции:

Используется синтаксис C

// Key-Value Query-Response associative memory proof of concept.
// Copyright (C) 2025 Gubanov Sergei. All rights reserved.

#define _CRT_NONSTDC_NO_WARNINGS 1
#define _CRT_SECURE_NO_WARNINGS 1

#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <inttypes.h>
#include <math.h>
#include <string.h>
#include <malloc.h>
#include <stdlib.h>

#ifdef __GNUC__
#define ALIGN(N) __attribute__((aligned(N)))
#else
#define ALIGN(N) __declspec(align(N))
#endif

#define CACHE_LINE_SIZE 64ull
#define THP_SIZE (2ull * 1024ull * 1024ull)

#define PI 3.14159265358979323846264338328

#define SPACE_DIM 32

typedef struct vector {
        float f[SPACE_DIM];
} vector_t;

typedef struct associative_memory {
        vector_t *K;
        vector_t *V;
        size_t N;
        void *raw_pointer;
} associative_memory_t;

typedef struct associative_processor {
        ALIGN(CACHE_LINE_SIZE) vector_t D;
        float *KQ;
        float *VR;
        float *W;
        associative_memory_t *memory;
        void *raw_pointer;
} associative_processor_t;

static float vector_similarity(vector_t *X, vector_t *Y) {
        float *x = X->f;
        float *y = Y->f;
        float XX = 1.0e-10f;
        float YY = 1.0e-10f;
        float XY = 0;
        for (size_t i = 0; i < SPACE_DIM; i++) {
                XX += (x[i] * x[i]);
                YY += (y[i] * y[i]);
                XY += (x[i] * y[i]);
        }
        return (XY / sqrtf(XX * YY));
}

static float get_vector_length(vector_t *X) {
        float *x = X->f;
        float sum2 = 0;
        for (size_t i = 0; i < SPACE_DIM; i++) {
                sum2 += (x[i] * x[i]);
        }
        return sqrtf(sum2);
}

static float set_vector_length(vector_t *X, const float new_length) {
        float *x = X->f;
        float sum2 = 0;
        for (size_t i = 0; i < SPACE_DIM; i++) {
                sum2 += (x[i] * x[i]);
        }
        float old_length = sqrtf(sum2);
        if ((old_length != new_length) && (old_length > 0)) {
                const float factor = (new_length / old_length);
                for (size_t i = 0; i < SPACE_DIM; i++) {
                        x[i] *= factor;
                }
        }
}

static float scalar_product(vector_t *X, vector_t *Y) {
        float *x = X->f;
        float *y = Y->f;
        float sum = 0;
        for (size_t i = 0; i < SPACE_DIM; i++) {
                sum += (x[i] * y[i]);
        }
        return sum;
}

static void vector_mul_add(vector_t *X, vector_t *Y, const float weight) {
        float *x = X->f;
        float *y = Y->f;
        for (size_t i = 0; i < SPACE_DIM; i++) {
                x[i] += (y[i] * weight);
        }
}

static void vector_add(vector_t *X, vector_t *Y) {
        float *x = X->f;
        float *y = Y->f;
        for (size_t i = 0; i < SPACE_DIM; i++) {
                x[i] += y[i];
        }
}

static void vector_sub(vector_t *X, vector_t *Y) {
        float *x = X->f;
        float *y = Y->f;
        for (size_t i = 0; i < SPACE_DIM; i++) {
                x[i] -= y[i];
        }
}

static inline uint32_t random_uint32(uint64_t *seed) {
        const uint64_t A = 3076458245721376007ull;
        const uint64_t B = 6418069888083187241ull;
        const uint64_t x = (*seed) * A + B;
        const uint32_t result = ((uint32_t)(x >> 32));
        (*seed) = x;
        return result;
}

static inline double random_double(uint64_t *seed) {
        return ((1.0 / 4294967296.0) * random_uint32(seed));
}

static void random_vector(vector_t *v, uint64_t *seed, float length) {
        for (int i = 0; i < SPACE_DIM; i++) {
                double p = 2 * PI * random_double(seed);
                double r = random_double(seed);
                v->f[i] = (float)(sin(p) * sqrt(-2 * log(r)));
        }
        set_vector_length(v, length);
}

static associative_memory_t *associative_memory_create(const size_t N) {
        const size_t size_K = ((sizeof(vector_t) * N) + CACHE_LINE_SIZE - 1) & ~(CACHE_LINE_SIZE - 1);
        const size_t size_V = ((sizeof(vector_t) * N) + CACHE_LINE_SIZE - 1) & ~(CACHE_LINE_SIZE - 1);
        const size_t size_am = (sizeof(associative_memory_t) + CACHE_LINE_SIZE - 1) & ~(CACHE_LINE_SIZE - 1);
        const size_t total_size = (THP_SIZE + size_am + size_K + size_V + THP_SIZE - 1) & ~(THP_SIZE - 1);
        uintptr_t address = (uintptr_t)malloc(total_size);
        if (address == 0) {
                printf("associative_memory_create() Out of memory");
                exit(1);
                return NULL;
        }
        void *raw_pointer = (void *)address; address = (address + THP_SIZE - 1) & ~(THP_SIZE - 1);
        associative_memory_t *memory = (associative_memory_t *)address; address += size_am;
        memory->K = (vector_t *)address; address += size_K;
        memory->V = (vector_t *)address; address += size_V;
        memory->N = N;
        memory->raw_pointer = raw_pointer;
        return memory;
}

static void associative_memory_delete(associative_memory_t *memory) {
        free(memory->raw_pointer);
}

static associative_processor_t *associative_processor_create(associative_memory_t *memory) {
        const size_t N = memory->N;
        const size_t size_KQ = ((sizeof(float) * N) + CACHE_LINE_SIZE - 1) & ~(CACHE_LINE_SIZE - 1);
        const size_t size_VR = ((sizeof(float) * N) + CACHE_LINE_SIZE - 1) & ~(CACHE_LINE_SIZE - 1);
        const size_t size_W = ((sizeof(float) * N) + CACHE_LINE_SIZE - 1) & ~(CACHE_LINE_SIZE - 1);
        const size_t size_am = (sizeof(associative_processor_t) + CACHE_LINE_SIZE - 1) & ~(CACHE_LINE_SIZE - 1);
        const size_t total_size = (CACHE_LINE_SIZE + size_am + size_KQ + size_VR + size_W + CACHE_LINE_SIZE - 1) & ~(CACHE_LINE_SIZE - 1);
        uintptr_t address = (uintptr_t)malloc(total_size);
        if (address == 0) {
                printf("associative_processor_create() Out of memory");
                exit(1);
                return NULL;
        }
        void *raw_pointer = (void *)address;
        address = (address + CACHE_LINE_SIZE - 1) & ~(CACHE_LINE_SIZE - 1);
        associative_processor_t *processor = (associative_processor_t *)address; address += size_am;
        processor->KQ = (float *)address; address += size_KQ;
        processor->VR = (float *)address; address += size_VR;
        processor->W = (float *)address; address += size_W;
        processor->memory = memory;
        processor->raw_pointer = raw_pointer;
        return processor;
}

static void associative_processor_delete(associative_processor_t *processor) {
        free(processor->raw_pointer);
}

static void inference(associative_processor_t *processor, vector_t *Q, vector_t *R) {
        vector_t *K = processor->memory->K;
        vector_t *V = processor->memory->V;
        size_t N = processor->memory->N;
        float *KQ = processor->KQ;
        float *VR = processor->VR;
        float *W = processor->W;
        vector_t *D = &processor->D;
        float query_length = get_vector_length(Q);
        for (size_t i = 0; i < N; i++) {
                KQ[i] = scalar_product(&K[i], Q);
        }
        memset(R, 0, sizeof(vector_t));
        for (size_t iteration = 0; iteration < 256; iteration++) {
                float max_KQ_VR = 0;
                for (size_t i = 0; i < N; i++) {
                        float VRi = scalar_product(&V[i], R);
                        VR[i] = VRi;
                        float KQ_VR = KQ[i] + VRi;
                        if (max_KQ_VR < KQ_VR) {
                                max_KQ_VR = KQ_VR;
                        }
                }
                float sum_w = 0;
                for (size_t i = 0; i < N; i++) {
                        float w = expf(KQ[i] + VR[i] - max_KQ_VR);
                        sum_w += w;
                        W[i] = w;
                }
                if (!(sum_w > 0)) {
                        printf("sum_w == NaN\n");
                        exit(1);
                }
                memset(D, 0, sizeof(vector_t));
                float norm_w = (1.0f / sum_w);
                for (size_t i = 0; i < N; i++) {
                        float w = W[i] * norm_w;
                        if (w > 1.0e-20f) {
                                vector_mul_add(D, &V[i], w);
                        }
                }
                vector_sub(D, R);
                const float dt = 1;
                vector_mul_add(R, D, dt);
                float D_length = get_vector_length(D);
                float R_length = get_vector_length(R);
                if (0) {
                        printf("[%" PRIu64 "] |Q|=%f, |R|=%f, |D|=%f, |D|/|R|=%f, max(exp(KQ+VR))=%e\n",
                                iteration, query_length, R_length, D_length, (D_length / R_length),
                                exp(max_KQ_VR));
                }
                if ((D_length * 10000) < R_length) {
                        if (0) {
                                printf("[%" PRIu64 "] |Q|=%f, |R|=%f, |D|=%f, |D|/|R|=%f, max(exp(KQ+VR))=%e\n",
                                        iteration, query_length, R_length, D_length, (D_length / R_length),
                                        exp(max_KQ_VR));
                        }
                        break;
                }
        }
}

int main(int argc, char *argv[]) {
        const size_t GROUND_TRUE_SAMPLES_COUNT = 1024;
        const size_t RANDOM_KEYS_SAMPLES_COUNT = (128 * 1024);
        const float ground_true_vector_length = (float)sqrtf(logf(8.0f * GROUND_TRUE_SAMPLES_COUNT));
        const float random_keys_vector_length = (float)sqrtf(logf(8.0f * RANDOM_KEYS_SAMPLES_COUNT));
        printf("#\n");
        printf("# Key-Value Query-Response associative memory proof of concept.\n");
        printf("# Copyright (C) 2025 Gubanov Sergei. All rights reserved.\n");
        printf("#\n");
        printf("# Dimension of vector space = %d\n", SPACE_DIM);
        printf("# Ground-True associative memory samples = %" PRIu64 "\n", GROUND_TRUE_SAMPLES_COUNT);
        printf("# Random-Keys associative memory samples = %" PRIu64 "\n", RANDOM_KEYS_SAMPLES_COUNT);
        associative_memory_t *ground_true_memory = associative_memory_create(GROUND_TRUE_SAMPLES_COUNT);
        associative_processor_t *ground_true_processor = associative_processor_create(ground_true_memory);
        uint64_t seed = 42;
        for (size_t i = 0; i < ground_true_memory->N; i++) {
                random_vector(&ground_true_memory->K[i], &seed, ground_true_vector_length);
                random_vector(&ground_true_memory->V[i], &seed, ground_true_vector_length);
        }
        printf("#\n");
        printf("# Ground-True associative memory self-verification:\n");
        {
                float min = 1;
                float max = -1;
                float sum = 0;
                for (size_t i = 0; (i < ground_true_memory->N); i++) {
                        vector_t response;
                        inference(ground_true_processor, &ground_true_memory->K[i], &response);
                        const float s = vector_similarity(&ground_true_memory->V[i], &response);
                        sum += s;
                        if (min > s) { min = s; }
                        if (max < s) { max = s; }
                        const size_t P = (ground_true_memory->N / 16);
                        if ((i % P) == (P - 1)) {
                                printf("[%8" PRIu64 "] Similarity Min=%f Avr=%f Max=%f\n", (i + 1),
                                        min, (sum / (i + 1)), max);
                        }
                }
                printf("# Similarity: Min=%f Avr=%f Max=%f\n", min, (sum / ground_true_memory->N), max);
        }
        printf("#\n");
        printf("# Random-Keys associative memory sampling:\n");
        associative_memory_t *random_keys_memory = associative_memory_create(RANDOM_KEYS_SAMPLES_COUNT);
        for (size_t i = 0; i < random_keys_memory->N; i++) {
                vector_t *key = &random_keys_memory->K[i];
                vector_t *value = &random_keys_memory->V[i];
                random_vector(key, &seed, random_keys_vector_length);
                inference(ground_true_processor, key, value);
                const size_t progress = (random_keys_memory->N / 16);
                if ((i % progress) == (progress - 1)) {
                        printf("[%8" PRIu64 "] progress %6.2f%%\n", (i + 1), ((100.0 * (i + 1)) / random_keys_memory->N));
                }
        }
        printf("#\n");
        printf("# Random-Keys associative memory verification with respect to Ground-True memories:\n");
        associative_processor_t *random_keys_processor = associative_processor_create(random_keys_memory);
        {
                float min = 1;
                float max = -1;
                float sum = 0;
                for (size_t i = 0; (i < ground_true_memory->N); i++) {
                        vector_t response;
                        inference(random_keys_processor, &ground_true_memory->K[i], &response);
                        const float s = vector_similarity(&ground_true_memory->V[i], &response);
                        sum += s;
                        if (min > s) { min = s; }
                        if (max < s) { max = s; }
                        const size_t P = (ground_true_memory->N / 16);
                        if ((i % P) == (P - 1)) {
                                printf("[%8" PRIu64 "] Similarity Min=%f Avr=%f Max=%f\n", (i + 1),
                                        min, (sum / (i + 1)), max);
                        }
                }
                printf("# Similarity: Min=%f Avr=%f Max=%f\n", min, (sum / ground_true_memory->N), max);
        }
        associative_processor_delete(ground_true_processor);
        associative_processor_delete(random_keys_processor);
        associative_memory_delete(ground_true_memory);
        associative_memory_delete(random_keys_memory);
        printf("#\n");
        return 0;
}
 

-- 03.11.2025, 00:12 --

Пояснение к коду.

Там сначала создаётся "ground-true" ассоциативная память из 1024 пар 32-мерных случайных векторов ключей и значений.

Потом создаётся "random-keys" ассоциативная память из 131072 пар 32-мерных случайных векторов ключей, а соответствующие им значения вычисляются согласно "ground-true" ассоциативной памяти.

Среди множества ключей "random-keys" памяти нет ключей из "ground-true" памяти.

Затем ассоциативной памяти "random-keys" задаются вопросы соответствующие ключам из "ground-true" памяти и ответы которые генерирует "random-keys" память сравниваются с ответами из "ground-true" памяти.

Эти ответы оказываются с большой точностью (пять девяток) совпадающими с ответами из "ground-true" памяти.

Таким образом, в Key-Value Query-Response модели ассоциативной памяти перенос знаний из одной памяти в другую возможен путём простого перевычисления векторов значений на другом множестве векторов ключей.

Разумеется новое множество ключей надо подобрать так, чтобы не сильно упустить что-то из старого множества ключей. В данном случае количество случайных ключей было в 128 раз больше чем количество "ground-true" ключей, поэтому вероятность что-то случайно упустить была не сильно велика. При меньшем количестве случайных ключей оно иногда "промахивается", но в среднем работает хорошо.

-- 03.11.2025, 00:24 --

Все $N$ пар векторов ключей и значений "ground-true" памяти нормированы так, что их длины одинаковы и равны $|K_{i}| = \sqrt{\log(8 N)}$, $|V_{i}| = \sqrt{\log(8 N)}$.

Все $M$ векторов ключей "random-keys" памяти нормированы так, что их длины одинаковы и равны $|K_{i}| = \sqrt{\log(8 M)}$. Векторы значений "random-keys" памяти никак не нормированы (этого нельзя делать), их длины ровно такие какие получились в результате дистилляции знаний из "ground-true" памяти в "random-keys" память.

[SergeyGubanov]

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

В предыдущем примере "random-keys" ассоциативная память задавалась 131072 парами ключей-значений, но её ранг был не выше чем 1024 потому, что она была порождена из "ground-true" ассоциативной памяти заданной 1024 парами ключей-значений.

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

Например, можно повыбрасывать те ключи-значения, которые и так вычисляются (воспроизводятся) благодаря оставшимся ключам-значениям. На каждом шаге "жадно" выбрасывать ту пару ключ-значение которая всех точнее воспроизводится за счёт оставшихся.

[SergeyGubanov]

Рассматриваемая Key-Value Query-Response модель ассоциативной памяти может быть использована для решения задачи предсказания следующего вектора $\vec{X}_{i + 1}$ последовательности векторов при условии, что следующий вектор последовательности зависит от небольшого количества предыдущих векторов этой последовательности (Марковская цепь).

Например, в случае зависимости от трёх предыдущих векторов:

$$ \vec{X}_{i + 1} = \vec{F}\left( \vec{X}_{i}, \vec{X}_{i - 1}, \vec{X}_{i - 2} \right). $$

При этом роль вектора запроса (Query Vector) выполняют предыдущие векторы последовательности, а роль вектора ответа (Response Vector) исполняет предсказываемый следующий вектор.

Функция потери (она же энергия если говорить в терминах Хопфильдовской модели) и дифференциальное уравнение для предсказания следующего вектора $\vec{X}_{i + 1}$ по трём предыдущим векторам этой последовательности принимает следующий вид:

$$ \vec{Q}^{(2)} = \vec{X}_{i - 2}, \qquad \vec{Q}^{(1)} = \vec{X}_{i - 1}, \qquad \vec{Q}^{(0)} = \vec{X}_{i}, \qquad \vec{R} = \vec{X}_{i + 1}, $$

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

$$ \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}^{(2)}_{n} \vec{Q}^{(2)} \right) + \left( \vec{K}^{(1)}_{n} \vec{Q}^{(1)} \right) + \left( \vec{K}^{(0)}_{n} \vec{Q}^{(0)} \right) + \left( \vec{V}_{n} \vec{R} \right) \right) }{ \sum_{n = 1}^{N} \exp \left( \left( \vec{K}^{(2)}_{n} \vec{Q}^{(2)} \right) + \left( \vec{K}^{(1)}_{n} \vec{Q}^{(1)} \right) + \left( \vec{K}^{(0)}_{n} \vec{Q}^{(0)} \right) + \left( \vec{V}_{n} \vec{R} \right) \right) }. $$

[SergeyGubanov]

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

Например, пусть система была в состоянии описываемом (hidden) вектором $\vec{H}_{in}$ потом она на вход получила векторный сигнал $\vec{X}$ и перешла в состояние описываемое (hidden) вектором $\vec{H}_{out}$, при этом система на своём выходе выдала сигнал $\vec{Y}$

$$ \vec{H}_{out} = \vec{F}\left( \vec{H}_{in}, \, \vec{X} \right), \qquad \vec{Y} = \vec{G}\left( \vec{H}_{out} \right). $$

Роль вектора запроса (Query Vector) выполняют вектор предыдущего внутреннего состояния рекуррентной системы $\vec{H}_{in}$ и вектор входного сигнала $\vec{X}$, а роль вектора ответа (Response Vector) исполняет вектор следующего внутреннего состояния системы $\vec{H}_{out}$.

Уравнения для вычисления следующего вектора состояния $\vec{H}_{out}$:

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

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

Уравнения для вычисления выходного сигнала $\vec{Y}$:

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

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

[SergeyGubanov]

С помощью рассмотренного выше механизма ассоциативной памяти можно предсказывать шум в диффузионой модели.

Упрощенно говоря, задача диффузионной модели состоит в том, чтобы для вектора запроса $\vec{X}$ предсказать вектор "шума" $\vec{\varepsilon}(\vec{X})$, такой, что если этот "шум" вычесть из вектора запроса, то получится "очищенный" вектор $\vec{X}' := \vec{X} - \vec{\varepsilon}(\vec{X})$.

Разумеется в явном виде векторы "шума" поместить в Key-Value список ассоциативной памяти в качестве $\vec{V}_{n}$ вектора нельзя так как шум может быть любой и список получился бы бесконечным. Поэтому в памяти надо хранить только ключи $\vec{K}_{n}$, а соответствующие векторы значений $\vec{V}_{n}$ вычисляются во время инференса вычитанием вектора ключа $\vec{K}_{n}$ из фактического вектора запроса $\vec{Q}$.

Итого:

Вектор запроса (Query vector) это теперь будет зашумлённый вектор $\vec{Q} = \vec{X}$.
Вектор ответа (Response vector) это теперь будет вектор шума $\vec{R} = \vec{\varepsilon}$.
Векторы ключей (Key vectors) имеют смысл эталонных незашумлённых векторов (для них надо предсказывать нулевой вектор шума).
Векторы значений (Value vectors) теперь "виртуальные" их не надо хранить в памяти, они вычисляются на лету как разность между фактическим вектором запроса и ключом $\vec{V}_{n} = \vec{X} - \vec{K}_{n}$.

$$ \sum_{n = 1}^{N} \exp \left( \left( \vec{K}_n \vec{Q} \right) + \left( \vec{V}_n \vec{R} \right) \right) \quad \to \quad \sum_{n = 1}^{N} \exp \left( \left( \vec{K}_n \vec{X} \right) + \left( \left( \vec{X} - \vec{K}_n \right) \vec{\varepsilon} \right) \right). $$

Для предсказания вектора шума $\vec{\varepsilon}$ связанного с вектором $\vec{X}$ функция потери (энергия Хоппфильда) имеет следующий вид:

$$ \mathcal{L}(\vec{X}, \vec{\varepsilon}) = \frac{1}{2} \vec{\varepsilon}^2 - \log \sum_{n = 1}^{N} \exp \left( \left( \vec{K}_n \vec{X} \right) + \left( \left( \vec{X} - \vec{K}_n \right) \vec{\varepsilon} \right) \right). $$

Уравнение для нахождения $\vec{\varepsilon}(\vec{X})$

$$ \frac{d \vec{\varepsilon} }{d t} = - \frac{\partial \mathcal{L}}{ \partial \vec{\varepsilon} }, $$

$$ \frac{d \vec{\varepsilon} }{d t} + \vec{\varepsilon} = \frac{ \sum_{n = 1}^{N} \left( \vec{X} - \vec{K}_n \right) \exp \left( \left( \vec{K}_n \vec{X} \right) + \left( \left( \vec{X} - \vec{K}_n \right) \vec{\varepsilon} \right) \right) }{ \sum_{n = 1}^{N} \exp \left( \left( \vec{K}_n \vec{X} \right) + \left( \left( \vec{X} - \vec{K}_n \right) \vec{\varepsilon} \right) \right) }. $$

Уравнения нелинейные, работают так как задумано только если нормировка ключей:

$$ | \vec{K}_{n} |^2 \ge 2 \log(N), $$

в противном случае всё излишне усредняется.

Если по указанной формуле вычислить вектор шума $\vec{\varepsilon}$ для случайного вектора $\vec{X}$, затем вычесть этот шум из него, получив "очищенный" вектор $\vec{X}' := \vec{X} - \vec{\varepsilon}(\vec{X})$, то для этого "очищенного" вектора $\vec{X}'$ шум вычисленный аналогичным образом равен нулю (шум очищенного вектора равен нулю), при этом сам вектор $\vec{X}'$ не равен ни одному из ключей $\vec{K}_{n}$, а если вдруг и равен, то чисто случайно.

[SergeyGubanov]

От дискретного случая распределения ключей $\vec{K}_n$ и значений $\vec{V}_n$ логично было бы перейти к рассмотрению непрерывного распределения ключей $\vec{K}(\xi)$ и значений $\vec{V}(\xi)$ по некоторому $d$-мерному многообразию $\Gamma$.

$$ \vec{K}_n \quad \to \quad \vec{K}_{\xi}, $$

$$ \vec{V}_n \quad \to \quad \vec{V}_{\xi}, $$

$$ \sum_{n = 1}^{N} \quad \to \quad \int_{\Gamma} D {\xi}, $$

$$ \sum_{n = 1}^{N} \exp \left( \left( \vec{K}_n \vec{Q} \right) + \left( \vec{V}_n \vec{R} \right) \right) \quad \to \quad \int_{\Gamma} D {\xi} \, \exp \left( \left( \vec{K}_{\xi} \, \vec{Q} \right) + \left( \vec{V}_{\xi} \, \vec{R} \right) \right), $$

$$ \mathcal{L}(Q, R) = \frac{1}{2}|\vec{Q}|^2 + \frac{1}{2}|\vec{R}|^2 - \log \int_{\Gamma} D {\xi} \, \exp \left( \left( \vec{K}_{\xi} \, \vec{Q} \right) + \left( \vec{V}_{\xi} \, \vec{R} \right) \right). $$