252 lines
8.7 KiB
C
252 lines
8.7 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/*
|
|
* test_pooling_loss.c — HDC Memory Viability Test 9
|
|
*
|
|
* Mean-Pooling Information Loss (synthetic, no model needed)
|
|
*
|
|
* Method:
|
|
* 1. Generate random embeddings (dim=768)
|
|
* 2. Group into batches of n_tokens (1, 4, 8, 16, 32)
|
|
* 3. For each pooling strategy:
|
|
* a. Pool batch → single representative embedding
|
|
* b. Store via fm_put()
|
|
* c. Query with each original single embedding
|
|
* d. Measure recall@5 (does the batch appear in top-5?)
|
|
* 4. Compare: mean-pool, max-pool, first-token, last-token
|
|
*
|
|
* Usage:
|
|
* ./examples/test_pooling_loss
|
|
*/
|
|
|
|
#include "hybridmem/hybridmem.h"
|
|
#include "hdc/hdc.h"
|
|
#include "tensor/tensor.h"
|
|
|
|
#include <math.h>
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
|
|
/* Simple xorshift PRNG */
|
|
static uint32_t rng_state = 42;
|
|
static float randf(void) {
|
|
rng_state ^= rng_state << 13;
|
|
rng_state ^= rng_state >> 17;
|
|
rng_state ^= rng_state << 5;
|
|
return (float)(rng_state & 0x7FFFFFFF) / (float)0x7FFFFFFF * 2.0f - 1.0f;
|
|
}
|
|
|
|
/* Generate random embeddings: out[n * dim] */
|
|
static void gen_random_embeddings(float *out, int n, int dim) {
|
|
for (int i = 0; i < n * dim; i++)
|
|
out[i] = randf();
|
|
}
|
|
|
|
/* Mean-pool: out[dim] = mean of rows[n * dim] */
|
|
static void mean_pool(const float *rows, int n, int dim, float *out) {
|
|
memset(out, 0, dim * sizeof(float));
|
|
for (int t = 0; t < n; t++)
|
|
for (int j = 0; j < dim; j++)
|
|
out[j] += rows[t * dim + j];
|
|
float inv_n = 1.0f / n;
|
|
for (int j = 0; j < dim; j++)
|
|
out[j] *= inv_n;
|
|
}
|
|
|
|
/* Max-pool: out[j] = max over rows */
|
|
static void max_pool(const float *rows, int n, int dim, float *out) {
|
|
memcpy(out, rows, dim * sizeof(float));
|
|
for (int t = 1; t < n; t++)
|
|
for (int j = 0; j < dim; j++)
|
|
if (rows[t * dim + j] > out[j])
|
|
out[j] = rows[t * dim + j];
|
|
}
|
|
|
|
/* First-token: out = rows[0] */
|
|
static void first_token_pool(const float *rows, int n, int dim, float *out) {
|
|
(void)n;
|
|
memcpy(out, rows, dim * sizeof(float));
|
|
}
|
|
|
|
/* Last-token: out = rows[n-1] */
|
|
static void last_token_pool(const float *rows, int n, int dim, float *out) {
|
|
memcpy(out, &rows[(n - 1) * dim], dim * sizeof(float));
|
|
}
|
|
|
|
typedef void (*pool_fn)(const float *rows, int n, int dim, float *out);
|
|
|
|
/* Run recall experiment for one pooling strategy at a given batch size.
|
|
* Returns recall@5 (fraction of query embeddings that find their batch in top-5). */
|
|
static float run_recall_experiment(
|
|
pool_fn pool, const char *pool_name,
|
|
int dim, int batch_size, int n_batches, int top_k)
|
|
{
|
|
/* Create HybridMem store */
|
|
char store_dir[256];
|
|
snprintf(store_dir, sizeof(store_dir), "/tmp/hm_pooling_test_%s_%d",
|
|
pool_name, batch_size);
|
|
char cmd[512];
|
|
snprintf(cmd, sizeof(cmd), "rm -rf %s", store_dir);
|
|
system(cmd);
|
|
|
|
fm_config cfg = fm_config_default(dim);
|
|
cfg.store_dir = store_dir;
|
|
fm_store *store = fm_store_new(&cfg);
|
|
if (!store) {
|
|
fprintf(stderr, "Failed to create store at %s\n", store_dir);
|
|
return -1;
|
|
}
|
|
|
|
/* Generate embeddings and store pooled versions */
|
|
int total_tokens = n_batches * batch_size;
|
|
float *all_embeddings = (float*)malloc(total_tokens * dim * sizeof(float));
|
|
gen_random_embeddings(all_embeddings, total_tokens, dim);
|
|
|
|
float *pooled = (float*)malloc(dim * sizeof(float));
|
|
|
|
for (int b = 0; b < n_batches; b++) {
|
|
float *batch = &all_embeddings[b * batch_size * dim];
|
|
pool(batch, batch_size, dim, pooled);
|
|
|
|
/* Store: keys = pooled (1 token), values = pooled (dummy) */
|
|
int64_t shape_k[] = {1, dim};
|
|
tm_tensor *keys = tm_from_data(pooled, shape_k, 2, TM_F32);
|
|
tm_tensor *keys_owned = tm_clone(keys);
|
|
tm_tensor *vals_owned = tm_clone(keys);
|
|
tm_free(keys);
|
|
fm_put(store, keys_owned, vals_owned, 0, 1.0f);
|
|
tm_free(keys_owned);
|
|
tm_free(vals_owned);
|
|
}
|
|
|
|
fm_store_flush(store);
|
|
|
|
/* Query with individual embeddings */
|
|
int hits = 0;
|
|
int queries = 0;
|
|
int sample_step = batch_size > 1 ? batch_size / 2 : 1; /* sample from each batch */
|
|
|
|
for (int b = 0; b < n_batches; b++) {
|
|
/* Query with one token from each batch */
|
|
float *query_emb = &all_embeddings[(b * batch_size + sample_step - 1) * dim];
|
|
int64_t shape_q[] = {dim};
|
|
tm_tensor *query_t = tm_from_data(query_emb, shape_q, 1, TM_F32);
|
|
|
|
fm_results *results = fm_query(store, query_t, 0, top_k);
|
|
tm_free(query_t);
|
|
|
|
if (results) {
|
|
/* Check if any result matches batch b (by checking similarity) */
|
|
float *batch_pooled = (float*)malloc(dim * sizeof(float));
|
|
pool(&all_embeddings[b * batch_size * dim], batch_size, dim, batch_pooled);
|
|
|
|
for (int r = 0; r < results->count; r++) {
|
|
float *rk = tm_data_f32(results->items[r].keys);
|
|
int n_ret = (int)tm_size(results->items[r].keys, 0);
|
|
for (int t = 0; t < n_ret; t++) {
|
|
/* Check if this matches our pooled embedding */
|
|
float diff_sq = 0.0f;
|
|
for (int j = 0; j < dim; j++) {
|
|
float d = rk[t * dim + j] - batch_pooled[j];
|
|
diff_sq += d * d;
|
|
}
|
|
if (diff_sq < 1e-4f) { hits++; goto next_query; }
|
|
}
|
|
}
|
|
next_query:
|
|
free(batch_pooled);
|
|
fm_results_free(results);
|
|
}
|
|
queries++;
|
|
}
|
|
|
|
free(pooled);
|
|
free(all_embeddings);
|
|
fm_store_free(store);
|
|
|
|
snprintf(cmd, sizeof(cmd), "rm -rf %s", store_dir);
|
|
system(cmd);
|
|
|
|
return queries > 0 ? (float)hits / queries : 0.0f;
|
|
}
|
|
|
|
int main(void) {
|
|
int dim = 768;
|
|
int top_k = 5;
|
|
int n_batches = 100;
|
|
|
|
int batch_sizes[] = {1, 4, 8, 16, 32};
|
|
int n_sizes = sizeof(batch_sizes) / sizeof(batch_sizes[0]);
|
|
|
|
const char *pool_names[] = {"mean", "max", "first", "last"};
|
|
pool_fn pool_fns[] = {mean_pool, max_pool, first_token_pool, last_token_pool};
|
|
int n_pools = 4;
|
|
|
|
fprintf(stderr, "==========================================\n");
|
|
fprintf(stderr, "HDC Memory Viability — Test 9\n");
|
|
fprintf(stderr, "Mean-Pooling Information Loss\n");
|
|
fprintf(stderr, "==========================================\n");
|
|
fprintf(stderr, "Embedding dim: %d\n", dim);
|
|
fprintf(stderr, "Batches: %d\n", n_batches);
|
|
fprintf(stderr, "Top-k: %d\n\n", top_k);
|
|
|
|
/* Results table */
|
|
fprintf(stderr, "%-8s", "batch");
|
|
for (int p = 0; p < n_pools; p++)
|
|
fprintf(stderr, " %-10s", pool_names[p]);
|
|
fprintf(stderr, "\n");
|
|
for (int i = 0; i < 8 + n_pools * 12; i++) fprintf(stderr, "-");
|
|
fprintf(stderr, "\n");
|
|
|
|
float results[5][4]; /* [batch_size_idx][pool_idx] */
|
|
|
|
for (int s = 0; s < n_sizes; s++) {
|
|
fprintf(stderr, "%-8d", batch_sizes[s]);
|
|
for (int p = 0; p < n_pools; p++) {
|
|
rng_state = 42 + s * 100 + p * 10; /* reproducible per experiment */
|
|
float recall = run_recall_experiment(
|
|
pool_fns[p], pool_names[p],
|
|
dim, batch_sizes[s], n_batches, top_k);
|
|
results[s][p] = recall;
|
|
fprintf(stderr, " %-10.1f%%", recall * 100.0f);
|
|
}
|
|
fprintf(stderr, "\n");
|
|
}
|
|
|
|
/* Analysis */
|
|
fprintf(stderr, "\n--- Analysis ---\n");
|
|
fprintf(stderr, "batch=1 is the baseline (perfect recall expected for all methods).\n");
|
|
fprintf(stderr, "Recall drop from batch=1 to batch=32 shows information loss from pooling.\n\n");
|
|
|
|
/* Check if mean-pool maintains >50% recall at batch=8 */
|
|
float mean_recall_8 = results[2][0]; /* batch_sizes[2] = 8, pool_fns[0] = mean */
|
|
float mean_recall_32 = results[4][0];
|
|
|
|
fprintf(stderr, "VERDICT:\n");
|
|
if (mean_recall_8 >= 0.50f) {
|
|
fprintf(stderr, " PASS — Mean-pool recall@5 at batch=8: %.1f%% (>50%%)\n",
|
|
mean_recall_8 * 100.0f);
|
|
} else {
|
|
fprintf(stderr, " CONCERNING — Mean-pool recall@5 at batch=8: %.1f%% (<50%%)\n",
|
|
mean_recall_8 * 100.0f);
|
|
fprintf(stderr, " Consider: store individual tokens or use multi-vector retrieval\n");
|
|
}
|
|
|
|
/* Find best pooling strategy at batch=8 */
|
|
int best_pool = 0;
|
|
float best_recall = results[2][0];
|
|
for (int p = 1; p < n_pools; p++) {
|
|
if (results[2][p] > best_recall) {
|
|
best_recall = results[2][p];
|
|
best_pool = p;
|
|
}
|
|
}
|
|
fprintf(stderr, " Best strategy at batch=8: %s (%.1f%% recall)\n",
|
|
pool_names[best_pool], best_recall * 100.0f);
|
|
fprintf(stderr, " Mean-pool at batch=32: %.1f%%\n", mean_recall_32 * 100.0f);
|
|
|
|
fprintf(stderr, "\n==========================================\n");
|
|
return 0;
|
|
}
|