Large Language Models Relearn Removed Concepts (2024)

Abstract: Advances in model editing through neuron pruning hold promise for removing undesirable concepts from large language models. However, it remains unclear whether models have the capacity to reacquire pruned concepts after editing. To investigate this, we evaluate concept relearning in models by tracking concept saliency and similarity in pruned neurons during retraining for named entity recognition tasks. Our findings reveal that models can quickly regain performance post-pruning by relocating advanced concepts to earlier layers and reallocating pruned concepts to primed neurons with similar semantics. This suggests that models exhibit polysemantic capacities and can blend old and new concepts in individual neurons. While neuron pruning provides interpretability into model concepts, our results highlight the challenges of permanent concept removal for improved model *safety*. Monitoring concept reemergence and developing techniques to mitigate relearning of unsafe concepts will be important directions for more robust model editing. Overall, our work strongly demonstrates the resilience and fluidity of concept representations in LLMs post concept removal.

## introduction large language models possess neurons that encode semantic concepts across different languages @xcite , architectures @xcite , and modalities @xcite . the primary objective when pruning such models is to eliminate redundant neurons, while preserving the most crucial ones @xcite . the underlying assumption is that removing important "concept neurons" will disrupt the model's structured internal representation of key concepts, causing performance degradation. however, it is in fact possible for models to regain high performance after pruning random or important neurons @xcite . models demonstrate a remarkable capacity to adapt and regain conceptual representations. this phenomenon, which we call "neuroplasticity", suggests a degree of adaptability in such models. neuroplasticity has significant implications for model editing. model editing promises to remove undesirable concepts, but neuroplasticity implies that those concepts may in fact reappear if retraining takes place. understanding how concepts are represented, redistributed, and recaptured is thus crucial for the development of safer, fairer and aligned models. furthermore, understanding how removed concepts are recovered in language models can significantly enhance their robustness. while the inherent plasticity of neural networks has been acknowledged @xcite , there is a limited understanding of where relearned concepts are redistributed, or what influences relearning in specific neurons. such insights could benefit models' ability to recover from partial damage and information loss. this paper investigates how large language models fine-tuned for named entity recognition relearn and redistribute concepts after the removal of important features, to explain how a model regains performance after experiencing damage to crucial neurons in its internal representation structure. we instigate neuroplasticity by pruning important concept neurons and retraining a model so that it regains its original performance on named entity recognition tasks. we then analyze how the distribution of concepts changes, and examine the relationship between concepts previously associated with a neuron and the concepts it relearns. our findings are four-fold: these findings contribute to a deeper understanding of how language models learn, adapt, and retain core conceptual representations. they underscore the potential of earlier layers to recapture fundamental representations, which has implications for model editing. furthermore, our exploration of how neuroplasticity increases polysemanticity, where a neuron can represent multiple concepts and meanings may inform strategies to enhance the transfer of learned representations, and facilitate interpretability analysis regarding polysemanticity. figure 1 : process of investigating neuroplasticity in a large language model. we identify concept neurons (dark blue) in the base model, and prune them (white). we then retrain the model until it regains its original performance and identify new concept neurons. ## related work in terms of analysing the distribution of concept representations, it is known that more complex concepts tend to be represented by neurons in higher layers of language models. @xcite found that neurons which capture word morphology were primarily located in the lower and middle layers, and those learning about advanced concepts such as syntax were found at the higher layers. however, dynamic changes in concept distribution over training are less well studied. prior works artificially redistributed concepts in a large language model by modifying the activations of specific neurons which capture desired properties, to control model outputs @xcite . however, these works focused on deliberately editing representations in a trained model. there is little work in which models redistribute concepts unsupervised, leaving a gap in understanding about how concept redistribution naturally occurs after pruning. in terms of performance recovery after pruning, @xcite demonstrated that randomly pruned sparse networks can be trained from scratch to match the accuracy of equivalent dense networks, particularly in larger networks. similarly, @xcite suggested that the inherent plasticity of deep convolutional neural networks allows them to recover from the loss of randomly pruned filters. their work supports the idea that large neural networks demonstrate plasticity, and can regain performance if important neurons might be pruned. however, their approach uses indiscriminate random pruning of neurons in image models. in contrast, we explicitly prune the most important neurons, which allows us to analyze causation between removing crucial features and performance. in order to clearly analyze changes in the model before and after pruning, our approach also trains, prunes and retrains the model in separate phases -as opposed to other approaches, which iteratively prune and adjust weights @xcite . our paper is most similar to @xcite , in which neurons that are redundant with regards to a specific task are pruned from llms, and model performance is re-examined. however, @xcite focus on identifying redundant neurons and do not damage key concept neurons. on the other hand, we focus on how the internal distribution of concept representations and model performance changes after pruning important neurons. ## problem setting consider a large language model m . let n be the set of all neurons in m . consider a concept c represented within m . a concept neuron is a neuron in n which activates strongly on tokens associated with c. let s ⊂ n be the set of concept neurons in m prior to ablation. let the set of remapped neurons r ⊂ n be the set of concept neurons associated with c in m , after pruning neurons in s and retraining. we analyze the redistribution of concepts in m after neuroplasticity occurs, in terms of concept saliency and concept similarity. ## method to explore the notion of neuroplasticity within a pretrained model m , we first fine-tune the model for a task that calls for the understanding of a specific concept c. subsequently, we employ a probeless neuron search method to pinpoint the set of concept neurons s within the model that correlate with c ( §4.1). in order to remove the concept from the model's internal representations, we prune the neurons in s ( §4.2). following the pruning process, the model is retrained until it attains its original performance level. the recovery of performance signifies that neuroplasticity has taken effect, and that the eradicated concept has been relearned and integrated throughout the model. we can use the same probeless search method to determine the set of remapped neurons r in the retrained model that have acquired the pruned concept. during the retraining period, we can examine how concept saliency and similarity evolve across neurons in both s and r at two-epoch intervals. the overall process of our methodology is shown in figure 1 . for detailed steps, refer to the pseudo-code algorithm 1 in appendix e. ## experimental setup we focus on pruning the specific concept of location names from a distilbert model, a distilgpt2 model @xcite and a gpt-2 model (radford et al., 2019) fine-tuned for named entity recognition. analysis of the models across different runs are in appendix b, and analysis of the models after pruning other concepts are in appendix c. ## results in our results, performance metrics indicate that neuroplasticity occurs early in the retraining process ( §6.2). specifically, neuroplasticity occurs when pruned concepts are relocated to earlier layers in a model ( §6.3). this may be because neurons which reacquire the pruned concept are primed for relearning, having captured similar concepts initially ( §6.4). concept similarity scores show that the pruned concept is relearned in addition to original concepts, in a polysemantic manner ( §c.5). ## conclusion we explored neuroplasticity in large language models (distilbert, distilgpt2 and gpt2) and implications on the distribution of concept representations. our analysis showed that these models redistribute concepts after the removal of important features, allowing them to regain performance within a few epochs of retraining. the recovery in performance can be attributed to how pruned concepts, initially present in later layers, are remapped to remaining neurons in earlier layers which cap-ture similar concepts. we further identified that neurons exhibit polysemantic properties as they relearn a blend of new and old concepts. these insights contribute to a deeper understanding of how language models learn, adapt, and retain core conceptual representations. they also suggest interesting research directions in model editing and transfer learning. for instance, models may need to be edited throughout retraining to ensure that pruned concepts are not recaptured by earlier layers during fine-tuning. studying the implications of neuroplasticity-induced polysemanticity can also aid the development of interpretable models and the enhanced transfer of learned representations.