Disentangling Transformer Language Models as Superposed Topic Models (2023)

Abstract: Topic Modelling is an established research area where the quality of a given topic is measured using coherence metrics. Often, we infer topics from Neural Topic Models (NTM) by interpreting their decoder weights, consisting of top-activated words projected from individual neurons. Transformer-based Language Models (TLM) similarly consist of decoder weights. However, due to its hypothesised superposition properties, the final logits originating from the residual path are considered uninterpretable. Therefore, we posit that we can interpret TLM as superposed NTM by proposing a novel weight-based, model-agnostic and corpus-agnostic approach to search and disentangle decoder-only TLM, potentially mapping individual neurons to multiple coherent topics. Our results show that it is empirically feasible to disentangle coherent topics from GPT-2 models using the Wikipedia corpus. We validate this approach for GPT-2 models using Zero-Shot Topic Modelling. Finally, we extend the proposed approach to disentangle and analyse LLaMA models.

## introduction the term 'superposition' in machine learning typically refers to overlapping concepts. @xcite allude to the property of superposition in transformers when discussing the difficulty of interpreting its multilayer perceptrons' (mlp) neurons in transformer language models (tlm). @xcite also find neurons in transformers having polysemantic behaviour. we believe it is possible to interpret neurons in superposition, motivated by three key insights. our first insight lies in the similarity between tlm and neural topic models (ntm) (see section 2). many analyses of transformers examine some variation of neuron activations paths @xcite . again, ntm also uses paths (albeit usually across one layer) and interprets the highest activated logits from a given neuron as a topic. it is thus natural to treat the interpretation of tlm as a topic model with topics in superposition. @xcite show that polysemous words, i.e., multiple meanings, are in linear superposition within word embeddings. similarly, we posit to linearly disentangle the logit interpretation of neurons in superposition and examine the neurons in the residual stream (see section 4). we treat neurons from mlp as belonging to the higherdimensional residual stream. our second insight acknowledges the decadelong interpretability research stemming from topic models. since @xcite , there has been a push to reconcile the evaluation of topic models with the human notion of coherence @xcite @xcite , and still debated @xcite . we show that their findings might help advance the interpretation of tlm (see section 3.1). our third insight contributes towards tackling the disentanglement problem from a novel angle. we map the disentanglement problem to an np-hard classical graph problem and propose a computationally feasible approach to solve it via heuristics and exact solving (see section 5). wary of mirages @xcite , we propose a non-trivial baseline and empirically show that the disentangled superposition interpretations are indeed meaningful and not due to chance (see section 6). contributions. this work aims to deploy topic modelling methodologies, specifically in the area of interpretability, to automate the search and evaluation of concepts within tlms. we propose a novel weights-based, corpus-agnostic approach to disentangle topics from decoder-only tlm with automated evaluation. these extracted topics are constrained to the themes in the chosen corpus. in a static setting, when considering each gpt-2's neurons (residual, mlp) individually, we empirically show that many are in superposition and may embody multiple distinct or related concepts. tlms are pre-trained on a large number of tokens across a plethora of themes. one possible application is to mine the tlm to identify and explain the neurons' role(s) concerning these learnt themes. to further validate our results dynamically, we employ a topic modelling task on gpt-2 in a zero-shot manner (zstm, see section 7). given a corpus as an input, we register top activating neurons and investigate the projected topics. finally, we extend our proposed approach 1 on llama (see section 8), detailing our intuition and approach in overcoming the barrier to interpretability introduced by its choice of tokenization. topic models. while earlier works utilised probabilistic graphical models @xcite , recent works focus on neural-based approaches. there are many ntms @xcite @xcite @xcite that incorporate an autoencoder framework (kingma and welling, 2014), while others @xcite @xcite utilizes graph neural networks @xcite . tlms have a considerable influence on such neural topic models. @xcite al. (2021a), grootendorst (2022 lingual s-bert document embeddings @xcite for use in a zero-shot crosslingual setting. @xcite propose a tlmagnostic method to add additional concept tokens, where from its representation, inferring topics after fine-tuning to any specific supervised task. for our zero-shot topic modelling task, we do not employ additional training or modifications to gpt-2. ## disentangling topics in superposition for any given neuron n in language model m , different inputs to m may activate the same neuron to express different concepts. we consider these concepts to be in superposition within n, and describable in a human-interpretable manner using groups of words -topics. if we can obtain its final logit representation f @xmath0, we posit that it is empirically possible to linearly disentangle f n into multiple interpretable vector representations i in vocabulary space v with a leftover abstract vector representation a n (equation 1 ). similar to the behaviours observed in @xcite , we observed that f n is dividable into three distinct behaviours: promotion, elimination, and mixture (figure 1 ). if i n exists, we expect i n to promote or eliminate specific themes. for uninterpretable mixture behaviour, we liken it to a n . ## path definitions transformer architecture. from @xcite , we adopt a similar notation style. however, we further differentiate the residual stream into higher and lower dimensions, respectively d h and d l , as illustrated in figure 3 . we analyse the activation paths of decoder-only transformers. given a pre-trained language model m with l layers, with tokens t from token-space s, we get its initial embeddings x 0 via embedding layer m e (equation 3). at each layer l, x l-1 ∈ r d l increases its dimensions via w l i , and is passed into transformer block t l . the d h -output of t l is then transformed back to d l space via w l o (equation 3 ). final logits @xmath1. presenting a higher-level abstraction, we omit details that differ between different models, such as positional embeddings and normalisation. usually, in each layer l, attention is applied via the various attention heads h l (equation 5 ) @xcite . attention layer a l is connected to mlp m l in a residual manner (equation 6 ). ## methodology in a word pool g, each word @xmath2, v j ) ∈ e weighted by corpus statistic w v i ,v j . we can formulate extracting b-sized cliques from g as disentangled topics into the following graph problem. order @xmath3.mean_score; chosen ← ∅; sets ← ∅; for @xmath4.@xmath5; maximum edge-weighted b-clique (mewc). given a graph g, with vertices v and edges e, find max-weighted b-sized clique in g. this problem is np-hard @xcite , as it reduces to the maximum clique problem @xcite . as g has a non-trivial amount of vertices, finding a top arbitrary number of non-overlapping mewc in g is computationally challenging. to enable feasible computation to search for highlyweighted b-cliques, we design a star heuristic (algorithm 1), inspired by @xcite 's facetdefining star inequalities, breaking g into smaller sub-graphs for exact solving. exact solving of subgraphs. gouveia and martins (2015) explored the feasibility of solving different convex formulations of mekc on g of varying degrees and sizes. we adapt @xcite 's formulation to fit our use case. we arrive at a mixed integer program (eq. 13) that maximises the weight of selected edges (eq. 13a), selecting edges (i, j) and vertex i by boolean variables e i,j (eq. 13f) and x i (eq. 13g) respectively. we constrain clique size to b (eq. 13d, 13e) and selection of edges within selected clique (eq. 13f, 13g). no further processing is required, as we recover a large portion of v wiki . baseline. we use a random distribution of logits to derive a non-trivial set of coherence topics that occur due to chance. in a random word pool, the number of possible b-sized topics scales exponentially with τ , resulting in favourable odds of producing topics that create a false illusion of interpretability. the random baseline uses the same sampling process s and hyper-parameters of gpt-2 to generate 50,000 random token pools. ## zero-shot topic modelling (zstm) in a traditional topic modelling task, we seek to derive k number of topics of size b for corpus analysis. in this zstm task, since the pre-training of the tlm is not specific to the selected corpora, the corpus becomes the analysis tool for tlm instead. if our interpretation of its neurons is valid, given a corpus, we expect neurons with similar concepts to be activated frequently. otherwise, there will be a contradiction. obtaining a competitive topic set will validate our previous findings. we do not require additional tokens, prompting, or fine-tuning. approach. we implement a voting scheme giving each token some votes for each layer, examining neurons that exceed the vote threshold (see appendix a.2). for gpt-2, we extract disentangled topics from paths lo-c and lo-i of the selected neurons. to select the final k topics, we use a greedy algorithm (lim and lauw, 2022) optimizing on corpus npmi c with diversity constraints. corpus. we generate sets of 20 topics for 20newsgroup foot_1 , bbc-news (lim and buntine, 2014), m10 @xcite , and dblp @xcite processed from octis @xcite , with additional sets of 50 topics for 20newsgroup and bbc-news. corpus details in appendix a.2. metrics. to account for rare occurrences, we set a minimum count of word pairs for npmi c calculation, measuring relevance to the corpus. npmi w evaluates the generality and coherence of the topic. desiring uniqueness amongst topics, topic uniqueness (tu) @xcite , measuring the proportion of unique words in k topics. to obtain competitive sets of results, we must overcome a three-way trade-off between the metrics. the top k extracted topics from gpt-2 models performed the best on 20newsgroup ( ## extending to llama llama @xcite is another popular decoder-only tlm. its 32k tokens s llama encoded using sentencepiece's bpe @xcite . cross-referencing s llama with v wiki produces few common words, as its tokens are sub-words which are uninterpretable as-is, necessitating additional processing of token pools. observing the possibility of recovering coherent topics directly from gpt-2's token space and considering a word as a topic of sub-words, we posit that it is possible to recover vocabulary from s llama . we decode each lower-cased @xmath12. projecting the tokens onto v wiki , we consider a word recovered when all of its constituent tokens are present in the token pool. it is possible to recover 36k vocabulary words in v wiki from s llama . 5 using the previous analysis on gpt-2, we examine all layers of llama models sized 7b and 13b. ## discussions aside from comparing across different sizes within the model family, we can also compare across model families. there are differences between the results between gpt-2 and llama. in this section, we discuss the possible cause and effect. topical sparsity. the larger sizes of tlms might induce greater topical sparsity, with larger models able to spread their learnt concepts over a larger number of neurons. with greater topical sparsity, there is greater difficulty in inferring concepts directly from the neuron as its characteristics may be similar to the random baseline. with millions of extracted pos pools topics from complete paths of both models, the topics from llama are fewer per neuron and less coherent when compared to gpt-2. hence, it is likely that the learnt concepts, with similar themes to wikipedia, are less concentrated within llama. ## conclusion we demonstrate that our approach from a novel angle, disentangling superposed topics from tlm via a graph-based formulation optimizing topic modelling metric npmi w , revealing and analysing superposed topics on gpt-2, evaluated on zstm, and extendable to llama. additionally, we show that tlm is quantifiable by topic coherence and the number of superposed topics in an automated manner. comparison using a random baseline gives us confidence when validating interpretations to be meaningful and not due to chance. these metrics might advance our understanding of architectural design decisions and warrant future investigation on other open-sourced tlm.