Towards Concept-Aware Large Language Models (2023)

Abstract: Concepts play a pivotal role in various human cognitive functions, including learning, reasoning and communication. However, there is very little work on endowing machines with the ability to form and reason with concepts. In particular, state-of-the-art large language models (LLMs) work at the level of tokens, not concepts. In this work, we analyze how well contemporary LLMs capture human concepts and their structure. We then discuss ways to develop concept-aware LLMs, taking place at different stages of the pipeline. We sketch a method for pretraining LLMs using concepts, and also explore the simpler approach that uses the output of existing LLMs. Despite its simplicity, our proof-of-concept is shown to better match human intuition, as well as improve the robustness of predictions. These preliminary results underscore the promise of concept-aware LLMs.

## introduction concepts are the glue that holds our mental model of the world together. it is hard to see how any intelligent agent could do without them. while there is no agreed-upon definition of concepts, one can think of them as the mental representations that enable us to identify objects and events as belonging to certain categories, communicate about them, and comprehend new situations in terms of previous ones: when we encounter a new situation (e.g., restaurant), we draw inferences about it using concepts we have already formed ("menu", "waiter"). concepts can be concrete ("soup") or abstract ("tasty"). they can also be complex, e.g., "good winter beach destinations". while there is a lively debate on their exact nature, researchers agree concepts play a pivotal role in various cognitive skills such as reasoning, categorization, learning, planning, and decision-making @xcite . thus, they are of interest to ai researchers wishing to endow machines with such abilities. the representation of concepts has been studied in nlp, ml, and knowledge representation @xcite @xcite , where they often view concepts as fixed, shallow structures representing some set of entities. for example, in the work on @xcite concepts are flat sets of context-independent entities. however, recent studies suggest concepts are more flexible and dynamic @xcite ; unfortunately, ai still struggles with accounting for the creative, contextsensitive manner in which people employ concepts. in this work we focus on adding concepts to large language models (llms). recently, llms @xcite @xcite @xcite gained immense popularity, achieving sota results across the board. however, they all work at the level of tokens, not concepts. this is problematic for even the most fundamental llm task -perplexity-based text completion. ranking by string probability is distorted by surface form competition: different tokens compete with each other, even if they represent the same concept ("mother" and "mom") @xcite . in other words, the probability mass of a concept is distributed across many different tokens, distorting the ranking. the problem runs deeper than mere synonyms. for example, consider the sentence "i can't get home for the holidays because of the @xcite ." the completions "snow", "blizzard", "weather", and "slippery roads" are not synonyms per se, but they correspond to the same concept -bad weather leading to hazardous driving conditions -and we believe that an llm should treat them as such. we stress that concepts are context-dependent; for example, while "snow" and "blizzard" are similar in the context of the sentence above, they are very different for the sentence "i love eating [mask] cones." 1 thus, we cannot rely on knowledge bases (such wordnet @xcite ) for generating static, context-free concepts to be used for training llms. we take the first step towards concept-aware llms, exploring the following questions: rq1: how well do llms capture concepts? rq2: how well do llms match human organization of concepts? rq3: how can we enhance an llm in terms of concepts, with or without retraining? we first show that contemporary llms capture human concepts to some extent, but they are still far from humans (rq1). we then find that llms violate many of the principles of human concept organization, exhibiting some inconsistency (rq2). lastly, we explore rq3 from two different angles: first, we sketch a method to pretrain conceptaware llms. next, we implement a proof-ofconcept model-agnostic method to shift any off-theshelf pretrained llm from token-to concept-level with no further training. our method improves both the ranking and robustness of the underlying llm. while we present here merely a promising proofof-concept, our underlying objective of endowing llms with concepts holds tremendous promise for the next-generation of llms. we hope our work will spur further research, paving new roads into this exciting new territory. ## rq1: how well do llms capture concepts? in this section, we explore to what extent llms grasp concepts. while concepts can be abstract and complex, here we focus on concrete and simple ones, since these are the basic building blocks of human concepts @xcite . ## rq2: how well do llms match human organization of concepts? in this section, we explore concept organization. we focus on three agreed-upon organization princi-ples of the human typeof hierarchy of concepts: ## related work there is little work on endowing llms with a notion of concepts. the work that is closest to ours in spirit is sensebert @xcite , which shifts bert from the form-to the sense-level by injecting wordnet token information, improving word sense disambiguation. unlike sensebert, we suggest learning the categories in a contextdependent manner (by using the categories that best assist in predicting the missing text), rather than their static, context-independent method. works on knowledge base embeddings learn representations from knowledge graphs @xcite , but they are also limited to simple concepts and kb relations, with no external context. aharoni and goldberg (2020) showed that llms capture the domain a sentence belongs to. while domains can also be thought of as concepts, we are after a different (and dynamic) granularity. conceptx is an interpretability framework to analyze how latent concepts are encoded in representations learned within pretrained llms @xcite . they found that roughly 50% of the encoded concepts adhere to their suite of humandefined linguistic concepts, which aligns with our results from rq1. however, their method fails to address complex, multi-faceted concepts. @xcite specifically explored the extent to which bert captures hypernymy, exploring different prompts and evaluation metrics. we note that their dataset was somewhat biased towards australian and new zealand culture, and also contained multiple errors. ## conclusions & future work in this work we explored the possibility of conceptaware llms, inspired by the importance of concepts in human cognition. today's llms all work at the level of tokens, not concepts; this is problematic, as different tokens from the same underlying concept compete for the probability mass. instead, we envision llms that first zero in on a likely concept, and only then pick the best surface form among the candidates. we started by analyzing to what extent contemporary pretrained llms capture concepts and match the human organization of concepts. we showed that llms do capture concepts to some (far-from-perfect) extent, while sometimes violating the human concept-organization principles of asymmetry, transitivity, and property inheritance. next, we explored directions for augmenting llms with concepts from two different angles. we first sketched our vision for training conceptaware llms. next, we presented a model-agnostic method to shift any off-the-shelf pretrained llm from the token-to the concept-level, without finetuning or adding any external information. we showed that our concept-bert (built on top of bert) outputs a ranked list of concepts which are relatively coherent and distinct, better match human intuition, and are more robust compared to the underlying model. the notion of concepts has been extensively studied in psychology, and we believe many of the questions posed there could inspire directions for future work. for example, prototype theory states there is a graded degree of belonging to a concept @xcite , which we have not taken into account here. other works tackled the problem of complex concepts and the composition of concepts. for example, @xcite showed that concepts can be combined ("tall" + "man") in a context-dependent way (e.g., the height for a man to be considered tall is not the same as for a building). for this direction, commonsense sources such as the distribution-overquantities dataset @xcite while this is only preliminary work, we believe that concept-aware llms hold immense promise for the next-generation of llms, and could benefit many downstream tasks around learning, planning, and reasoning. one place where we believe this approach will be particularly useful is whenever there is disambiguity (speech recognition, word sense, spell check, etc.). consider rare words that need to be identified by automatic speech recognition: by working at concept-level, these words could be clustered together with completions that are likely given the rest of the sentence but do not sound similar to the audio, thus increasing the acoustic model's certainty score. we hope to draw the community's attention to this exciting new research area of concept-aware llms.