AEA: Adaptive Expert Allocation Improves Sentence Embeddings from Mixture-of-Experts LLM
Paper · 2026 · Shufan Yang, Zifeng Cheng, Zhiwei Jiang et al.
Summary
The paper addresses training-free sentence embeddings taken directly from Mixture-of-Experts (MoE) LLMs. Earlier work, such as MoEE, keeps the same Top-K number of experts for every layer and every token. The authors show that this uniform budget wastes capacity. In DeepSeekMoE-16B and OLMoE-1B-7B, expert outputs are much more alike in the lower and upper layers than in the middle ones, and attention is concentrated on only a few tokens.
Their AEA framework redistributes the existing expert budget in two ways, without adding experts:
- Across layers: each layer gets a homogeneity score, computed once on a calibration set from pairwise similarity of expert outputs. More homogeneous layers receive fewer experts.
- Across tokens: tokens that receive weaker attention get fewer experts. Importance is measured by the highest attention a token receives from any token and any head.
AEA is plug-and-play, works with prompt-based methods such as PromptEOL, and adds no runtime cost. On the STS benchmarks it improves embedding quality across several MoE models, so stronger embeddings come with no fine-tuning or extra data.
Key points
- Models tested include OLMoE-1B-7B, DeepSeekMoE-16B and Qwen1.5-MoE-A2.7B. Evaluation uses STS12–STS16, SICK-R and STS-B (Spearman correlation).
- About 90% of attention weights fall below one-tenth of their row's maximum, which the authors take as evidence of strong token-level sparsity.
- On OLMoE-1B-7B, the following were compared against the hidden-state baseline (average 72.88):
- Token-wise allocation using max pooling over both tokens and heads reached 75.44.
- Token-wise allocation alone reached 74.42.
- Reversing the token-wise rule (more experts to low-attention tokens) dropped performance to 49.87.
- With three different prompts on STS-B (OLMoE-1B-7B), AEA gains between 2.59 and 5.45 points. The penultimate layer gives the best embeddings for all three models.
- The authors list two limitations: AEA needs internal access to expert outputs and attention maps, and it was evaluated only on English benchmarks.
Fields, topics and concepts
📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 18,614 worksArtificial Intelligence · 17,965 worksArts and Humanities · 17,450 worksLanguage and Linguistics · 17,296 works
🧭 Topics 1Research subjects and languages the work is about.Language models & pretraining · 3,894 works
Cite
Cite as Shufan Yang, Zifeng Cheng, Zhiwei Jiang, Qingfeng Qi, Yafeng Yin, Cong Wang, Ao Zhou and Qing Gu. 2026. AEA: Adaptive Expert Allocation Improves Sentence Embeddings from Mixture-of-Experts LLM. Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers).
@inproceedings{yang2026aea,
title = {{AEA: Adaptive Expert Allocation Improves Sentence Embeddings from Mixture-of-Experts LLM}},
author = {Shufan Yang, Zifeng Cheng, Zhiwei Jiang, Qingfeng Qi, Yafeng Yin, Cong Wang, Ao Zhou and Qing Gu},
booktitle = {Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)},
year = {2026},
doi = {10.18653/v1/2026.acl-long.1331},
url = {https://doi.org/10.18653/v1/2026.acl-long.1331}
}References
- Eneko Agirre, Carmen Banea, Claire Cardie, Daniel M. Cer, Mona T. Diab, Aitor Gonzalez-Agirre, Weiwei Guo, Iñigo Lopez-Gazpio, Montse Maritxalar, Rada Mihalcea, German Rigau, Larraitz Uria, and Janyce Wiebe. 2015. Semeval-2015 task 2: Semantic textual similarity, english, spanish and pilot on interp…
- Eneko Agirre, Carmen Banea, Claire Cardie, Daniel M. Cer, Mona T. Diab, Aitor Gonzalez-Agirre, Weiwei Guo, Rada Mihalcea, German Rigau, and Janyce Wiebe. 2014. Semeval-2014 task 10: Multilingual semantic textual similarity. In Proceedings of the 8th International Workshop on Semantic Evaluation, Sem…
- Eneko Agirre, Carmen Banea, Daniel M. Cer, Mona T. Diab, Aitor Gonzalez-Agirre, Rada Mihalcea, German Rigau, and Janyce Wiebe. 2016. Semeval-2016 task 1: Semantic textual similarity, monolingual and cross-lingual evaluation. In Proceedings of the 10th International Workshop on Semantic Evaluation, S…
- Eneko Agirre, Daniel M. Cer, Mona T. Diab, and Aitor Gonzalez-Agirre. 2012. Semeval-2012 task 6: A pilot on semantic textual similarity. In Proceedings of the 6th International Workshop on Semantic Evaluation, SemEval@NAACL-HLT 2012, pages 385–393.
- Eneko Agirre, Daniel M. Cer, Mona T. Diab, Aitor Gonzalez-Agirre, and Weiwei Guo. 2013. *sem 2013 shared task: Semantic textual similarity. In Proceedings of the Second Joint Conference on Lexical and Computational Semantics, *SEM 2013, pages 32–43.
- Parishad BehnamGhader, Vaibhav Adlakha, Marius Mosbach, Dzmitry Bahdanau, Nicolas Chapados, and Siva Reddy. 2024. Llm2vec: Large language models are secretly powerful text encoders. arXiv preprint arXiv:2404.05961.
- Daniel M. Cer, Mona T. Diab, Eneko Agirre, Iñigo Lopez-Gazpio, and Lucia Specia. 2017. Semeval-2017 task 1: Semantic textual similarity - multilingual and cross-lingual focused evaluation. CoRR, abs/1708.00055.
- Sachin Chanchani and Ruihong Huang. 2023. Composition-contrastive learning for sentence embeddings. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 15836–15848. find
- Zifeng Cheng, Zhonghui Wang, Yuchen Fu, Zhiwei Jiang, Yafeng Yin, Cong Wang, and Qing Gu. 2025. Contrastive prompting enhances sentence embeddings in llms through inference-time steering. arXiv preprint arXiv:2505.12831.
- Damai Dai, Chengqi Deng, Chenggang Zhao, RX Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Y Wu, et al. 2024. Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models. arXiv preprint arXiv:2401.06066. find
- Yuchen Fu, Zifeng Cheng, Zhiwei Jiang, Zhonghui Wang, Yafeng Yin, Zhengliang Li, and Qing Gu. 2025. Token prepending: A training-free approach for eliciting better sentence embeddings from llms. In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Lon…
- Tianyu Gao, Xingcheng Yao, and Danqi Chen. 2021. Simcse: Simple contrastive learning of sentence embeddings. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 6894–6910.
- Ting Jiang, Shaohan Huang, Zhongzhi Luan, Deqing Wang, and Fuzhen Zhuang. 2024. Scaling sentence embeddings with large language models. In Findings of the Association for Computational Linguistics: EMNLP 2024, pages 3182–3196.
- Ting Jiang, Jian Jiao, Shaohan Huang, Zihan Zhang, Deqing Wang, Fuzhen Zhuang, Furu Wei, Haizhen Huang, Denvy Deng, and Qi Zhang. 2022. Promptbert: Improving bert sentence embeddings with prompts. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 8826–8… find
- Chankyu Lee, Rajarshi Roy, Mengyao Xu, Jonathan Raiman, Mohammad Shoeybi, Bryan Catanzaro, and Wei Ping. 2024. Nv-embed: Improved techniques for training llms as generalist embedding models. arXiv preprint arXiv:2405.17428.
- Yibin Lei, Di Wu, Tianyi Zhou, Tao Shen, Yu Cao, Chongyang Tao, and Andrew Yates. 2024. Meta-task prompting elicits embeddings from large language models. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), ACL 2024, pages 10141–10157. find
- Junxian Li, Beining Xu, Simin Chen, Jiatong Li, Jingdi Lei, Haodong Zhao, and Di Zhang. 2025. Iag: Input-aware backdoor attack on vlm-based visual grounding. arXiv preprint arXiv:2508.09456. — [arXiv:2508.09456](https://arxiv.org/abs/2508.09456) find
- Pengyu Li, Lingling Zhang, Zhitao Gao, Yanrui Wu, Yuxuan Dong, Huan Liu, Bifan Wei, and Jun Liu. 2026. AGT AO : Robust and stabilized llm unlearning via adversarial gating training with adaptive orthogonality. arXiv preprint arXiv:2602.01703. — [arXiv:2602.01703](https://arxiv.org/abs/2602.01703) find
- Xianming Li and Jing Li. 2024. Bellm: Backward dependency enhanced large language model for sentence embeddings. In Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers), pages 792–804. find
- Ziyue Li and Tianyi Zhou. 2025. Your mixture-of-experts LLM is secretly an embedding model for free. In The Thirteenth International Conference on Learning Representations, ICLR 2025. find
- Tian Yu Liu, Matthew Trager, Alessandro Achille, Pramuditha Perera, Luca Zancato, and Stefano Soatto. 2024a. Meaning representations from trajectories in autoregressive models. In The Twelfth International Conference on Learning Representations, ICLR 2024. find
- Zhu Liu, Cunliang Kong, Ying Liu, and Maosong Sun. 2024b. Fantastic semantics and where to find them: Investigating which layers of generative llms reflect lexical semantics. In Findings of the Association for Computational Linguistics ACL 2024, pages 14551– 14558. find
- Marco Marelli, Stefano Menini, Marco Baroni, Luisa Bentivogli, Raffaella Bernardi, and Roberto Zamparelli. 2014. A SICK cure for the evaluation of compositional distributional semantic models. In Proceedings of the Ninth International Conference on Language Resources and Evaluation, LREC 2014, pages…
- Niklas Muennighoff, Luca Soldaini, Dirk Groeneveld, Kyle Lo, Jacob Morrison, Sewon Min, Weijia Shi, Pete Walsh, Oyvind Tafjord, Nathan Lambert, et al. 2024a. Olmoe: Open mixture-of-experts language models. arXiv preprint arXiv:2409.02060. find
- Niklas Muennighoff, Hongjin Su, Liang Wang, Nan Yang, Furu Wei, Tao Yu, Amanpreet Singh, and Douwe Kiela. 2024b. Generative representational instruction tuning. arXiv preprint arXiv:2402.09906.
- Niklas Muennighoff, Nouamane Tazi, Loïc Magne, and Nils Reimers. 2022. Mteb: Massive text embedding benchmark. arXiv preprint arXiv:2210.07316.
- Jianmo Ni, Gustavo Hernandez Abrego, Noah Constant, Ji Ma, Keith Hall, Daniel Cer, and Yinfei Yang. 2022. Sentence-t5: Scalable sentence encoders from pre-trained text-to-text models. In Findings of the Association for Computational Linguistics: ACL 2022, pages 1864–1874.
- Jacob Mitchell Springer, Suhas Kotha, Daniel Fried, Graham Neubig, and Aditi Raghunathan. 2025. Repetition improves language model embeddings. In The Thirteenth International Conference on Learning Representations, ICLR 2025.
- Hongjin Su, Weijia Shi, Jungo Kasai, Yizhong Wang, Yushi Hu, Mari Ostendorf, Wen-tau Yih, Noah A Smith, Luke Zettlemoyer, and Tao Yu. 2023. One embedder, any task: Instruction-finetuned text embeddings. In Findings of the Association for Computational Linguistics: ACL 2023, pages 1102–1121.
- Qwen Team. 2024. Qwen1.5-moe: Matching 7b model performance with 1/3 activated parameters". find
- Liang Wang, Nan Yang, Xiaolong Huang, Linjun Yang, Rangan Majumder, and Furu Wei. 2024. Improving text embeddings with large language models. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), ACL 2024, pages 11897–11916.
- Shufan Yang, Zifeng Cheng, Zhiwei Jiang, Yafeng Yin, Cong Wang, Shiping Ge, Yuchen Fu, and Qing Gu. 2026. Regionmarker: A region-triggered semantic watermarking framework for embedding-asa-service copyright protection. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 40, page… find
- Bowen Zhang, Kehua Chang, and Chunping Li. 2024. Simple techniques for enhancing sentence embeddings in generative language models. CoRR, abs/2404.03921. — [arXiv:2404.03921](https://arxiv.org/abs/2404.03921) find
- Hao Zhang, Mengsi Lyu, Zhuo Chen, Xingrun Xing, Yulong Ao, and Yonghua Lin. 2025. Pdtrim: Targeted pruning for prefill-decode disaggregation in inference. arXiv preprint arXiv:2509.04467.
Updated
This page in Markdown → https://md.research.elamur.ai/free/en/papers/yang-2026-aea-adaptive-expert-allocation-improves/