📓 Notebook: attns / attn_ranker.ipynb

The goal is a lightweight "head" that acts as a gatekeeper for a temporary knowledge base, sitting between a vector database and a large language model.

The motif, rephrased from my heart

Take the simplest case: remembering things from one chat session between a user and a bot. Assume that nothing matters except these three parts:

1. Base model: the LLM that generates the dialogue.
2. Session cache: short-term memory, managing KV-caches or attention patterns so the LLM stays coherent and in persona.
3. Datastore: long-term memory, like a RAG-style vector store holding the bot's "passport" or biography.

The objective is to build dynamic passages between these three, so data cycles through them. At its simplest that means syncing short-term and long-term memory regularly. The harder version is picking the features that meet the latent space of large models.

The other route I considered was a Hopfield network trained with Hebbian learning, combined with an HMM or PageRank for ranking. Hebbian learning is appealing because its weights are already an adjacency matrix, the same thing graph models work on. This notebook tries plain attention first.

Ordinal targets

A rank of 2 out of 5 becomes [1, 1, 0, 0, 0]. That way the model is predicting "how far up the ladder" an item goes, rather than one class out of five:

create_ordinal_target([2, 4], 5)
# tensor([[1, 1, 0, 0, 0],
#         [1, 1, 1, 1, 0]])

Component	Description	Shape
Q, K, V	Input sequences for attention	(batch_size, seq_len, emb_dim)
target	Ordinal target labels in binary format	(batch_size, 1, num_ranks)
output	Model predictions (sigmoid outputs)	(batch_size, 1, num_ranks)

The Ranker

nn.MultiheadAttention (2 heads, no bias, dropout 0.6), then a softmax and a linear layer to num_ranks, summed over the sequence and passed through a sigmoid. Nearly every line has a shape assert after it, because I kept getting the shapes wrong.

attn_output, attn_weights = self.attn(Q, K, V)
output = attn_output.softmax(-1)
x = self.fn(output)                  # (batch, seq_len, num_ranks)
x = x.sum(axis=1).unsqueeze(1)       # (batch, 1, num_ranks)
return x.sigmoid()

Trial run

Synthetic data from make_classification: batch 10, seq_len 25, emb_dim 50, 5 ranks, split into Q/K/V. Trained with Adam (lr 1e-3) and MSE loss for 100 epochs. A prediction counts as a 1 when the output is above 0.65.

	Predicted ranks	Correct (of 10)	Loss
Epoch 1	[1, 1, 1, 1, 1, 1, 1, 1, 1, 1]	2	0.2580
Epoch 100	[4, 4, 0, 2, 1, 1, 2, 3, 0, 3]	10	0.0134

It goes from predicting rank 1 for everything to getting all 10 right.

⚠️ Warning for future self / reader: this is ten rows of synthetic data. Getting 10 out of 10 means the model can fit. It doesn't show that it can rank. The results are probably biased to the synthetic dataset.

Self notes

• Try nn.BCEWithLogitsLoss instead of MSE. If I do, remove the sigmoid, because that loss expects raw logits.
• This is a version of the "small ranking model inspired by PageRank" that ended up inside Gaby, and part of the memory work for Tom in Tom's memory networks.