A Mathematical Theory of the Jacobian Lens (arxiv.org)

Anthropic has unveiled a groundbreaking interpretability tool for AI language models called the Jacobian lens (J-lens), as detailed in a recent paper by Gurnee et al. This mathematical framework sheds light on how the J-lens identifies linear directions in hidden states of language models that share characteristics with working memory-like representations. Central to this innovation is the discovery that the J-lens preferentially detects coherent features—Lyapunov modes—that amplify through model layers and remain stable across different contexts and token positions. This advancement holds significant implications for enhancing our understanding of AI behaviors and predictions. By validating the theory using a Bayes filter and a toy transformer model, the research demonstrates that adjusting task parameters can correspondingly modify the J-lens's responses, providing deeper insights into feature coherence in predictive models. Furthermore, the introduction of a family of J-lenses, termed $\omega$-J-lenses, allows for focusing on various temporal scales of information processing in language models, opening new avenues for research in AI interpretability and functionality.