The Problem: Opaque Agent Decisions
For today, we want to focus on what we call the Kiji Inspector. It is a tool that is helping you to understand how agents "Think." The fundamental problem that we have with agents is that we don't know how they actually make their decisions. We have a number of tools out there. We have prompt engineering, behavioral testing, or probably the most sophisticated processes like chain of thought. None of them are actually providing you causality, they're often just giving you plausible explanations or plausible narratives of what the root cause could be. When we were asked to figure out how can we make sure that the model that is running an agent is picking the right tool, how do we know that it's thinking about the right context, we decided that we needed to look into the model.
Part I - The Solution
The way we do this is by using sparse autoencoders. Before we get there, let me start explaining what an autoencoder is. This is a concept that has been around for a very long time, probably three decades by now. Autoencoders take a set of inputs, reduce this to a smaller latent space vector, and then reconstruct them through a decoder. The whole point of that autoencoder is that you want to minimize the error between the input and the reconstructed output. In an optimal system, there would be zero error, and the output here would be exactly as the input. Then you know that the representation of your vector z here in your latent space is most optimal. This has been used many times in various things like embeddings and vector searches and information compression, dimensionality reductions of various sorts of kinds. This is a really interesting way of forcing your model to reduce the amount of information to a smaller vector space.
The cool thing about this concept is that it's self-supervised. You don't need any labeling. You use the input as the output, that's your target, and then in between, that's what you would use in production. With sparse autoencoders, you would flip this concept around a little bit. Here in this case, we wouldn't create a bottleneck for the information, but we want to go wide here, with a little catch, and I'll explain the catch. Let's say you grab an activation from a model, or a hidden state, and you pass it through the sparse autoencoder, you can go wide in the information, and that means you force the model to pick certain features to represent the information. When you take a look at the activations inside of a model, there's something that we call polysemantic representation of the context, which means that one neuron that fires can have multiple meanings or topics it would fire on.
In our feature representation, if we want to understand the model, we need to separate them and put them apart from each other, and so that's why we're going wide in the feature representations here. Then again, we try to then have a decoder where we narrow that information piece again, and we try to reconstruct the activations that we used as an input. The catch here is that in this activation function, we are then forcing the model to only use a small set of features to be actually active, and the rest of them, roughly 95%, that's the ideal number, we want to force them to be zero. In that moment, we actually force the model during the training process to just be really selective on a handful of features, and then we can be absolutely certain that those features are representing the themes and the context of the model itself, what the model was thinking in this given layer at the processing of a certain token.
Then we can go one step further. Once you have those feature vectors you generated, those indices don't mean anything to us, to humans. What we can then do is, once we have a corpus, we can grab all the documents that fire highly on, let's say, dimension 3 and dimension 6, and then we can look at those documents and say, what do they have in common? It could easily be that they say like, the first one has something to do with Python, so then over the course of the entire corpus, we could say, those two features, they represent a programming language or an API vocabulary or things like this.
Training Pipeline and Contrastive Pairs
What we have tried to do in this project is combining a number of things and package them up so that everybody can use them in production. The first thing that we want to do is, sparse autoencoders have been around for a number of years now, and OpenAI published some research, and Eleuther published some research on this, but they focused on the entire LLM generation. What we changed here in our approach is we were only interested in the agentic decision-making. We focused on the decision token rather than the entire prompt, and then averaging those activations across the entire prompt. Then we put a bunch of concepts together. For example, Eleuther came out with this concept of token-level fuzzing, and so we added this to our framework so anybody who wants to use that framework and trains their own models, they can use that token testing. Then also we added causal validation testing as well so that you can make sure, is it actually causal or is it just correlation?
Then the other sweet thing, what we changed here in this project, what is novel, is that we're not just training on a number of prompts, but we created those pairs, training pairs. We have a given intent — I'll show you an example later — but for a given intent, you have something that we call the anchor that's the correct answer and the correct tool use, and then there's an adversarial reply, which is very similar, but very different. It's basically like look something up in an internal database versus look up something online. Two different use cases, very similar in the prompt response, but very different outcome in terms of the tool selection.
Looking at the training pipeline, let's say we have all these pairs generated, which I will show you how we do this. We then take our target model, and little caveat here, those have to be open-source models. We can't do this with OpenAI's models because we don't have access to their weights, but you can do this with any open-source model, and that's another novelty in our framework that we're not just focusing on Gemma models from Google or Nemotron models from NVIDIA. That tool, that framework here is working with various open-source models that are out there. You pick your subject model, that's the model you want to use the explainability for, and then you run all your prompts through the model itself, grab all the activations from the particular layers you want to watch. Then with those activations, you can train your sparse autoencoder as I just showed you with putting them in, widening them up, and then forcing the decoder again to reconstruct that activation.
Then once you have the model trained, you can do a contrastive analysis, and I'll show you this later, how we then figure out through those pairs of what's actually relevant in the context or not. Or you can go to a large language model and say, tell me what those documents have in common, and then you can generate those themes. This all happens in that package where we call the Kiji Inspector. You put in your corpus, you nominate those models, and out you get a feature definition of what those features mean when you run those LLMs and agent requests.
Layer Selection and Results
Earlier, I just mentioned you have to pick a certain layer. There's a little bit of an art what layer to pick rather than science. There's a little bit of an explanation here that if you're early in the stack, like maybe in the first third, the model hasn't fully conceptualized all the concepts of your prompt request. You're basically activating things that are highly relevant to the prompt itself. If you're further down the stack, there's a lot of intricacies of the model itself. Here in this particular model, we're looking at the Nemotron 3 from NVIDIA. That's an MoE model. That's because every different request goes through different routings, you have crazy errors. Learning something later on is very difficult. Also, the concepts you extract are so abstract in some cases that even for humans, it's impossible to understand what's going on. Then also for the models, it's difficult.
What we wanted to do is, for our particular experiments, we picked layer 20 because it is already down the pipe. It's like the model extracted some abstract concepts. It's not too far down the pipe that it would be too non-human understandable. In our open-source models, what we're publishing, we're publishing multiple layers. You can pick wherever you want to be in your model layer setup. Again, this is another representation of the training pipeline. Basically, you have your input. You calculate those representations. Then you do the reconstruction back to, in our case, it's 2688 dimensions. Then you calculate the loss. It's a very standard autoencoder model architecture. The key thing, and we show this in the following slide, we're using a particular activation, what we call the Jump Regularized Linear Unit. I will explain this here why this is in this particular case.
Then, how do we actually focus on the decision token? The way we do this is by forcing the agent to respond to a given prompt and says, I will use blank. Whatever tool, whatever the model then responds is expected to be the tool name, because we've provided the selection of tools in the prompt itself. Then in the response, we would see what tool the model has chosen. We do this a million times. We generated a training set with a million samples, basically 500,000 pairs, good and bad examples. The pairs come in later when we try to evaluate the model. During the training process itself, we use them raw, and we have those 1 million samples. Why are we doing this contrastive pair design again? Here's a little example. Let's say you have a given intent, resolve a password issue, and you can have the request, how do I reset my password? That would trigger the tool knowledge base. You could also say, I tried to reset my password three times, and that could be a completely different tool. It's very similar in the response or in the request, but it's a very different tool outcome in this particular case. Later on, what you will see is we will then use those pairs and figure out which of the features we generated are actually relevant for those given tools. When we generate the features for both tools, the dimensions where the features are different have mostly contributed to the decision-making between knowledge base versus ticket lookup.
Jump Regularized Linear Unit
I don't want to go through the math here, but there's an interesting concept regarding the JumpReLU. As you know, Regularized Linear Units basically are, anything negative is zeroed out and anything positive is the input. It's like a linear function, but it's not differentiable because of that bend. Here in the JumpReLU, that's another derivative of the ReLU activation function. It's basically the output is the input, but there's a threshold. The interesting thing here is when we train the model of, figure out which activation is actually relevant and not relevant, we can then tune that threshold during the training process as the output activation. That's another tweak we can have to teach the model which feature is relevant for which given activation.
When we look at some results per layer, then we see some really interesting side effects of early on, the error is low, but the model hasn't learned much. Later on, the error is crazy high because this is logarithmic. Then we have too many, basically we learn something, but the error is exorbitant high. That's why we ended up in our practical example, we ended up somewhere in between at layer 20. Again, if we look at some stats here, like total, I think we had like 2,600-and-something input features, we widened this to 10,000 dimensions. Of those 10,000 dimensions, a little bit above 80% of those dimensions are alive. That means somewhere through the training, they have been nonzero. Then only less than 0.2% are completely dead, meaning those features have never been used through the entire training process. Then when we do the testing later on, we actually figured out that out of those 10,000 features, 668 are actually highly relevant to all the tool selections we wanted to investigate. This is for a given model for a given layer and a given dataset. If you retrain this, you might see different results.
Why Sparse Autoencoders?
Then somebody could say, why do we have to go through all this craziness of sparse autoencoders and training another machine learning model in all these crazy setups? As you can imagine, it's a lot of computation there. This is basically a couple comparisons. If you take the activations and then train a linear model on this, could we then figure out is this tool correct or incorrect? If you compare this and a binary classification here, we get to almost 80% accuracy. It's not bad, but it's also not perfect. Then if you do some just accuracy versus F1 score, that's a little bit lower. Then somebody could say, what happens if we just take those activations with those 2,000 dimensions and we just cluster them? Could we from the cluster then figure out which tool is the correct one? You're better off flipping a coin than actually clustering.
I think that was for us the interesting thing to see with the sparse autoencoders. I'll explain what fuzzing is. We can bridge the gap between interpretability and also having some causality testable features.