Towards Improved Sentence Representations using Token Graphs

Krishna Sri Ipsit Mantri; Carola-Bibiane Schönlieb; Zorah Lähner; Moshe Eliasof

cs.LGMar 3, 2026

Towards Improved Sentence Representations using Token Graphs

Authors: Krishna Sri Ipsit Mantri, Carola-Bibiane Schönlieb, Zorah Lähner, Moshe Eliasof

Organizations: University of Bonn · Lamarr Institute for Machine Learning and Artificial Intelligence · University of Cambridge · Ben-Gurion University of the Negev

Abstract

Obtaining a single-vector representation from a Large Language Model's (LLM) token-level outputs is a critical step for nearly all sentence-level tasks. However, standard pooling methods like mean or max aggregation treat tokens as an independent set, discarding the rich relational structure captured by the model's self-attention layers and making them susceptible to signal dilution. To address this, we introduce GLOT, a lightweight, structure-aware pooling module that reframes pooling as relational learning followed by aggregation. Operating on the outputs of a frozen LLM, GLOT first constructs a latent token-similarity graph, then refines token representations with a graph neural network, and finally aggregates them using a readout layer. Experimentally, our approach is remarkably robust and efficient: on a diagnostic stress test where 90% of tokens are random distractors, GLOT maintains over 97% accuracy while baseline methods collapse. Furthermore, it is competitive with state-of-the-art techniques on benchmarks like GLUE and MTEB with 20x fewer trainable parameters and speeds up the training time by over 100x compared with parameter-efficient fine-tuning methods. Supported by a theoretical analysis of its expressive power, our work shows that learning over token graphs is a powerful paradigm for the efficient adaptation of frozen LLMs. Our code is published at https://github.com/ipsitmantri/GLOT.

Figures & tables

Figure 1: Fine-tuning large language models for sentence embeddings is computationally expensive. Our pooling method, Glot , constructs a latent token-similarity graph from the outputs of a frozen model. It then refines token representations with a graph neural network before aggregation. This technique enables decoder-only models (like Mistral-7B), typically optimized for next-token prediction, to produce powerful sentence-level representations without requiring any fine-tuning.

Figure 2: An overview of the Glot pooling architecture. Given token hidden states from a frozen language model, our trainable module performs three stages : (1) it constructs a latent token-similarity graph, (2) a TOKEN-GNN performs relational learning to refine token representations, and (3) a readout layer aggregates the refined vectors into a final sentence representation, ${\mathbf{z}}$

Figure 3: Robustness to signal dilution on the diagnostic stress test. Each of the four panels displays the classification accuracy for all pooling methods at a specific distractor ratio, which increases from 20% to 90%. Within each panel, backbone models are arranged along the x-axis by their parameter count.

Method	# Trainable Params	GPU Memory (GB) $\downarrow$	Batch Runtime (ms) $\downarrow$	MCC $\uparrow$
Full FT $+$ EOS	7.11B	32.59	$1318.8\pm 1.1$	49.63
LoRA ( $r=64$ ) $+$ EOS	167.8M	33.50	$1454.6\pm 1.1$	48.23
Glot (ours)	8.92M	0.42	$\mathbf{\phantom{2}13.4\pm 3.0}$	53.29

Table 5: A comparison of training methods by resource consumption and performance on the CoLA task, using the Mistral-7B backbone. We contrast our frozen-backbone approach ( Glot ) against full fine-tuning (Full FT) and LoRA. Batch runtime is reported as the mean

\pm

standard deviation over 10 measurements.

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

Hyperparameter	Search Space
Optimization
Learning Rate	{1e-3, 2e-4, 2e-5}
Weight Decay	{0.0, 1e-5, 5e-5}
Token-GNN Architecture
GNN Layers ( $K$ )	{2, 4}
GNN Hidden Dimension	{64, 128, 256}

Appendix

Table 6: Hyperparameter search space for the Glot pooling head. The final model configuration was determined via a grid search over these values. The search was performed consistently across all backbone models and datasets.

Figure 4: Z-score normalized performance on the GLUE benchmark, aggregated by task category. Performance, represented as a z-score, is plotted against the number of parameters in the frozen backbone model (log scale). A higher z-score indicates better relative performance compared to the average of all tested methods for that setting.

Model	Method	20% Distractors	50% Distractors	80% Distractors	90% Distractors
BERT	[CLS]	70.8	58.2	57.2	67.6
	Mean	68.0	58.6	64.2	53.4
	Max	57.4	50.8	51.6	50.4
	AdaPool	91.4	78.8	65.6	61.6
	Glot	97.2	97.0	97.8	98.8
RoBERTa	[CLS]	83.6	63.4	51.6	48.6

Appendix

Table 7: Full results for the diagnostic stress test , which evaluates robustness to signal dilution. The table reports the classification accuracy for all pooling methods across six backbones as the ratio of distractor tokens in the input sequence increases from 20% to 90%. The best results for each model are in bold.

Threshold ( $\tau$ )	STSB (Spear. $\uparrow$ )	CoLA (MCC $\uparrow$ )	RTE (ACC $\uparrow$ )	Linear Probing (ACC $\uparrow$ )
BERT (Encoder-only)
$\tau=0.0$	81.88	39.62	50.90	76.41
$\tau=0.2$	82.12	39.82	50.90	76.51
$\tau=0.4$	82.25	47.49	52.34	77.18
$\tau=0.6$	83.86	43.16	59.21	77.56
$\tau=0.8$	83.85	43.16	52.70	77.18

Appendix

Table 8: Impact of Graph Topology (

\tau

) on Performance and Linearity. We compare the sensitivity of BERT (top) and Mistral (bottom) to the threshold

\tau

. Linear Probing indicates the accuracy of a linear SVM trained on frozen Glot embeddings, serving as a proxy for linear separability.

Method	Trainable	GPU Mem.	Runtime	CoLA	STSB	RTE
Method	Params	(GB) $\downarrow$	(ms) $\downarrow$	(MCC) $\uparrow$	(Spear.) $\uparrow$	(ACC) $\uparrow$
Mistral-7B (Decoder-only)
Full FT + [EOS]	7.11B	32.59	$1318.8\pm 1.1$	49.63	55.68	55.23
LoRA ( $r=64$ )	167.8M	33.50	$1454.6\pm 1.1$	48.23	54.54	53.43
Glot (Ours)	8.92M	0.42	13.4 $\pm$ 3.0	53.29	80.51	59.21
BERT (Encoder-only)

Appendix

Table 9: Cross-Model Efficiency Benchmarks. We compare Glot against Full Fine-Tuning and LoRA. MCC : CoLA, Spear. : STSB, ACC : RTE. Runtimes are measured with batch size 32. Glot provides consistent efficiency gains across architectures.

Backbone	Max Context ( $L_{max}$ )	Graph Const. (ms)	Total Runtime (ms)	Overhead (%)
BERT	512	$0.043\pm 0.002$	$5.36\pm 0.06$	$0.8\%$
TinyLlama-1.1B	2048	$0.672\pm 0.001$	$143.15\pm 0.05$	$0.5\%$
SmolLM2	8192	$4.46\pm 0.05$	$772.30\pm 0.11$	$0.6\%$
Llama-3B	8192	$15.05\pm 0.37$	$2041.77\pm 0.19$	$0.7\%$
Mistral-7B	32768	$303.47\pm 1.02$	$23460.29\pm 4.77$	$1.3\%$

Appendix

Table 10: Scalability Stress Test. Graph construction time vs. total inference time (per sample) at maximum context lengths (

L_{max}

). All times are in milliseconds (ms). Even at

L=32\textsc{K}

, the graph construction overhead is

\approx 1.3\%

of the total runtime.

Method	CoLA (MCC)	STSB (Spearman)	RTE (ACC)
Mistral-7B (Decoder-only)
AdaPool (No GNN)	48.00	79.55	54.87
Glot (GCN)	52.65	79.74	57.04
Glot (GAT) (Ours)	54.30	80.51	59.21
Glot (GIN)	59.30	79.73	59.30
BERT (Encoder-only)

Appendix

Table 11: Performance comparison of different GNN backbones. We evaluate the effectiveness of different graph variants against the AdaPool baseline using both Mistral-7B and BERT embeddings.

Dataset	[EOS]	Mean	MaxPool	AdaPool	Glot
Classification (ACC $\uparrow$ )	0.4910	0.5178	0.3917	0.4907	0.4727
AmazonCounterfactualClassification	0.5767	0.6301	0.6052	0.6076	0.6210
AmazonPolarityClassification	0.7858	0.6912	0.6447	0.6364	0.6299
AmazonReviewsClassification	0.3865	0.3544	0.2842	0.3068	0.3048
Banking77Classification	0.3127	0.4482	0.2704	0.4425	0.4730
EmotionClassification	0.2662	0.2995	0.2142	0.2832	0.3016

Appendix

Table 12: Comprehensive MTEB benchmark results. Datasets are grouped by task category, with the category average reported in the shaded rows. The evaluation metric for each category is indicated in parentheses: Classification (Accuracy), Retrieval (NDCG@10), Clustering (V-Measure), STS (Spearman Correlation), Reranking (MAP), PairClassification (Average Precision), and Summarization (Spearman Correlation). † Evaluated in half precision due for datasets with

>100

K sentences, due to resource constraints. First , Second , Third .

Method (Encoder Architectures)	STS-B
BERT (Encoder-only)
PromptBERT	70.60
Glot (Ours)	83.85
Mistral-7B (Decoder-only)
PromptEOL	75.77
Pretended CoT	76.66

Appendix

Table 13: Performance comparison on the STS-B benchmark (Spearman correlation) between Glot and prompting-based methods across Encoder and Decoder architectures.

Model	Runtime (ms) $\downarrow$	Params (M) $\downarrow$	CoLA (MCC) $\uparrow$	STS-B (Spear.) $\uparrow$	RTE (ACC) $\uparrow$
SBERT (MPNet-v2) ( Song et al., 2020 )	52.99 $\pm$ 0.08	109.50	17.70	86.70	54.51
SimCSE (BERT-sup) ( Gao et al., 2021 )	47.21 $\pm$ 0.09	109.50	24.91	87.27	52.70
Glot (BERT)	13.40 $\pm$ 3.00	8.92	47.49	83.86	59.21

Appendix

Table 14: Performance comparison between Glot (frozen backbone) and fully fine-tuned contrastive baselines. Runtime is reported in milliseconds per batch. Best performance is in bold .

Method	Params (M) $\downarrow$	CoLA (MCC) $\uparrow$	STS-B (Spear.) $\uparrow$	RTE (ACC) $\uparrow$
MLP	9.2	51.33	74.12	57.76
Glot (Ours)	8.9	54.30	80.51	59.21

Appendix

Table 15: Ablation study comparing Glot against a parameter-matched MLP baseline using the Mistral-7B backbone. Best results are in bold .

Method	# Params	GPU Mem (GB) $\downarrow$	Runtime (s) $\downarrow$	MCC $\uparrow$	Spear. $\uparrow$	ACC $\uparrow$
[EOS]	8.2K	32.58	3.227 $\pm$ 0.006	38.63	72.36	50.90
Mean	8.2K	32.58	3.244 $\pm$ 0.003	38.61	77.96	53.07
Max	8.2K	32.58	3.259 $\pm$ 0.009	10.78	70.72	53.07
AdaPool	2.1M	32.58	3.249 $\pm$ 0.011	48.00	79.55	54.87
Glot (Ours)	8.92M	32.64	3.252 $\pm$ 0.007	53.29	80.51	59.21

Appendix

Table 16: Inference-time cost and performance benchmark using the Mistral-7B backbone. Runtime is measured in seconds per batch. Best performance is in bold .

Category	Method	Trainable Params	GPU Mem (GB)	CoLA (MCC)	STS-B (Spear)	RTE (Acc)
Low Cost	Mean Pooling	0	$<0.1$	38.61	77.96	53.07
	Max Pooling	0	$<0.1$	10.78	70.72	53.07
	AdaPool	2.1 M	$<0.1$	48.00	79.55	54.87
	MLP Baseline	9.2 M	0.42	51.33	74.12	57.76
High Cost	LoRA ( $r=64$ )	167.8 M	33.5	48.23	54.54	53.43
	Full Fine-Tuning	7,110 M	$>40.0$	49.63	55.68	55.23

Appendix

Table 17: Unified comparison of training efficiency and performance using the Mistral-7B backbone. We categorize methods by computational cost. Best results are in bold .

Panel of Figure 5 — Figure 5: Token Contribution Analysis with frozen BERT. Visualization of learned token weights ( $\pi$ ) on 2 examples. The orange highlights on the X-axis indicate the top-3 scoring tokens identified by Glot . While Mean Pooling (green) is uniform and AdaPool (grey) tends to over-index on high-frequency functional words, Glot (blue) consistently up-weights the semantic anchors essential for determining sentence equivalence.

Sentence Embeddings

Sentence embeddings map sentences to vectors for semantic similarity, retrieval, clustering, and downstream language tasks.

See all 35 papers

When Graph Tokens Sink: A Mechanistic Analysis of Graph Language Models
Jun 2, 2026Ding Zhang, Runtao Zhou, Wenqing Zheng +3LLM Interpretability Mechanistic Interpretability
University of Virginia · Capital One
Is One Token All It Takes? Graph Pooling Tokens for LLM-based GraphQA
Apr 1, 2026Ankit Grover, Lodovico Giaretta, Rémi Bourgerie +1LLM Reasoning with Graphs Graph Pooling
KTH Royal Institute of Technology, Stockholm, Sweden · RISE Research Institutes of Sweden, Stockholm, Sweden
Toward Graph-Tokenizing Large Language Models with Reconstructive Graph Instruction Tuning
Mar 2, 2026Zhongjian Zhang, Xiao Wang, Mengmei Zhang +2Graph Reconstruction LLMs for Graph Learning
Beijing University of Posts and Telecommunications, Beijing, China · Beihang University, Beijing, China · China Telecom Bestpay, Beijing, China

Jun 2, 2026J/K move · Enter open · S save

When Graph Tokens Sink: A Mechanistic Analysis of Graph Language Models

Ding Zhang, Runtao Zhou, Wenqing Zheng, Rizal Fathony+2

University of Virginia · Capital One

LLM Interpretability Mechanistic Interpretability Language Modeling Attention Sinks LLMs for Graph Learning

Open paper AskPDF arXiv

Graph Language Models (GLMs) have become a promising direction for adapting Large Language Models (LLMs) to graph learning tasks. By transforming graph topology and node information into graph tokens, GLMs allow LLMs to jointly process structured graph inputs and textual instructions. Yet, it remains unclear how LLMs internally interpret these graph tokens and whether graph tokens act as meaningful carriers of graph structure. In this work, we analyze how LLMs process graph information through graph-token behavior in representative GLM architectures. Findings. We find that the internal saliency of graph tokens in GLMs is not equivalent to graph information utilization. Graph sink tokens consistently emerge as activation-level outliers: they can be identified by massive activation values along a small set of hidden-state dimensions and are biased toward early graph-token positions. However, this activation-level saliency does not imply that these tokens are the main carriers of graph information. Unlike classical attention sinks in language and vision-language models, graph sink tokens do not necessarily attract the largest attention weights from query tokens. Through pruning, repositioning, and swapping interventions, we show that graph sink tokens are not the most important semantic or structural tokens for downstream prediction. Implications. Together, these results suggest that after current GLMs map graph structure into the LLM token space, the resulting graph-token representations do not naturally form a fully usable topology-aware internal representation; instead, they exhibit a decoupling between activation-level saliency and graph-semantic utility. This decoupling points to limitations in existing graph-token construction, placement, and alignment mechanisms.

Abstract

Sentence Embeddings

Explore similar work

When Graph Tokens Sink: A Mechanistic Analysis of Graph Language Models

Is One Token All It Takes? Graph Pooling Tokens for LLM-based GraphQA

Toward Graph-Tokenizing Large Language Models with Reconstructive Graph Instruction Tuning