Neural Decompiling of Tracr Transformers

TL;DR

This work tackles automatic interpretability for transformers by training TraDe, an encoder–decoder decompiler, to translate Tracr-generated transformer weights back into readable RASP code. It builds a large synthetic dataset of weight–RASP pairs by generating and validating random RASP programs, then trains a Transformer-based model to map weights to RASP tokens using a structured, one-hot line representation and weight-tokenization. Empirical results show that TraDe achieves exact token reproduction for about 30% of test programs in non-autoregressive mode and up to 26% in autoregressive mode, with substantial rates of compilability and functional equivalence (up to 73%), indicating meaningful progress toward end-to-end interpretability. The study highlights significant challenges for scaling to large, real-world models and suggests future work on decompiling learned weights, compression, and feature-focused analyses to narrow the interpretability gap.

Abstract

Recently, the transformer architecture has enabled substantial progress in many areas of pattern recognition and machine learning. However, as with other neural network models, there is currently no general method available to explain their inner workings. The present paper represents a first step towards this direction. We utilize \textit{Transformer Compiler for RASP} (Tracr) to generate a large dataset of pairs of transformer weights and corresponding RASP programs. Based on this dataset, we then build and train a model, with the aim of recovering the RASP code from the compiled model. We demonstrate that the simple form of Tracr compiled transformer weights is interpretable for such a decompiler model. In an empirical evaluation, our model achieves exact reproductions on more than 30\% of the test objects, while the remaining 70\% can generally be reproduced with only few errors. Additionally, more than 70\% of the programs, produced by our model, are functionally equivalent to the ground truth, and therefore a valid decompilation of the Tracr compiled transformer weights.

Figures (8)

• Figure 1: The core concept of the present method is based on the idea that Tracr can be used to generate training data for an automated end-to-end interpretability system.
• Figure 2: Visualization of the three major steps of Algorithm \ref{['prog_gen_algorithm']} that are repeated until the pool of available inputs converges to one entry. 1: Select a function; 2: Fill the function with variables; 3: Add the newly created variable to the pool of available inputs.
• Figure 3: Four examples of generated programs that are rejected
• Figure 4: Illustration of the vector representation of the RASP code. It works by dividing each line into four components, which are in turn represented by a one-hot-encoded vector, denoting one of the options for this component.
• Figure 5: Illustration of the process of translating the weights of a compiled transformer into a set of tokens.