Towards Improved Sentence Representations using Token Graphs
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
| Method | # Trainable Params | GPU Memory (GB) | Batch Runtime (ms) | MCC |
| Full FT EOS | 7.11B | 32.59 | 49.63 | |
| LoRA ( ) EOS | 167.8M | 33.50 | 48.23 | |
| Glot (ours) | 8.92M | 0.42 | 53.29 |
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 ( ) | {2, 4} |
| GNN Hidden Dimension | {64, 128, 256} |
| 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 |
| Threshold ( ) | STSB (Spear. ) | CoLA (MCC ) | RTE (ACC ) | Linear Probing (ACC ) |
| BERT (Encoder-only) | ||||
| 81.88 | 39.62 | 50.90 | 76.41 | |
| 82.12 | 39.82 | 50.90 | 76.51 | |
| 82.25 | 47.49 | 52.34 | 77.18 | |
| 83.86 | 43.16 | 59.21 | 77.56 | |
| 83.85 | 43.16 | 52.70 | 77.18 | |
| Method | Trainable | GPU Mem. | Runtime | CoLA | STSB | RTE |
| Params | (GB) | (ms) | (MCC) | (Spear.) | (ACC) | |
| Mistral-7B (Decoder-only) | ||||||
| Full FT + [EOS] | 7.11B | 32.59 | 49.63 | 55.68 | 55.23 | |
| LoRA ( ) | 167.8M | 33.50 | 48.23 | 54.54 | 53.43 | |
| Glot (Ours) | 8.92M | 0.42 | 13.4 3.0 | 53.29 | 80.51 | 59.21 |
| BERT (Encoder-only) | ||||||
| Backbone | Max Context ( ) | Graph Const. (ms) | Total Runtime (ms) | Overhead (%) |
| BERT | 512 | |||
| TinyLlama-1.1B | 2048 | |||
| SmolLM2 | 8192 | |||
| Llama-3B | 8192 | |||
| Mistral-7B | 32768 |
| 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) | |||
| Dataset | [EOS] | Mean | MaxPool | AdaPool | Glot |
| Classification (ACC ) | 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 |
| 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 |
| Model | Runtime (ms) | Params (M) | CoLA (MCC) | STS-B (Spear.) | RTE (ACC) |
| SBERT (MPNet-v2) ( Song et al., 2020 ) | 52.99 0.08 | 109.50 | 17.70 | 86.70 | 54.51 |
| SimCSE (BERT-sup) ( Gao et al., 2021 ) | 47.21 0.09 | 109.50 | 24.91 | 87.27 | 52.70 |
| Glot (BERT) | 13.40 3.00 | 8.92 | 47.49 | 83.86 | 59.21 |
| Method | Params (M) | CoLA (MCC) | STS-B (Spear.) | RTE (ACC) |
| MLP | 9.2 | 51.33 | 74.12 | 57.76 |
| Glot (Ours) | 8.9 | 54.30 | 80.51 | 59.21 |
| Method | # Params | GPU Mem (GB) | Runtime (s) | MCC | Spear. | ACC |
| [EOS] | 8.2K | 32.58 | 3.227 0.006 | 38.63 | 72.36 | 50.90 |
| Mean | 8.2K | 32.58 | 3.244 0.003 | 38.61 | 77.96 | 53.07 |
| Max | 8.2K | 32.58 | 3.259 0.009 | 10.78 | 70.72 | 53.07 |
| AdaPool | 2.1M | 32.58 | 3.249 0.011 | 48.00 | 79.55 | 54.87 |
| Glot (Ours) | 8.92M | 32.64 | 3.252 0.007 | 53.29 | 80.51 | 59.21 |
| Category | Method | Trainable Params | GPU Mem (GB) | CoLA (MCC) | STS-B (Spear) | RTE (Acc) |
| Low Cost | Mean Pooling | 0 | 38.61 | 77.96 | 53.07 | |
| Max Pooling | 0 | 10.78 | 70.72 | 53.07 | ||
| AdaPool | 2.1 M | 48.00 | 79.55 | 54.87 | ||
| MLP Baseline | 9.2 M | 0.42 | 51.33 | 74.12 | 57.76 | |
| High Cost | LoRA ( ) | 167.8 M | 33.5 | 48.23 | 54.54 | 53.43 |
| Full Fine-Tuning | 7,110 M | 49.63 | 55.68 | 55.23 |
Sentence Embeddings
Sentence embeddings map sentences to vectors for semantic similarity, retrieval, clustering, and downstream language tasks.
See all 35 papers