Currently: Dissecting a Coding Model

Four weeks into my sabbatical! My first sabbatical project originally was going to be “embed a bunch of codebases, and render some visualizations.” This led me to the Nomic Code Embedder, which is designed/trained to place code in an embedding space for vector-based code search. And I could have just used UMAP (or this new EVoC thing!) to embed repos, and render them. Would have been straightforward.

Where’s the fun in that though. Not nearly convoluted enough. So of course I got sidetracked.

  • I learned that Nomic Code Embed was trained using a method called contrastive training.
  • I then trained a Sparse Auto Encoder (SAE) on Nomic Code Embed (first ever model training!) so that I could “get cooler annotations” on the code map. Imagine highlighting the chunks of code in the repo associated with a feature! A region lights up green because more of that feature activates there!
  • …but features don’t come labeled. You have to label them yourself. So I started building a labeling pipeline.
  • Oh and Nomic Code Embed might have trained out some of the characteristics I want, so I built another SAE on Qwen Coder 2.5 instead (which is what Nomic Code Embedder is based on)

This is where things went truly off the rails. I started asking Claude about how SAEs actually work. Learned a whole bunch of jargon:

  • Superposition - there are more concepts to learn than there are dimensions/neurons (?) to learn them, so models squeeze multiple concepts into one dimension.
  • Therefore, neurons are polysemantic (i.e. represent multiple semantic concepts at once)
  • SAEs use a sparsity constrained model to try to learn to break those down into single concept “features” (or latents)

Took a while to really understand what this all means. What I really wanted though, was to label maps, and annotate code. I’m a designer! Give me something I can use to color stuff. So I went and read some more.

  • Well, an SAE feature is technically a derivation of the value of the residual stream
  • A feature has two faces! What it reads from the residual stream (and its KV value cache)
  • …and what it writes back to the residual stream, in the form of a decoder direction. (?)
  • All of this corresponds to a particular token in the input.
  • …but also influenced by the tokens that came before in the context.

All of which makes SAE features more difficult to interpret, and difficult to make claims about. After weeks of wrestling with the model, the SAE, and the literature, one key is:

Some SAE features are about content. Most of them are not.

Not terribly profound. Critical though, for my purposes, i.e. to build a cool code feature visualization. So now my research splits into two halves

  1. For the features that are content oriented, how do I validate that they are indeed what they appear to be? We ablate to check they seem to stop firing when the token or concept disappears. We find “nearby” chunks to check if we are accurately bounding the feature, i.e. not too narrow or general.
  2. For the features that are not content oriented… well what are they?

Ok. Back to listening to Neel Nanda podcasts.