Axis-LoRA: Probing and Adapting Embedding Models for Personal-Memory Retrieval

AI summary

Overview

Research area: Natural Language Processing — dense retrieval and embedding-model adaptation for personalized memory question answering.

Technical level: Intermediate (assumes familiarity with embedding models, contrastive retrieval, and low-rank adapters).

Scope: The paper diagnoses why frozen dense embedding models fail to distinguish personal memories along four categorical axes (time, place, person, event), and proposes a per-axis LoRA adapter with a per-query gate to fix it.

What This Paper Is About

When a user asks an LLM something like "What did I discuss with Marco at Tuesday's meeting?", the system must retrieve the right memory from a long personal corpus — and that requires telling apart memories that share content but differ in date, place, person, or event type. The authors show that current off-the-shelf embedding models barely change their output when a memory's date or place is swapped, so the correct memory does not rank above content-similar distractors. The goal is to locate where inside the encoder this signal is lost and to adapt the encoder so that each of these four axes becomes discriminative at retrieval time.

Key Contributions

1. Characterizes "axis-asymmetric under-exposure." The authors show across three embedding models that frozen encoders expose the four memory axes to very different degrees, measured both as axis-swap cosine distance and as axis-anchored R@10 retrieval.
2. Two diagnostic probes that localize the problem inside the model. A projection probe shows that a trained head on top of the frozen embedding cannot close the retrieval gap, and a head-spectrum probe shows each axis is concentrated in a distinct attention-head pattern, mostly in later layers.
3. Proposes Axis-LoRA. Four independent LoRA branches (one per axis) inside the embedding model, restricted to the last 20% of attention layers, combined through a per-query sparsemax gate, scored as frozen-base similarity plus a gated per-axis term.
4. Introduces AxisCF and demonstrates gains. A synthetic axis-counterfactual triplet corpus (4,196 triplets across four memory formats) is used for training, yielding +1.28 to +1.56 pp 4-benchmark mean R@10 across three encoders while training ≤0.6% of base parameters.

Main Findings

• TEMPORAL is essentially invisible to frozen encoders. Against a non-axis 1-word edit floor, TEMPORAL swap distance is 0.34x on MiniLM (floor 0.023), 1.12x on bge-base (floor 0.012), and 1.04x on Qwen3-Embedding-0.6B (floor 0.010) — effectively at noise level. SPATIAL, ENTITY, and EVENT clear the floor by 2.6x to 8.0x, but all four axes sit at only 1% to 14% of the inter-memory distance μ_inter (0.78 for MiniLM, 0.47 for bge-base, 0.65 for Qwen3).

• The pattern holds on real memory data. Repeating the measurement on 200 ATM-Bench, 200 LoCoMo, and 100 LongMemEval memories reproduces it, with TEMPORAL at or below the 1-word edit floor on all three encoders and every axis at ≤11% of μ_inter.

• Axis-anchored retrieval is correspondingly weak. Raw cosine R@10 ranges from 0.039 (Qwen3, TEMPORAL) to 0.312 (Qwen3, EVENT), with TEMPORAL lowest on every encoder at 0.039 to 0.072.

• An external projection helps only a little. A per-axis 2-layer MLP trained with InfoNCE lifts TEMPORAL R@10 by +8.9 to +14.4 pp, but the absolute ceiling stays at 0.15 to 0.18. ENTITY lifts +2.5 to +6.1 pp; SPATIAL and EVENT are mixed, with drops as large as −8.1 pp.

• Each axis has its own head footprint, concentrated in late layers. The fraction of attention heads covering 50% of per-axis response is 41.7%/36.1%/29.2%/26.4% (TEMP/EVENT/SPAT/ENT) on MiniLM, 39.6%/35.4%/29.2%/30.6% on bge-base, and 32.1%/25.4%/19.6%/18.5% on Qwen3-0.6B. TEMPORAL is the most diffuse axis and SPATIAL/ENTITY the most concentrated, consistently across models.

• Axis-LoRA improves the 4-benchmark mean on every encoder. +1.56 pp on MiniLM-L6-v2 (0.4608 → 0.4765), +1.38 pp on bge-base-en-v1.5 (0.5635 → 0.5773), +1.28 pp on Qwen3-Embedding-0.6B (0.5892 → 0.6020), as 4-seed means with cross-seed std ≤ 0.017 per cell.

• The largest per-benchmark gain is on the multi-evidence ATM-Bench Hard split. +3.75 pp on Qwen3-0.6B, +2.45 pp on MiniLM, +2.37 pp on bge-base. Two of twelve (encoder, benchmark) cells regress: Qwen3/ATM-Normal −1.24 pp and bge-base/LoCoMo −0.66 pp, which the appendix attributes to a narrow rank-10 boundary effect.

• Gains are statistically significant. Bootstrap 95% confidence intervals over 20k query resamples are [+0.54, +2.70] on MiniLM, [+0.25, +2.39] on bge-base, and [+0.15, +2.52] on Qwen3-0.6B. The pooled ATM-Hard lift is +2.86 pp [+0.55, +5.36] (n = 31 queries).

• General retrieval is largely preserved. On a 6-task MTEB BEIR subset (NFCorpus, SciFact, ArguAna, SCIDOCS, FiQA2018, TRECCOVID), nDCG@10 drops from 0.5449 to 0.5399 (−0.50 pp) with Qwen3-Embedding-0.6B.

• Where the adaptation happens matters more than that it happens. Axis-prefix prompting gives +0.9 pp (avg 0.596), TF-PCA gives −0.4 pp (0.583), an external MLP+InfoNCE gives +0.5 pp (0.592), and a single shared-trunk LoRA collapses to 0.493 (−9.6 pp below Base). Full Axis-LoRA reaches 0.602.

• Every component earns its place. Removing axis-CF supervision drops the average to 0.593 (−0.9 pp), removing the MS MARCO KD anchor drops it to 0.588 (−1.4 pp), and replacing the learned gate with uniform routing drops it to 0.599 (−0.3 pp).

Methodology in Plain English

The authors start by building a controlled test set called AxisCF. For each personal memory, they create four counterfactuals that change only one thing: the date, the place, the person, or the event type. Keeping everything else identical means any change in the embedding must come from that single axis value.

They then measure, on three off-the-shelf embedding models ranging from a 22M-parameter BERT-style model to a 600M-parameter LLM-based one, how far a memory's embedding moves when one axis is swapped. They compare that movement against two yardsticks: how much the embedding moves when a single ordinary content word is replaced, and how far apart two different memories sit in the same space.

To find out whether the missing information could be recovered after the fact, they train a per-axis projection head on top of the frozen embeddings using contrastive learning and see how much retrieval improves. To find out where the information lives, they perturb the input by one axis token at a time and record how much each individual attention head's pooled output changes, producing a map of layers against heads for each axis.

Those probes point inside the model, so the fix operates there. Axis-LoRA attaches four separate low-rank adapters — one per axis — to the query and value projections of the last 20% of attention layers, with rank 16. A small network reads the frozen query embedding and outputs four weights through sparsemax, which forces exact zeros so irrelevant axes contribute nothing; the temperature of this gate is annealed from 2.0 down to 0.5 during training so all four axes get early signal. The final retrieval score is the frozen model's similarity plus 0.25 times the gated sum of the four per-axis similarities.

Training is a single joint stage with the base model frozen. The main loss is contrastive over AxisCF, where the axis counterfactuals serve as hard negatives, and a distillation term pulls each adapter back toward the base model on MS MARCO text so the adapters do not corrupt general-purpose embeddings.

Why This Matters

Impact on research. The paper reframes personal-memory retrieval as a structural problem — memories are indexed along categorical axes, and general-purpose embedding training optimized for topic overlap under-trains exactly those attributes. The two probes give a reusable diagnostic template: test whether an output-side head can recover a signal, and if not, map which attention heads carry it before deciding where to intervene. The ablation showing shared-trunk LoRA collapses to 0.493 (−9.6 pp below base) is a concrete warning that per-axis structure must live inside the model, not just at its output.

Real-world applications:

• Personal AI assistants that answer questions about a user's own history across long chat logs, where distinguishing "Tuesday's meeting" from "Friday's meeting" is the whole task.
• Multimodal personal archives spanning images, videos, and emails, the setting ATM-Bench targets, where the Hard split's multi-evidence queries show the biggest gains.
• Memory layers in agentic systems that store and recall past interactions over months, as in LoCoMo's multi-month conversations.
• Any retrieval pipeline where records share boilerplate content and differ mainly in metadata slots.

Industry relevance. The adaptation is cheap: ≤0.6% of base parameters, 4,196 training triplets, a frozen backbone, and a −0.50 pp cost on general retrieval as measured on the BEIR-6 subset. That combination makes it practical to drop into existing memory pipelines built on standard embedding models rather than replacing the retriever.

Future Directions

• Extending beyond four axes. The current schema is tied to the benchmarks' annotation schemas; new domains such as clinical or legal would need their own axis definitions and counterfactual construction.
• Handling implicit axis cues. AxisCF edits explicit values, so relative time expressions, vague locations, and entangled multi-entity events are under-represented — a known weak spot.
• Richer routing and supervision. The authors describe Axis-LoRA as a simple baseline and point to mixture-of-experts gates, stronger inside-model supervision, and joint training with retrieval-aware objectives.
• Per-user personalization. Whether the same recipe can quickly fit a retrieval model to a single user from a small amount of their own interaction data is posed as a natural next step.

Target Audience

Researchers and engineers working on retrieval-augmented generation, long-term agent memory, or personalization of language systems. It is most useful to readers who already understand contrastive embedding training and want a diagnostic method plus a parameter-efficient fix, and to practitioners deciding whether to adapt an embedding model internally versus bolting on an output-side projection.

Authors’ abstract

As LLMs are increasingly asked questions about a user's personal experiences, retrieval must distinguish memories by time, place, person, and event. Current personal-memory pipelines rely on frozen dense embedding models to embed memories for retrieval, but we find that these models do not capture the four axes well. Swapping a memory's date or place barely shifts its embedding, and a trained projection on top of the frozen embedding does not recover the missing signal. The problem lives inside the model, where each axis is encoded by a distinct attention-head pattern. Guided by these findings, we propose Axis-LoRA, a per-axis adapter that places four LoRA branches inside the embedding model and combines them via a per-query gate. Across three standard embedding models evaluated on LongMemEval, ATM-Bench, and LoCoMo, Axis-LoRA improves 4-benchmark mean R@10 by $+1.28$ to $+1.56$~pp while training only a small fraction of the base parameters on a small axis-counterfactual training set. Code is available at https://github.com/dohoonkim-ai/axis-lora.