Mechanistic Interpretability of Large Language Models using Kolmogorov-Arnold Networks

This paper introduces a methodology for mechanistic interpretability in language models that integrates Kolmogorov-Arnold Networks (KANs) and Sparse Autoencoders (SAEs). Our solution substitutes KAN layers for the feed-forward blocks typically found in Transformer architectures, replacing them with interpretable polynomial functions to provide symbolic insights into Large Language Models (LLMs). We extract sparse latent features from model activations using SAEs and trace them back to their respective KAN polynomials. This allows for human-understandable interpretations of linguistic representations. Our experiments demonstrate that KAN-based Transformers perform similarly to conventional Multi-Layer Perceptrons (MLPs) while offering greater transparency. On the TinyStories dataset, KAN Transformers achieve a validation loss of 2.89, nearly identical to the baseline GPT-2 model despite having fewer parameters. Additionally, SAE analysis reveals a tunable trade-off between sparsity and reconstruction accuracy, confirming the feasibility of structured feature extraction. Our case study identifies the polynomials responsible for sentiment processing and proposes a selective manipulation technique that manipulates edge polynomials to change alignment-related outputs. ... mehrRemarkably, modifying fewer than 0.0025% of the coefficients in a single layer is sufficient to shift the sentiment logit difference by over 1.1. Our results suggest that the robust explainability of KAN models could support future audit and threat detection mechanisms for LLMs.