Multi-Sense Embeddings for Language Models and Knowledge Distillation (2025)

Abstract: Transformer-based large language models (LLMs) rely on contextual embeddings which generate different (continuous) representations for the same token depending on its surrounding context. Nonetheless, words and tokens typically have a limited number of senses (or meanings). We propose multi-sense embeddings as a drop-in replacement for each token in order to capture the range of their uses in a language. To construct a sense embedding dictionary, we apply a clustering algorithm to embeddings generated by an LLM and consider the cluster centers as representative sense embeddings. In addition, we propose a novel knowledge distillation method that leverages the sense dictionary to learn a smaller student model that mimics the senses from the much larger base LLM model, offering significant space and inference time savings, while maintaining competitive performance. Via thorough experiments on various benchmarks, we showcase the effectiveness of our sense embeddings and knowledge distillation approach.

## introduction in recent years, transformer-based large language models (llms) have revolutionized natural language processing by providing powerful capabilities for a wide range of applications @xcite . these language models rely on continuous contextual embeddings, which allow for an infinite number of representations for each token. while this approach has proven effective, it contrasts with the way humans perceive language, where each word often carries only a limited number of distinct senses. for example, consider the visualization of the different contextual embeddings of the token "novel" in figure 1 . we can observe a good amount of variation in the embeddings, but there are only two main senses: i) as a noun indicating a fictitious prose narrative, and ii) as an adjective denoting new or original, with some interesting subclasses (e.g., graphic/visual novels). unlike llms, which learn languages by processing vast corpora with token-based vocabularies, humans use a different methodology: we first learn the meaning of a new word by remembering a limited number of senses and storing them in our brains; when encountering new text, we determine the appropriate sense of each word based on the context and then integrate these meanings @xcite . this observation raises an intriguing question: can we effectively replace the virtually unlimited context-based embeddings generated by large language models for the same token, with a finite set of embeddings more in-tune with human language understanding? inspired by this notion, we explore the potential of multi-sense discrete embeddings, or just sense embeddings for short, as a means to compress and store continuous embeddings as multiple distinct vectors for each token. more specifically, our approach collects the embeddings produced from an llm encoder and feeds them to a clustering algorithm, e.g., k-means clustering @xcite , which returns the k cluster centers for each token in the vocabulary. these cluster centers serve as the sense embeddings, and are stored in a sense dictionary to provide compact representations that preserve essential semantic information, and are readily used in subsequent layers of llms. furthermore, we propose a novel sense-based knowledge distillation technique. knowledge distillation refers to learning a smaller student model that is trained to replicate the behavior of a larger teacher model, thereby facilitating computational time and space savings during inference time. in our sensebased distillation, we use the sense dictionary to replace the last hidden layer in a large (teacher) encoder to yield a much smaller (student) encoder model. we conduct empirical experiments on both encoder-based and decoder-based models. to show figure 1 : contextual embeddings of the token novel generated by the deberta-v3-large llm model @xcite (scatter plot is shown in 2d via t-sne (van der @xcite on embedding vectors with dimension @xmath0, which have a total variance of 220.8). two main senses can be observed: i) fictitious prose narrative (center group), and ii) new and unusual (top right group). we also observe an auxiliary sense: iii) visual/graphic novel (leftmost group), which can be considered a subclass of sense i). the effectiveness of our multi-sense embeddings we evaluate them as a drop-in replacement in lieu the llms' continuous embeddings, and we also study their performance on the word similarity task. in addition, we evaluate our knowledge distillation framework using the glue @xcite and mteb @xcite benchmark tasks and demonstrate that significant space savings can be achieved without incurring a major decrease in the observed accuracy. our main contributions can be summarized as follows: the proposed knowledge distillation model outperforms the finetuned deberta-v3-xsmall (he et al., 2021) on glue benchmarks, and achieves 92% of the average accuracy while utilizing only 19% of the gpu memory compared to llama-3-8b-instruct (touvron et al., 2023) on mteb classification tasks. ## background and related work language models: there are mainly two types of language models: discrete embedding-based language models and context embedding-based language models @xcite . discrete embedding-based language models, such as word2vec @xcite , drg2vec @xcite , dict2vec @xcite , and hg2vec @xcite utilize sources like wikipedia or dictionaries as the input corpus. these models slide a context window along the input text and maximize the similarity of words within each context window. after training, each word is assigned a discrete embedding, allowing users to leverage them directly during inference. but they typically assign one embedding vector per word (or token), and therefore fail to capture context-specific meanings for polysemous words, i.e., those that have different meanings in different contexts (e.g., the word 'bank' can refer to a financial institution, a rising ground, an airplane's incline and so on). contextual language models @xcite can be grouped into three main types: pure encoder-based models, such as bert @xcite and deberta @xcite ; pure decoder-based models, such as gpt2 @xcite and llama @xcite ; and encoder-decoder-based models, such as t5 @xcite . all of them are based on the attention mechanism @xcite . encoderbased models transform the input text into a contextualized representation by applying self-attention mechanisms to capture dependencies between tokens across the entire input. decoder-based models auto-regressively leverage the self-attention mechanism to generate new tokens. encoder-decoderbased models integrate an encoding mechanism to capture and represent the contextual semantics of the input and then sequentially generate the output sequence through decoders. nonetheless, even though llms excel at understanding the meaning of the entire corpus and demonstrate superior performance on various tasks, they suffer from high space and computational costs for both training and inference. our proposed approach aims to leverage the strengths of both discrete embedding-based language models and context embedding-based language models. to achieve this, we gather the output embeddings from llms and cluster them into groups to discern the different senses of a token. this approach preserves the semantics, at the same time significantly reduces the time and space requirements during inference. sense embeddings: in contrast to the traditional concept of embeddings, where each word is represented by a single vector, sense embeddings associate multiple vectors per word, where each one of the vectors aims to capture a different meaning @xcite . however, different models utilize different definitions of sense embeddings. mssg @xcite extends the traditional skip-gram model to effectively capture multiple senses of a word by clustering the context of word occurrences and associating each cluster with a distinct sense embedding. fasttext @xcite introduces multi-sense embedding per word with gaussian components. @xcite do word sense induction but use discrete word representations, i.e., as a probability distribution over words, instead of using continu-ous embeddings, and therefore their sense vectors are not suitable for use in llm models. autoextend @xcite focues on earning embeddings for synsets (a group of synonyms) and lexemes (associating spellings with particular meanings), treating both words and synsets as a sum of their lexemes, whereas our focus is on individual token-based sense embeddings, rather than group-based embeddings. @xcite use ego-network graph-based clustering of word vectors to create sense vectors, which is complimentary to our approach. lmms @xcite generates embeddings per sense key defined by annotated resource, such as wordnet @xcite . as a result, lmms can achieve more precise representations for each sense key but lacks support for most datasets consisting of plain text. in contrast, our work develops a general solution to generate multiple embeddings per token without requiring annotation. vector quantization (vq): vq is a data compression technique where a large set of vectors is represented by a smaller set of reference vectors, called a codebook (with each vector called a code). @xcite significantly reduces memory and computation cost for searching by introducing product quantization. extending this, vq-vae @xcite incorporates vq into deep generative models through variational autoencoders, showcasing its potential for learning discrete latent representations. transformer-vq (lingle, 2023) proposes a linear time attention model for decoders with using vector-quantized keys and a novel caching mechanism within the attention process. our work is closely related to vq, but instead of a global codebook, our sense dictionary can be considered as using a fine-grained token-based codebook to elicit and capture a finite set of meanings or senses for each token. incidentally, our work is complementary and compatible with the precision-based quantization approaches that use low-bit representation to optimize learning cost @xcite . knowledge distillation: knowledge distillation is a key model compression technique that trains a smaller student model to replicate the outputs of a larger teacher model. the primary advantage of this approach is its ability to significantly reduce model size without substantially sacrificing performance. various methods have been developed to enhance this training process. the first approach focuses on replicating the distribution of the output layer, such as distilbert @xcite and lightpaff @xcite . the second category leverages the benefits of mimicking the hidden states. tinybert @xcite and pkd @xcite leverage the information in the teacher's hidden layers, encouraging the student model to learn through a multi-step distillation process. minilmv2 @xcite introduces an additional strategy of mimicking the attention layers. recently, with the emergence of extremely large language models containing hundreds of billions of parameters, methods like alpaca @xcite and vicuna @xcite have been developed to mimic the outputs generated by a larger llm via instruction or conversation-based fine-tuning. we propose a novel approach to knowledge distillation, where the much smaller student model learns to choose the correct sense for each token based on the guidance from a much larger teacher model. this can significantly cut down on the space cost with little loss in performance. ## multi-sense embeddings we first tackle the issue of whether we can extract meaningful senses for the tokens and use a discrete set of sense embeddings per token, rather than having an unlimited number of representations for the same token. we next tackle the issue of using our multi-sense embeddings for a novel sense-based knowledge distillation approach whereby a smaller student model learns to mimic the sense output of the larger teacher (base) llm, thereby saving both space and time during inference. as a motivating example, we consider the deberta-v3-large llm @xcite , which is an encoder-based model with hidden dimensionality @xmath1, and it improves over both bert @xcite and roberta @xcite ) models using disentangled attention, enhanced mask decoder, and gradient disentangled embedding sharing. we randomly selected 5,000 occurrences of the token novel, from the english wikipedia dump @xcite ; it has two primary meanings: new and unusual and a fictitious prose narrative. figure 1 shows the distribution of contextual embeddings for novel from the last hidden layer of deberta-v3-large @xcite , where we use t-sne (van der @xcite to visualize the embeddings in 2d. we see a considerable spread of these embeddings in the hidden latent space. nevertheless, tokens with similar meanings naturally form clusters, rather than being randomly distributed. the larger cluster in the center corresponds to the meaning a fictitious prose narrative, while the smaller cluster on the top right represents new and unusual. we also observe subgroups depending on the context, e.g., the smaller cluster on the left represents phrases such as graphic novel and visual novel is a subgroup of the first meaning. we conclude that clustering can capture semantic information based on context or usage, and the distinct centers can represent different senses of each token. ## empirical evaluation for our experiments we employ up to 6 nvidia v100 gpus (32gb ram each) with 20-core ibm power 9 processors (512gb ram). experimental details on the various benchmarks and models are given in the appendix. our code is available at https://github.com/ qitong-wang/sensedictionary . ## conclusion and future work in this paper, we make two main contributions. the first is the notion of a sense dictionary that encapsulates multi-sense embeddings that capture a finite number of meanings, as opposed to the infinite number of continuous contextual embeddings in llm models. we show that our sense embeddings can serve as drop-in replacement for the embeddings in both encoder-and decoder-based llms, and retain the virtually the same performance. we also propose a novel sense-based knowledge distillation approach, where the significantly smaller student model learns the correct sense based on the teacher llm, resulting in a significant reduction in space and inference time, but with better performance compared to other baseline distilled models. there are still areas for improvement, such as developing an end-to-end approach to learn the sense embeddings. our sense embedding approach offers the intriguing possibility of learning from a small but diverse corpus that encapsulates most of the senses of each token, rather than relying on trillions of tokens as done by the current generation of llms; we plan to explore this in the future.