For now, I tried a few small experiments in Colab:
I think there is a real, testable core here. In one pinned T4 setup, I was able to extract all of the token-attention tensors from jinaai/jina-embeddings-v2-base-en, map token spans back to chunks, and build chunk-to-chunk matrices from them. So the basic extraction path is workable, at least under one reproducible configuration.
The main open question looks less like whether a graph can be built and more like what the edge is supposed to mean.
In the small probes I ran, an all-layer/all-head average behaved mainly like a document-proximity graph. A few individual heads looked much more interesting on controlled synthetic examples, but those head sets did not transfer cleanly to a small held-out natural-text test based on HotpotQA distractors. On that natural sample, query-conditioned dense retrieval was much stronger at useful candidate budgets.
My current default route would therefore be:
- Define and inspect the edge before adding traversal or updates. Keep the full
[layer, head, source_chunk, target_chunk] tensor, define token-to-chunk pooling explicitly, and preserve directionality initially.
- Establish incremental value over simple controls. At minimum: token distance, independent chunk cosine, late-chunk cosine, order permutation, lexical/entity hard negatives, random expansion, and dense query retrieval.
- Evaluate traversal as a budget curve, not one chosen budget. Separate oracle entry points from query-derived entry points, and compare tight budgets with nearly exhaustive retrieval.
- Treat embedding mutation as a separate experiment. Preserve the original text and original embedding, and compare no-update, routing-only, edge-reinforcement, and content-vector updates independently.
A compact decision tree might be:
What should an edge mean?
├─ document-local context / discourse proximity
│ └─ raw attention may be useful; compare directly with distance and order controls
├─ broad semantic association
│ └─ compare with independent and late-chunk embedding cosine
├─ provenance / “this node was derived from that source”
│ └─ generation-time attention or explicit relation extraction is easier to interpret
└─ relevance to the current query
└─ use a query-conditioned score rather than relying only on a persistent edge
Does the edge add value beyond the controls?
├─ no
│ └─ keep dense retrieval, or use the graph for organization / inspection only
└─ yes
├─ only at large candidate budgets
│ └─ calibrate against random and near-exhaustive retrieval
└─ at tight budgets and on held-out data
└─ proceed to traversal design, then test mutation as a separate ablation
What I tested in Colab, and the limitations
Pinned execution path
The main runs used:
- model:
jinaai/jina-embeddings-v2-base-en
- model commit:
322d4d7e2f35e84137961a65af894fda0385eb7a
- remote-code commit:
f3ec4cf7de7e561007f27c9efc7148b0bd713f81
- Transformers:
4.47.1
- PyTorch runtime reported by the run:
2.11.0+cu128
- GPU: Tesla T4
- FP16 model loading
In that setup, output_attentions=True returned 12 layer tensors with 12 heads. The first small document produced a stacked tensor of shape:
[layer=12, head=12, token=106, token=106]
That matches the usual Transformers model-output contract: attention is returned per layer as a token-to-token tensor, normally shaped (batch, heads, sequence, sequence).
Because the model uses remote custom code, reproducing the result may require pinning both the model and code_revision. An older Jina discussion reports a tuple-index error with output_attentions=True, but the pinned run succeeded, so I would treat that as version/backend-dependent rather than a universal blocker.
Probe 1: small synthetic contract and edge-shape check
I used nine short chunks containing:
- nearby but unrelated statements;
- separated but related statements;
- a lexical/entity distractor;
- a temporal update;
- reverse-order, random-order, and chunk-length controls.
The all-layer/all-head mean had semantic-label AUC 0.60, Spearman correlation 0.924 with inverse token distance, and 0.526 with chunk cosine. Its strongest edges were mostly adjacent chunks, including unrelated ones. Reverse order preserved much of the matrix, while random permutation changed it sharply. This is consistent with a strong position/discourse component. Jina V2 uses a symmetric bidirectional ALiBi variant, although order changes discourse as well as position.
A few middle-layer heads looked much less position-dominated and retained the labelled pairs across those small controls. That motivated a second held-out probe rather than treating the uniform mean as the final graph.
Probe 2: held-out templated relations and hard negatives
The second probe used 30 deterministic documents:
- 18 for head selection;
- 12 held out;
- six relation families: coreference, temporal update, causal chain, organizational relation, correction, and process dependency.
Against distance-matched unrelated negatives, the strongest results on the 12 held-out documents were:
| Edge |
Mean held-out document AUC |
| Train-selected heads, symmetric |
0.9375 |
| Heads preregistered from the first probe, symmetric |
0.9167 |
| Chunk cosine |
0.9167 |
| Uniform attention, symmetric |
0.4583 |
| Inverse token distance |
0.4375 |
Four of the six selected heads overlapped the head set preregistered from the first probe, which was encouraging.
However, the result reversed when I compared the direct positive relation with a lexical/entity hard negative such as a different person who shared the same city, company, role pattern, or event vocabulary:
| Edge |
Positive-vs-hard-negative AUC |
| Chunk cosine |
0.708 |
| Uniform attention, symmetric |
0.708 |
| Preregistered heads, symmetric |
0.208 |
| Train-selected heads, symmetric |
0.125 |
So the selected heads were good at separating a direct relation from a distance-matched unrelated statement, but often scored a topically or entity-related false path above the intended relation.
That changed my interpretation. The signal may be closer to broad association, entity continuity, or a coreference candidate than to a universal typed edge such as “causes,” “supersedes,” or “is the next required process step.”
The retrieval result was also highly budget-sensitive. Starting from an oracle seed and requiring both target nodes:
| Candidate budget |
Selected/preregistered attention |
Cosine graph |
| 3 nodes |
2/12 |
3/12 |
| 4 nodes |
9/12 |
4/12 |
The attention graph often put the useful relation somewhere near the top, but a hard negative could consume one of the very limited slots. Increasing the budget by one changed the conclusion dramatically. That made the proposal’s budget mechanism look like a central experimental variable rather than a secondary tuning parameter.
Probe 3: natural HotpotQA distractors
For a small natural-text check, I used the HotpotQA distractor setting, which provides multiple Wikipedia paragraphs, two supporting documents, sentence-level supporting facts, and both bridge and comparison questions.
I selected 24 eligible examples deterministically:
- train: 6 bridge + 6 comparison;
- held out: 6 bridge + 6 comparison;
- ten paragraphs per question;
- two supporting paragraphs and eight natural distractors.
To materialize full attention on a free T4, each paragraph was capped at 64 tokens, and I retained only examples whose annotated supporting sentences remained inside that cap. This produced a filtered subset, so the result should not be read as a HotpotQA benchmark result.
On the 12 held-out examples, gold-support-pair ranking looked like this:
| Edge method |
Mean gold-pair percentile |
Pooled gold-vs-all-pairs AUC |
| Independently embedded paragraph cosine |
0.669 |
0.632 |
| Uniform attention after subtracting a train-fitted distance component |
0.623 |
0.617 |
| Natural-text train-selected attention heads |
0.475 |
0.483 |
| Late-chunk contextual cosine |
0.473 |
0.466 |
| Synthetic preregistered attention heads |
0.453 |
0.473 |
| Uniform attention |
0.424 |
0.422 |
| Inverse token distance |
0.394 |
0.385 |
The natural-text train-selected head set overlapped only 2/6 with the first synthetic preregistration and 0/6 with the second probe’s selected set. The selected natural heads were also close to chance on held-out pair ranking.
Original-vs-shuffled paragraph-order matrix correlations were approximately:
| Method |
Mean Spearman correlation |
| Independent paragraph cosine |
1.000 |
| Late-chunk contextual cosine |
0.087 |
| Synthetic preregistered attention |
0.009 |
| Natural-selected attention |
-0.008 |
| Uniform attention |
-0.014 |
| Distance-residual attention |
-0.078 |
Again, shuffling changes discourse coherence as well as position, so I would not interpret this as a pure position-invariance test. It does, however, make a raw contextual-attention matrix look fragile as a query-independent, persistent relation between memories whose order may later change.
The query-retrieval result was clearer. The independently embedded query selected one of the two gold paragraphs as its top seed in all 12/12 held-out examples. The fraction for which both gold paragraphs were present in the candidate set was:
| Candidate budget out of 10 paragraphs |
Dense query retrieval |
Paragraph-cosine graph |
Distance-residual graph |
Natural-selected attention |
Uniform attention |
| 2 |
8/12 |
3/12 |
2/12 |
2/12 |
1/12 |
| 3 |
10/12 |
5/12 |
3/12 |
2/12 |
2/12 |
| 4 |
11/12 |
7/12 |
4/12 |
2/12 |
2/12 |
| 5 |
12/12 |
8/12 |
6/12 |
3/12 |
3/12 |
| 6 |
12/12 |
8/12 |
8/12 |
4/12 |
4/12 |
| 7 |
12/12 |
9/12 |
11/12 |
7/12 |
6/12 |
| 8 |
12/12 |
11/12 |
12/12 |
9/12 |
8/12 |
Because one gold paragraph was already the seed, an unstructured expansion that admits budget - 1 of the nine remaining paragraphs has an expected probability of (budget - 1) / 9 of including the second gold paragraph. The measured random-graph curve was close to that expectation. Several graph variants became successful only when most of the ten paragraphs were already admitted.
There was one narrow positive result: for the six bridge questions, the paragraph-cosine graph found both supports in 6/6 by budget 4, versus 5/6 for dense query retrieval. This used embedding cosine rather than attention. For comparison questions, dense query retrieval was clearly stronger: 5/6 at budget 2 and 6/6 at budget 4.
That split seems conceptually useful. A bridge question can connect two pages through a shared entity or intermediary, so pairwise graph expansion can help. Two comparison pages do not necessarily need to resemble or attend to one another; they may be jointly relevant only because the current query asks for the comparison.
Limits of these observations
These were deliberately small sanity checks, not a claim about the general usefulness of attention graphs.
- The first two corpora were synthetic or templated.
- The natural held-out set contained only 12 filtered examples.
- Head selection can overfit the relation definitions and negative set.
- Paragraph shuffling changes both order and discourse coherence.
- I evaluated edge ranking and candidate retrieval, not final generated-answer quality.
- I did not test learned aggregation, query-conditioned head weights, explicit entity links, a reranker trained on hard negatives, or a large memory stream.
- I deliberately did not test embedding mutation yet, so none of these results evaluate the proposed update rule.
The token-attention to chunk-edge contract
One distinction that seems important here is that Late Chunking and attention pooling are two different operations.
The Late Chunking reference implementation first obtains contextual token representations for the whole input, then pools the token representations inside each chunk span to produce chunk embeddings.
By contrast, output_attentions=True produces token-to-token matrices. A chunk graph therefore needs a second, explicit aggregation contract. A simple version is:
token attention:
A[layer, head, source_token, target_token]
chunk pooling:
E[layer, head, source_chunk, target_chunk]
= aggregate A[layer, head, tokens(source_chunk), tokens(target_chunk)]
That leaves several design choices which can materially change the graph:
1. Span mapping
- How are character chunks mapped to token spans?
- Are special tokens excluded?
- What happens when a chunk or supporting sentence is partially truncated?
- Are titles and separators part of the chunk?
2. Pooling
- Mean over the source-span × target-span rectangle?
- Sum with source/target-length normalization?
- Maximum or top-k mean?
- Separate treatment for punctuation, stopwords, or entity tokens?
A plain mean is easy to reproduce, but it can dilute a sparse useful relation inside a long chunk. A plain sum creates the opposite length bias.
3. Direction
It may be worth retaining both A → B and B → A before deciding whether to symmetrize. In the synthetic probe, directed and symmetric results were noticeably different.
Possible variants include:
- directed;
- arithmetic mean of both directions;
- maximum of both directions;
- mutual-only edges where both directions exceed a threshold;
- two separately typed directions.
4. Layer/head aggregation
The full intermediate object is useful:
[layer, head, source_chunk, target_chunk]
If it is averaged immediately, it becomes impossible to tell whether a small group of useful heads was cancelled by many positional or delimiter-oriented heads.
Reasonable comparators include:
- all-layer/all-head uniform average;
- layer-wise averages;
- fixed preregistered heads;
- train-only learned softmax weights;
- relation-conditioned weights;
- query-conditioned weights.
The SproutRAG pooling code is a concrete implementation example: it converts token attention into a layer/head/chunk/chunk tensor by span pooling. Its attention aggregator includes uniform aggregation as a baseline and learnable weights across layers and heads. SproutRAG then uses the result to build a tree, so it is not an end-to-end validation of this memory-graph design, but the tensor contract is highly relevant.
5. Normalization and graph sparsification
A matrix used for visualization is not automatically a calibrated transition matrix.
Possible choices include:
- row normalization;
- fixed top-k outgoing edges;
- global or row-wise thresholds;
- matched graph density across baselines;
- calibration per document or per relation type.
This is also where long chunks, dense hubs, and nearly uniform low-valued edges can create unexpected traversal behavior.
Related systems that seem close to different parts of the idea
I found several systems that seem useful as implementation or vocabulary references. I would not call any of them identical to this proposal; each constrains the meaning of its attention relation differently.
PECAN
PECAN constructs a hierarchical weighted graph with edges derived from LLM attention and uses dynamic progress control to vary how much information is retrieved for each query. Its repository is available here.
It shows that an attention-derived graph can be part of a working retrieval architecture. The difference is edge semantics: PECAN’s attention is tied to creating higher-level Information Points from source information, so the source-dependence interpretation is narrower than treating arbitrary bidirectional encoder self-attention as a general persistent memory relation.
SAKI-RAG
SAKI-RAG uses a SentenceAttnLinker to model inter-sentence attention relationships and a Dual-Axis Retriever that expands and filters candidates using both semantic similarity and contextual relevance.
That dual-axis design seems especially relevant here. It does not require one persistent attention edge to simultaneously represent:
- structural/contextual association;
- semantic similarity;
- and relevance to the current query.
SproutRAG
SproutRAG converts sentence-level attention into a hierarchical tree for long-document retrieval. Its code makes several hidden choices explicit:
- token-span to sentence-span pooling;
- a complete layer/head/chunk/chunk tensor;
- diagonal masking;
- direction symmetrization;
- uniform aggregation as an initialization/baseline;
- optional training of layer/head weights;
- explicit retrieval and reranking stages.
It builds a tree rather than an evolving general graph, but it is probably the closest reusable code example for the token-attention-to-chunk-relation part.
Late Chunking
The Late Chunking paper and implementation are useful for the chunk-embedding side: the chunk vector can retain whole-document context because pooling happens after contextualization.
That does not automatically make the token-attention matrix a sentence-level relation matrix. The two signals can be compared or combined, but they should be measured separately.
Graph propagation references
If the graph itself proves useful, systems such as HippoRAG provide another useful comparison: query-related seeds are introduced into a graph, then relevance is propagated rather than treating the stored edge as the complete retrieval decision.
Possible meanings of an edge, and the controls for each
A single scalar attention_weight(A, B) may be asked to represent several different things. I suspect it will be easier to evaluate if the intended relation is stated first.
| Intended edge meaning |
Plausible signal |
Essential control |
Typical failure mode |
| Local discourse/context proximity |
Raw contextual attention |
Token/chunk distance and order permutations |
Adjacent but unrelated text dominates |
| Broad semantic association |
Independent cosine, late-chunk cosine, selected attention heads |
Lexical/entity hard negatives |
Same topic or entity type mistaken for the intended relation |
| Coreference/entity continuity |
Selected heads plus explicit entity identity |
Same-type, same-city, or same-role distractors |
Two different entities are conflated |
| Temporal update or supersession |
Typed relation plus timestamps/validity |
Old/new fact controls |
Both versions blend into one representation |
| Provenance or derivation |
Generation-time attention or explicit extraction |
Source attribution |
Edge has no readable causal/source meaning |
| Current-query relevance |
Query-conditioned scorer or transition weight |
Dense query retrieval |
Fixed graph misses comparison-style relevance |
This suggests a multi-edge or multi-signal design may be more natural than a single universal matrix. For example:
- persistent context edge;
- entity/coreference edge;
- temporal
supersedes edge;
- provenance edge;
- query-conditioned transition score.
An edge that is useful for candidate expansion need not also be best for final ranking, explanation, or memory update.
A low-cost evaluation sequence would label a small held-out set for one relation, separate ordinary from hard negatives, compare distance/cosine/attention/random edges at matched density, test alternative order, and only then evaluate traversal.
Traversal: budget, queue discipline, and baselines
The decaying-budget idea seems testable, but the search rule matters.
If every edge traversal has exactly the same cost, literal breadth-first search is coherent. If each edge has a different penalty and the objective is to preserve the most remaining budget, a priority queue that expands the best current path is a natural comparator. This is analogous to the usual distinction in NetworkX’s shortest-path guidance: BFS for unweighted hop distance and Dijkstra-like methods for non-negative weighted cost.
The objective may be spreading activation rather than one shortest path, but making the frontier rule explicit would make the comparison easier.
Questions that change the result include:
- Is path cost additive, multiplicative, or based directly on remaining budget?
- If a node is reached again through a stronger path, is it reopened?
- Can a strong cycle repeatedly reinforce the same region?
- Is there a maximum out-degree or edge threshold?
- Is budget counted in nodes, tokens, total context characters, or computation?
- Are nodes ranked by best path, sum of paths, or first visit?
The evaluation should separate seed quality, edge quality, traversal quality under a fixed budget, and final answer quality.
Useful controls are:
- oracle seed vs query-derived seed;
- dense query top-k;
- the same graph with random weights or permuted edges;
- a random expansion with the same candidate count;
- a budget curve from tight retrieval to nearly exhaustive retrieval;
- context tokens consumed, not only node count.
The budget curve is particularly important. In the HotpotQA probe, dense query retrieval found both supporting paragraphs in 8/12 examples at budget 2 and 12/12 by budget 5. Several graph variants only approached full recall at budget 8 of 10 paragraphs. Reporting only the large-budget endpoint would make a weak expansion method look much stronger than it was at useful context sizes.
Why I would separate persistent association from query-conditioned relevance
The bridge/comparison split in the small HotpotQA probe gave a concrete example of this distinction.
For a bridge question, the first supporting paragraph may mention an intermediary entity that connects naturally to the second paragraph. A persistent entity, semantic, or discourse edge can help expand from one support to the other.
For a comparison question, the two supporting paragraphs may have little pairwise similarity. They are connected because the query asks to compare two things. A query-independent edge between the paragraphs is not necessarily expected to be strong.
That leaves several reasonable architectures:
Option A: persistent graph for candidates, query scorer for ranking
- Dense query retrieval chooses one or more seeds.
- Persistent graph edges expand structural candidates.
- Query relevance and hard-negative reranking choose the final context.
Option B: typed persistent edges with query-conditioned transitions
- Store entity, temporal, provenance, and discourse relations separately.
- The current query decides which edge types and directions receive weight.
Option C: use graph expansion only for bridge-like uncertainty
- If dense retrieval already returns several independently relevant supports, keep it.
- If the query appears to require an intermediary relation, expand through the graph.
Option D: route between dense-only, graph, and hybrid retrieval
- The routing signal could be query type, score margin, retrieval entropy, or failure to find enough independent evidence.
The probes do not identify the best option, but they suggest that a fixed raw attention graph should not be expected to replace query relevance by itself.
The embedding-update part looks like a separate, higher-risk experiment
The update rule is interesting, but it changes a different contract from graph construction.
Moving an existing content embedding toward a newly retrieved embedding changes both:
- how the node will be routed to later;
- and the vector that is supposed to represent the original evidence.
That creates several possible feedback loops:
- a mistaken retrieval changes the node so it becomes even easier to retrieve incorrectly next time;
- frequently visited nodes drift into general-purpose hubs;
- old and new facts become geometrically blended even when the newer fact supersedes the older one;
- the current embedding may no longer be reproducible from the stored source text.
I did not implement the update rule, so these are reasons to isolate the experiment rather than observed failures.
A safer decomposition could be:
- immutable raw text / episode;
- reproducible original content embedding;
- mutable routing embedding;
- mutable typed edge weights;
- timestamps and validity intervals;
- update lineage;
- regression tests and rollback.
Graphiti is a useful implementation reference for this separation. It keeps raw episodes as provenance, gives facts temporal validity windows, preserves superseded history, and combines semantic, keyword, and graph retrieval. It is a much heavier architecture than the proposal here, but the provenance and temporal boundaries are relevant even in a small implementation.
LongMemEval is also a useful evaluation map because it separates long-term memory into information extraction, multi-session reasoning, temporal reasoning, knowledge updates, and abstention. A retrieval improvement on the newest query does not by itself show that older memories remain correct.
A clean mutation ablation might be:
- no update;
- edge reinforcement only;
- mutable routing-vector update while preserving the original content vector;
- direct content-vector nudge;
- gated update, accepted only if both the new query and a regression set of older queries improve;
- rollback when they do not.
I would keep the retrieval graph frozen while comparing those variants. Otherwise an edge change, traversal change, and vector mutation can compensate for one another and make the source of an improvement difficult to identify.
Implementation and reproducibility checklist
For this specific idea, I think the following small record would make results much easier to compare across machines and future revisions:
- exact model ID;
- model commit SHA;
- remote-code repository and
code_revision SHA;
- Transformers and PyTorch versions;
- attention backend actually used by each module;
- returned attention shapes;
- input token count and truncation status;
- chunk character spans and token spans;
- special-token handling;
- token-to-chunk pooling formula;
- directed or symmetric edge rule;
- layer/head aggregation rule;
- normalization and graph sparsity rule;
- dataset IDs and selection filters;
- train/held-out split used for head selection;
- ordinary and hard-negative definitions;
- random seed;
- seed-selection rule;
- candidate budgets;
- random and dense baselines;
- candidate recall, final reranked recall, answer quality, and context cost as separate metrics.
Saving the config, chunk spans, full layer/head chunk tensor, embeddings, query scores, edge metrics, budget curve, and environment/errors would allow many later aggregation and traversal experiments to run on CPU without another GPU quota.
Possible next routes, depending on the goal
If the immediate goal is a conceptual prototype
The smallest convincing demonstration may be:
- choose one intended edge meaning;
- show the directed per-head chunk matrix for one readable document;
- compare it with token distance and cosine;
- include one ordinary negative and one hard negative;
- show how the result changes under order permutation.
That is enough to establish what the edge is doing without building the full evolving memory system first.
If attention tensors and pooling are already implemented
The highest-value additions would be:
- retain and optionally publish the full layer/head matrix rather than only the mean;
- document the span-pooling and normalization rules;
- treat uniform averaging as a baseline;
- add held-out hard negatives;
- compare directed, symmetric, and learned/query-conditioned aggregation.
If the goal is retrieval performance
I would keep dense query retrieval as the primary baseline and test whether the graph adds candidates which dense retrieval misses at the same context budget.
The most informative measurements would separate bridge-like from comparison-like queries and report budget curves, oracle/query seeds, random expansion, candidate recall, post-reranking recall, answer accuracy, and context-token cost.
If the goal is long-term evolving memory
I would add provenance, timestamps, and supersession before direct embedding mutation. The first update experiment could change only routing state or edge strength while leaving evidence immutable.
If the goal is visualization or interpretability
Directed per-head attention graphs may still be useful even if they are not the best retriever. A graph can be diagnostically valuable without also serving as the final persistent semantic graph.
Overall, I think the proposal is experimentally tractable, and the closest related systems suggest that attention-derived relations can be useful when the edge semantics, aggregation, query conditioning, and traversal objective are explicit.
The small results changed my own default from “average the attention and then tune the graph” to:
retain the full tensor, define one relation at a time, test it against strong controls and hard negatives, measure the whole budget curve, and only then decide which part deserves to become persistent memory state.
Even if the raw attention ultimately turns out to represent mostly local discourse structure rather than a universal semantic relation, that could still be a useful edge type. It just may need to live alongside semantic, temporal, provenance, and query-conditioned signals rather than replacing all of them with one scalar.