Axis-LoRA: Probing and Adapting Embedding Models for Personal-Memory Retrieval: what it means for business leaders

Common embedding models barely register a memory's date, place, person or event; Axis-LoRA adds small gated adapters inside a frozen model and improves average personal-memory recall on three benchmarks, though not in every test.

1. What the method is

Axis-LoRA is a way to adapt an existing text embedding model for retrieving personal memories, the records an assistant keeps about what a user did, when, where and with whom. It leaves the base model frozen and adds four separate low-rank adapters, one each for time, place, entity and event, inside the model's late attention layers. A small per-query gate decides which adapters count for a given question, and the final score is the frozen model's similarity plus the gated per-axis similarities. The paper also offers a diagnosis: three common embedding models barely respond to changes in these four attributes, and the authors conclude that the fix belongs inside the model rather than on its output. The adapters train on a small synthetic set of counterfactual memories, and the authors have published their code.

2. Why the method was developed

Memory systems for assistants have invested in storage and hierarchy but, according to the authors, still rely on general-purpose embedding models trained mainly on topic and content overlap. Personal records often differ only in a detail: the same project discussed on Tuesday or Friday, in Boston or Seattle. The authors swapped single attributes and measured the shift in the embedding. Date swaps moved it no more than replacing a random non-attribute word, and all four attribute types moved it far less than the typical distance between two memories. An obvious cheap remedy failed too: a trained two-layer projection on top of the frozen output lifted date-anchored recall somewhat, but its ceiling stayed far short of useful retrieval and its effect on place and event queries was mixed. Instruction prefixes gave only a small lift, which pushed the authors toward changing the encoder itself.

3. Who should care

The ML engineers who own our retrieval and agent-memory stacks are the primary audience, because this is a change to the embedding layer that does not require replacing the base model. Product leaders with assistants that answer questions about a user's history should care because failures on when, where and who are visible to customers and may trace to retrieval rather than to the language model. Teams buying memory products should ask vendors how they test retrieval on attribute-specific queries. Infrastructure and finance owners should note that, as we read the design, memories need per-axis representations beyond the base vector, which would raise index size and re-embedding work; the text we reviewed gives no serving cost. Teams in clinical, legal or other specialist domains should note that the four axes are tied to these benchmarks, so they would need their own schema and training data.

4. How the method works

The authors measured how much each attention head's pooled output changes when one attribute token is edited. Response clustered in late layers, and each attribute drew on its own pattern of heads: date signal was the most diffuse, needing 32.1% to 41.7% of heads to reach the paper's coverage threshold, while entity needed 18.5% to 30.6%. Each of the four adapters therefore attaches to the query and value projections of the late attention layers and shares no parameters with the others. A small neural network reads the query and produces sparsemax weights, which can be exactly zero, so unused adapters contribute nothing. Training is one joint stage with two losses: a contrastive loss that ranks the true memory above versions with one swapped attribute, and a distillation term that keeps each adapter close to the base model on a subset of MS MARCO general-retrieval text.

5. How it was evaluated

The authors tested retrieval on four splits from three benchmarks: LongMemEval and LoCoMo, which cover long multi-session conversations, and ATM-Bench Normal and Hard, which cover multi-source personal data such as images, videos and emails. Hard carries more multi-evidence queries but is small. They applied the method to three embedding models: MiniLM-L6-v2, bge-base-en-v1.5 and Qwen3-Embedding-0.6B. The primary metric was recall at ten, averaged over random seeds, with query-level bootstrap confidence intervals. Training used AxisCF, a synthetic counterfactual corpus generated with Qwen3.5-27B and screened against a blacklist of entities and dates mined from test memories. Comparisons on Qwen3-Embedding-0.6B included axis-specific prompts, a training-free subspace, a trained projection and a single shared adapter, plus removal of each main component. General retrieval was checked with nDCG@10 on six BEIR tasks from MTEB, on that model only.

6. How it performed

The copy of the paper we reviewed lost most inline figures and was truncated, so we do not quote gain sizes; Table 5 of the paper has them. The four-benchmark mean recall at ten improved on all three models, with confidence intervals excluding zero, and the largest single gain came on ATM-Bench Hard. Two of twelve model-benchmark results regressed: Qwen3-Embedding-0.6B on ATM-Bench Normal and bge-base on LoCoMo. The authors report that general retrieval was preserved on the one model tested, and that a shared adapter fell below the unmodified base model. They acknowledge that implicit cues such as relative dates are under-represented in training and call the method a simple baseline. An appendix heading also indicates that the free gate collapses to one branch per encoder, a detail we could not read and would check first. Our call: pilot on our own memory data, measuring answer quality, index growth and latency against the unmodified model. (Source: arXiv 2611.07820, 2026)

← Back to dossier index