Classifying hallucinations in a small language model
My first model ever: a classifier that predicts when a small LLM is hallucinating.
This was the very first time I ever trained a model. We wanted to predict whether Liquid AI's LFM2-1.2B was hallucinating, by interviewing it on a factual dataset and learning a classifier over signals we collected while it answered.
The idea
We interviewed LFM2-1.2B on a factual dataset and, as it answered, captured model internals. Because we know whether the model got each answer right, we could train a classifier on those internals to predict factual errors.
What we captured
For the first generated token, we grabbed the last attention layer and the last fully-connected layer, plus the token probabilities, and trained a classifier on top of that.
The quantization struggle
My biggest challenge was dealing with quantization. I had never heard of that, and I tried to grab internals while running the model on my laptop. But an optimized, quantized runtime doesn't expose those first-token internals the way I was expecting (I later learned that FlashAttention was at fault). Figuring out why was a big lesson in how models actually run in deployment.
What I learned
- My first end-to-end experience training a model and working with local LLMs.
- What quantization is, and why it limits which internal signals are available in deployment.
What I'd do differently
- Use the quantized moodel as it'd actually be deployed, and build the classifier on full-generation token-probability distributions instead of hard-to-access internals.
- We defined a hallucination as getting the final fact of a question wrong, but you can define “wrong” per use case: a chess tutor making an illegal move, for example.
- Treat hyperparameter search and regularization (dropout, more data) systematically instead of by trial and error. I later learned how to do this in my CS231n project.