compmech-bowman-lu

This repository contains the code base for "Unsupervised Recovery of Hidden Markov Models from Transformers with Evolutionary Algorithms", a technical report from Colin Lu and Dylan Bowman for research conducted during the Computational Mechanics Hackathon organized by PIBBSS and Simplex from June 1, 2024 to June 3, 2024.

Update 6/15: This project won first prize in the hackathon.

Abstract

Prior work finds that transformer neural networks trained to mimic the output of a Hidden Markov Model (HMM) embed the optimal Bayesian beliefs for the HMM’s current state in their residual stream, which can be recovered via linear regression. In this work, we aim to address the problem of extracting information about the underlying HMM using the residual stream, without needing to know the MSP already. To do so, we use the $R^2$ of the linear regression as a reward signal for evolutionary algorithms, which are deployed to search for the parameters that generated the source HMM. We find that for toy scenarios where the HMM is generated by a small set of latent variables, the $R^2$ reward signal is remarkably smooth and the evolutionary algorithms succeed in approximately recovering the original HMM. We believe this work constitutes a promising first step towards the ultimate goal of extracting information about the underlying predictive and generative structure of sequences, by analyzing transformers in the wild

Code

Setup

Assuming you already have a Python environment set up, clone the epsilon-transformers (GitHub) repository and follow the setup instruction for installation with pip. Then, clone this repository.

Exploratory Analysis

• The src folder contains two scripts, generate_paths_and_beliefs.py and experiment.py. These scripts are meant to be run sequentially.
  • generate_paths_and_beliefs.py: For each set of parameters in src/msp_cfg.yaml, generate the optimal Bayesian belief states for every possible input sequence. The output tensor is cached in a .pt file in src/cached_belief_store. src/cached_belief_store comes pre-populated with the beliefs from our grid search so there isn't a need to run this script unless you're generating beliefs for new sets of parameters.
  • experiment.py: For each set of parameters in src/msp_cfg.yaml and for each of the pretrained models (Mess3(0.15, 0.6) and Mess3(0.05, 0.85)), we train a new probe and print out its $R^2$ and MSE scores. This script produces visualization/r2.pkl for seeing how well a probe fits transformer activations to different Mess3 processes. By default, this script will also produce reconstruction visualizations in src/images folder for each set of parameters.
    • --no-image to disable reconstruction image. Image takes >1 min to generate, so this flag increases performance substantially for the probe training.
    • --device to set the PyTorch device, which speeds up probe training substantially when set to a GPU.
• The visualization folder contains two IPython notebooks, MSP_visualization.ipynb and r2_visualization.ipynb.
  • MSP_visualization.ipynb: The primary section is titled Mess3 MSP: Chaos Game Fractal, and plots Mess3 MSPs for various alpha and x values. See examples for usage. The code also supports arbitrary 3-state HMMs defined by emission_and_transition_pi.
  • r2_visualization.ipynb: Plot the $R^2$ values (via colors) for how good the linear fit is from transformer activations to Mess3 MSP geometries, for different values of alpha and x. Can observe that $R^2$ is approximately unimodal and convex, and maximized close to the actual Mess3 parameter values that generated the transformer's training sequence.

Evolutionary Algorithms

• evo_alg_demo.ipynb runs evolutionary search over parameter sets using probe $R^2$ as a reward signal.