Understanding How CodeLLMs (Mis)Predict Types with Activation Steering

CodeLLMs are transforming software development as we know it.

This is especially true for tasks where rule-based approaches fall short, like type prediction.

The type prediction task consists in adding a new type annotation to a partially typed program, such that the resulting program is closer to being fully typed.

The intractability of rule-based approaches and high cost of manual annotation make CodeLLMs an attractive solution to the problem.

However, CodeLLMs are still far from being deployed on the large-scale due to doubts surrounding their reliability.

To shed some light on how CodeLLMs approach type prediction, we investigate what happens when a model mispredicts a type.

We show that by applying semantics-preserving edits to code, CodeLLMs are eventually misled into mispredicting type annotations.

However, by leveraging activation steering we are able to "steer" the model back to the correct prediction, making models more robust against semantically irrelevant prompt features.

We show that steering achieves comparable performance to fine-tuning directly on the type prediction task.

Furthermore, we find that steering vectors computed from Python code are effective at correcting TypeScript mispredictions, and vice versa.

To our knowledge, this is the first evidence of its kind to suggest that CodeLLMs learn task representations that transfer across languages.

;Comment: 14 pages, 7 figures