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

EMNLP 2026 Main

Dohoon Kim¹, Junho Park¹, Hyunwoong Bae², Seyoon Kim^2,3, Taesup Moon^1,2,4†

¹Department of ECE, Seoul National University ²Interdisciplinary Program in AI, Seoul National University
³Mobile eXperience Business, Samsung Electronics ⁴ASRI/INMC/AIIS, Seoul National University

^†Corresponding author

Paper Code AxisCF Data

TL;DR. Frozen embedding models barely notice when a memory's time, place, person, or event changes. Axis-LoRA adds four axis-specific LoRA branches inside the model, weighted per query, and raises 4-benchmark mean R@10 on all three embedding models we test.

✓ One LoRA branch per axis ✓ ≤ 0.6% of base parameters ✓ 4,196 synthetic training triplets

Left: a frozen retriever ranks four dentist appointments at nearly equal distance, so the query's time cue does not lift the correct May visit. Right: Axis-LoRA adds four axis LoRA branches to the last 20% of layers and a gate that weights them per query.

(a) A frozen retriever places same-content memories at nearly equal distance, so the time cue in the query does not separate the right memory from its neighbors. (b) Axis-LoRA adds four axis-conditional LoRA branches inside the embedding model and a per-query sparsemax gate that decides how much each axis contributes.

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.

Do Embedding Models Encode the Four Axes?

We describe a personal memory by four axes: TEMPORAL (dates, times), SPATIAL (places), ENTITY (people, objects), and EVENT (actions, event types). An axis counterfactual changes the value of one axis and keeps the rest of the memory verbatim, which isolates how much a model responds to that axis alone.

ORIGINAL	[2023-11-22 19:00] Mom: Thanksgiving dinner went well. Everyone loved the turkey, and we played board games until midnight.
TEMPORAL	[2023-12-01 10:00] Mom: Thanksgiving dinner went well. Everyone loved the turkey, and we played board games until midnight.
SPATIAL	[2023-11-22 19:00] Mom: Thanksgiving dinner went well at the office. Everyone loved the turkey, and we played board games until midnight.
ENTITY	[2023-11-22 19:00] Dad: Thanksgiving dinner went well. Everyone loved the turkey, and we played board games until midnight.
EVENT	[2023-11-22 19:00] Mom: Thanksgiving lunch went well. Everyone loved the turkey, and we played board games until midnight.

A training triplet from AxisCF: one memory and one counterfactual per axis, with the changed span underlined.

An axis edit barely moves the embedding

Violin plots of cosine distance between a memory and its axis counterfactual for three embedding models. TEMPORAL sits at the one-word edit floor on every model; the other axes are higher but small.

Cosine distance between a memory and its axis counterfactual. The dashed line is the distance caused by swapping one non-axis word.

Across MiniLM-L6-v2, bge-base-en-v1.5, and Qwen3-Embedding-0.6B, swapping one axis moves a memory's embedding far less than the typical distance between two memories. TEMPORAL is the weakest axis. A date swap moves the embedding about as much as replacing one unrelated word, and SPATIAL, ENTITY, and EVENT clear that floor but stay small. A per-axis projection trained on the frozen output recovers part of the TEMPORAL signal, but retrieval by date stays far from useful.

Axis response concentrates in the late layers

Heatmaps of per-layer, per-head response to axis edits for three models and four axes. Bright cells cluster in later layers. TEMPORAL and EVENT spread over many heads, SPATIAL and ENTITY over fewer.

Per-(layer, head) response to single-token axis edits. Rows within a panel are layers (input at top), columns are heads.

We measure how much each attention head's pooled output changes when a single token of one axis is edited. On all three models the response concentrates in the later layers. How widely it spreads across heads depends on the axis. TEMPORAL is spread across the most heads, EVENT comes next, and SPATIAL and ENTITY are concentrated in the fewest. This points to adapting the late layers of the model internally rather than post-processing its output.

Axis-LoRA

The embedding model stays frozen. Axis-LoRA adds a gated per-axis term to its similarity score:

\[ \mathrm{score}(q, m) = \langle \tilde{\mathbf{z}}_b(q),\, \tilde{\mathbf{z}}_b(m) \rangle + \alpha \sum_{A \in \mathcal{A}} w_A(q)\, \langle \tilde{\boldsymbol{\phi}}_A(q),\, \tilde{\boldsymbol{\phi}}_A(m) \rangle \]

\(\mathbf{z}_b\) is the frozen embedding model, \(\boldsymbol{\phi}_A\) the same model with the LoRA branch of axis \(A\), \(\tilde{\cdot}\) denotes \(\ell_2\) normalization, \(w_A(q)\) is the per-query gate weight, and \(\alpha = 0.25\).

Four axis branches

One LoRA adapter per axis (rank 16) on the query and value projections of the last 20% of attention layers. The branches share no parameters, so each axis learns its own direction.

Per-query gate

An MLP on the frozen query embedding, followed by sparsemax, gives exact zeros to axes the query does not need. The gate temperature anneals from 2.0 to 0.5 during training.

Training

InfoNCE on AxisCF, where each memory's own counterfactuals are hard negatives, plus a distillation anchor that keeps every branch close to the base model on MS MARCO passages.

Results

R@10 of the frozen embedding model and Axis-LoRA (mean of 4 seeds). Small numbers are the change over the frozen model in percentage points.

Embedding model	Method	LongMemEval	ATM-Bench Normal	ATM-Bench Hard	LoCoMo	4-bench Mean
MiniLM-L6-v2	Frozen	0.4741	0.6465	0.2806	0.4422	0.4608
MiniLM-L6-v2	Axis-LoRA	0.4872+1.31	0.6571+1.06	0.3051+2.45	0.4564+1.42	0.4765+1.56
bge-base-en-v1.5	Frozen	0.6099	0.7286	0.2911	0.6244	0.5635
bge-base-en-v1.5	Axis-LoRA	0.6397+2.98	0.7368+0.82	0.3148+2.37	0.6178−0.66	0.5773+1.38
Qwen3-Embedding-0.6B	Frozen	0.7174	0.7544	0.2994	0.5856	0.5892
Qwen3-Embedding-0.6B	Axis-LoRA	0.7238+0.64	0.7420−1.24	0.3370+3.75	0.6053+1.97	0.6020+1.28

Axis-LoRA raises the 4-benchmark mean on all three embedding models while training at most 0.6% of the base parameters on 4,196 AxisCF triplets. ATM-Bench Hard improves on all three models (+2.37 to +3.75 pp).

Confidence intervals

	MiniLM-L6-v2	bge-base-en-v1.5	Qwen3-Embedding-0.6B	ATM-Bench Hard (3 models pooled)
Gain (pp)	+1.56	+1.38	+1.28	+2.86
95% CI	[+0.54, +2.70]	[+0.25, +2.39]	[+0.15, +2.52]	[+0.55, +5.36]

Gain in 4-benchmark mean R@10, with query-level bootstrap 95% CIs (20k resamples). ATM-Bench Hard has 31 queries, so its gain is tested pooled over the three embedding models.

BibTeX

@inproceedings{kim-etal-2026-axis,
    title = "Axis-{L}o{RA}: Probing and Adapting Embedding Models for Personal-Memory Retrieval",
    author = "Kim, Dohoon  and
      Park, Junho  and
      Bae, Hyunwoong  and
      Kim, Seyoon  and
      Moon, Taesup",
    editor = "Lee, Hung-Yi  and
      Dev, Sunipa  and
      Li, Jessy  and
      Elliott, Desmond",
    booktitle = "Proceedings of the 2026 Conference on Empirical Methods in Natural Language Processing",
    month = oct,
    year = "2026",
    address = "Budapest, Hungary",
    publisher = "Association for Computational Linguistics",
    url = "https://aclanthology.org/2026.emnlp-main.542/",
    pages = "12473--12495",
    ISBN = "979-8-89176-459-0"
}

Axis-LoRA: Probing and Adapting Embedding Modelsfor Personal-Memory Retrieval