Activation Steering for Robust Type Prediction in CodeLLMs

Abstract

Contemporary LLMs pretrained on code are capable of succeeding at a wide variety of programming tasks. However, their performance is very sensitive to syntactic features, such as the names of variables and types, the structure of code, and presence of type hints. We contribute an inference-time technique to make CodeLLMs more robust to syntactic distractors that are semantically irrelevant. Our methodology relies on activation steering, which involves editing internal model activations to steer the model towards the correct prediction. We contribute a novel way to construct steering vectors by taking inspiration from mutation testing, which constructs minimal semantics-breaking code edits. In contrast, we construct steering vectors from semantics-preserving code edits. We apply our approach to the task of type prediction for the gradually typed languages Python and TypeScript. This approach corrects up to 90% of type mispredictions. Finally, we show that steering vectors calculated from Python activations reliably correct type mispredictions in TypeScript, and vice versa. This result suggests that LLMs may be learning to transfer knowledge of types across programming languages.

1 Introduction

Large Language Models trained on code (CodeLLMs) are one of the most successful applications of LLMs to date. They are the foundation for both developer tools (e.g., Weiss & Yahav (2013); git (2021)) and for reasoning agents (e.g., Zelikman et al. (2024); Yang et al. (2024); Hu et al. (2024); Li et al. (2022)). However, despite their utility, CodeLLMs can still be unreliable. Recent work shows that even the most capable models are not robust to small semantics-preserving changes such as refactorings and variable renamings (Hooda et al., 2024; Tambon et al., 2024). Thus CodeLLMs often cannot be deployed without careful supervision, such as a human-in-the-loop or robust code sandboxing.

In this paper, we present an approach to make CodeLLMs more robust using activation steering (Li et al., 2024; Rimsky et al., 2023; Turner et al., 2023; Subramani et al., 2022). Activation steering is an inference-time model editing technique that modifies the intermediate computations (activations) of a model using steering vectors to steer model behavior toward desired outcomes. We present a novel technique for constructing steering vectors for the code domain by taking inspiration from mutation testing (DeMillo et al., 1978). Mutation testing automatically constructs minimal semantics-breaking code edits. In contrast, we construct steering vectors from minimal semantics-preserving code edits that lead to CodeLLM mispredictions. The nature of code allows us to construct these code edits in a sound and scaleable way.

This paper focuses on robust type prediction for gradually typed programming languages, specifically Python and TypeScript. A gradually typed programming language allows programs to freely mix typed and untyped code, giving programmers more flexibility than traditional static ty** affords (Siek & Taha, 2006; Tobin-Hochstadt & Felleisen, 2006). Given a partially typed program $p$ written in a gradually typed language, the type prediction task is defined as follows (Migeed & Palsberg, 2020). Choose an untyped variable binding $\mathit{var}\in p$, predict a type annotation $\mathit{var}:T$, and insert the annotation back into the program to get a new program $p^{\prime}$. The task is successful if $p^{\prime}$ continues to pass the type-checker.

We present an approach for steering type prediction from pairs of prompts $(x_{1},x_{2})$ that represent two type prediction tasks for the same variable binding $\mathit{var}:T$. However, the CodeLLM successfully predicts $T$ for $x_{1}$ but mispredicts it for $x_{2}$. In each pair, $x_{1}$ is from natural data, whereas $x_{2}$ is constructed from $x_{1}$ using semantics-preserving program edits (§ 3.1). We carefully construct steering vectors that narrowly target particular types of program edits. Moreover, we also combine edits to study the robustness of steering for different types of mispredictions. In our experiments, we find that steering activations at multiple adjacent layers can significantly improve accuracy (§ 4.1).

Finally, we are surprised to find that type prediction steering transfers across programming languages. In particular, the steering vectors that we construct from TypeScript prompts are very effective at steering Python type predictions, and vice versa. We conduct a careful analysis and conclude that is it likely that multi-lingual CodeLLMs are learning a representation of types that is shared across the programming languages we study (§ 4.5).

2 Background and Related Work

⬇

return s[::-1] == s

   ⬇ return s[::-1] == s<fim_middle>

Figure 1: An example of a code-infilling prompt using FIM special tokens. The correct prediction expected after the FIM middle token is str.

CodeLLMs and fill-in-the-middle

CodeLLMs are capable of performing a variety of downstream tasks such as generating code from natural language, explaining code, generating tests, and more (Nam et al., 2024; Schäfer et al., 2024). Contemporary CodeLLMs are decoder-only transformers trained on vast amounts of source code (Li et al., 2023; Rozière et al., 2024; Guo et al., 2024). The aforementioned models are also trained to fill-in-the-middle (FIM) (Bavarian et al., 2022; Fried et al., 2023). FIM training does the following (1) It splits $\approx 50\%$ of training items into three chunks—prefix, middle, and suffix—of random lengths; (2) It adds special tokens to the start of each chunk to demarcate chunk boundaries; and (3) It reorders the middle chunk to appear last. The language modelling training objective remains unchanged. Thus FIM inference allows us to generate the middle chunk, conditioned on the prefix and the suffix (figure 1). The resulting model can still be used for conventional left-to-right generation.

Neural type prediction

In this paper we study activation steering for type prediction, formulated as a fill-in-the-middle task. Prior work uses this formulation to evaluate base models (Yee & Guha, 2023; Fried et al., 2023), but does not study models’ internal mechanisms or try to improve task performance in a targeted way. There is prior work that trains smaller, specialized models for type prediction (Hellendoorn et al., 2018; Jesse et al., 2022; 2021; Pandi et al., 2021; Wei et al., 2020), but contemporary CodeLLMs outperform these specialized models (Yee & Guha, 2023; Fried et al., 2023).

Classical type prediction and type inference

Type prediction is also known as type migration, and is distinct from type inference as found in languages such as OCaml and Haskell. In those languages, every variable is typed, even if the types are implicit (Harper & Mitchell, 1993). In contrast, gradual type prediction can change program semantics (Phipps-Costin et al., 2021). There is work on gradual type prediction that uses human-written constraints and constraint solvers (Phipps-Costin et al., 2021; Rastogi et al., 2012; Siek & Vachharajani, 2008; Campora et al., 2018; Henglein & Rehof, 1995; Cartwright & Fagan, 1991). But, these papers present algorithms for variations of the lambda calculus or simple functional languages such as Scheme, and have not been scaled to more complex, modern programming languages.

Mechanistic interpretability

Interpretability seeks to understand the causal mechanisms behind model inference. In particular, mechanistic interpretability studies how model behaviors arise from specific components within the transformer architecture. One of the driving motivations of the field is that by understanding how a model works, we can understand how to align models to desired behavior.

Prior work in mechanistic interpretability includes localizing and editing factual associations within the transformer (Meng et al., 2022), as well as probing hidden representations for knowledge of high-level concepts (Li et al., 2024; Dong et al., 2023). These works are forms of implicit evaluation of model ability, as opposed to explicit evaluation using benchmarks (Dong et al., 2023). Mechanistic interpretability has contributed a wealth of methods for implicit evaluation, for example activation patching (Vig et al., 2020; Variengien & Winsor, 2023), which applies patches to model activations in order to produce a change in behavior. Research leveraging activation patching has suggested the existence of task vectors (Hendel et al., 2023; Ilharco et al., 2022). Task vectors are representations within the transformer which encode a high-level notion of the task described by the prompt. The existence of task vectors is further supported by work in activation steering. Steering has been employed to diminish model deceitfulness and sycophancy (Rimsky et al., 2023; Li et al., 2024), which suggests that steering works by modifying the activations corresponding to task vectors. We thus draw on the idea of task vectors to perform our steering experiments. Importantly, steering relies on the linear representation hypothesis, which argues that concepts are represented as directions in the embedding space of a model (Park et al., 2023). We posit that successful steering performs transformations over model activations that shift the task representation towards a target direction.

Despite much work on the mechanistic interpretability of LLMs, similar research on CodeLLMs has been limited. In this paper we focus on types, which are a fundamental feature of programming languages. One obstacle in this domain is that within a program, a type annotation is semantically constrained by both the tokens that precede it and those that succeed it. For example, the type of a function argument is constrained by how the argument is subsequently used in the function body. To overcome this, we use CodeLLMs that are trained on FIM. This allows us to construct prompts for type prediction.

Mutation testing and program transformations

We take inspiration from mutation testing (DeMillo et al., 1978), but our technique is different. The goal of mutation testing is to test a program’s test suite. To do so, a mutator injects small bugs that alter the semantics of a program, such as changing a $0$ to a $1$ or turning $x>y$ into $x<y$. The hypothesis is that a good test suite should be able to catch these artificial bugs, and there is a substantial evidence that the ability to catch both artificial and real-world bugs is strongly correlated (Just et al., 2014). In this paper, we make program edits that would not affect test cases. Instead, we make syntactic edits that lead to type mispredictions for a given CodeLLM.

3 Methodology

3.1 Constructing Steering Datasets

Model choice

We build steering datasets for 1B and 7B parameter StarCoderBase models. These models are trained to fill-in-the-middle (§ 2), which is necessary for the type prediction task.

Source datasets

A distinguishing feature of our work is that we construct steering pairs from natural data (source code from GitHub) instead of using templates to construct simple examples. This provides more evidence that activation steering is robust. For Python, we use ManyTypes4Py (Mir et al., 2021), a dataset for neural type prediction. It features code from 5,382 Python projects that use Python type annotations and successfully type-check. For TypeScript, we filter The Stack (Kocetkov et al., 2023) to find 1.1M TypeScript files that type-check, of which we sample up to 2,000 for steering.

⬇

def __init__(self, x, y):

self.x = x

self.y = y

def delta_x(p: Point , x: float):

p.x = p.x + x

(a) The original program.

⬇

def __init__(self, x, y):

self.x = x

self.y = y

def delta_x(p: Type0, x: float):

p.x = p.x + x

(b) Type renaming.

⬇

def __init__(self, x, y):

self.x = x

self.y = y

def delta_x(p , x: float):

p.x = p.x + x

(d) Type annotation removal.

        

⬇

def __init__(self, x, y):

self.x = x

self.y = y

def delta_x(p: Point, tmp: float):

p.x = p.x + tmp

(c) Variable renaming.

Figure 2: Examples of three semantics-preserving edits. The type prediction site is float. We carefully ensure that each edit is internally consistent. E.g., in (2(c)), when we rename the binding x to tmp, we rename references to the binding. To construct a steering pair, we repeatedly apply edits until the CodeLLM mispredicts.

⬇

class KafkaAvroBackend(RepositoryBackend):

def __init__(

self, config__tmp0: dict, producer=AvroProducer, loader=AvroMessageLoader,

value_serializer: Callable = to_message_from_dto,

get_producer_config: Callable = get_producer_config,

get_loader_config: Callable = get_loader_config

) -> None:

producer_config = get_producer_config(config__tmp0)

loader_config = get_loader_config(config__tmp0)

...

Figure 3: A fragment of a Python steering pair. The original code is 70 lines of text. The dict is the expected prediction. But, renaming config to __tmp0 makes the model mispredict Repository, which is a hallucinated type.

Semantics-preserving code edits

For a given model $M$, our steering dataset is a set of triples $(x^{+},x^{-},t)$ where $t$ is a type and $x^{+},x^{-}$ are prompts that represent fill-in-the-middle type prediction tasks with $t$ as the expected type for both prompts. However, the maximum likelihood generation for $x^{+}$ is $M(x^{+})=t$ and for $x^{-}$ it is $M(x^{-})=k$, where $k\neq t$.

To construct the positive prompt and determine its target type $t$, we select a Python or TypeScript file from one of our source datasets. Within that file, we select the type on a type-annotated variable binding $\mathit{var}:t$ as the target type. With these selected, we build a fill-in-the-middle prompt. Let the text of the file be $p\cdot t\cdot s$, where $p$ is the text that precedes the selected type annotation $t$ and $s$ is the text that succeeds $t$. The positive prompt thus becomes $\langle\mathtt{fim\_prefix}\rangle p\langle\mathtt{fim\_suffix}\rangle s\langle\mathtt{fim\_middle}\rangle$ (figure 1).

To construct the negative prompt, we incrementally apply semantics-preserving code edits to the positive prompt until $M$ mispredicts $t$. We implement the following edits. 1) Rename variable:We select a variable binding (from a function definition) and rename it to an arbitrary name that does not conflict with other variables in the program. We also rename all bound occurrences of the same variable so that the program’s semantics remain unchanged. 2) Remove type annotation:We select a type annotation (excluding the target $t$) and delete it. In a gradually typed language, removing or relaxing an annotation does not alter program semantics. 3) Rename user-defined type:We select an arbitrary type definition (e.g., a class name or a type alias) and rename it to an arbitrary name that does not conflict with other names in the program. 4) Rename builtin type:We also support renaming builtin types by introducing type aliases. Figure 2 illustrates several separate edits to a program.

By incrementally applying edits, the model eventually mispredicts $t$ (or else we fail to find an $x^{-}$ and thus discard $x^{+}$). Figure 3 illustrates a real example from our dataset where StarCoderBase-1B mispredicts after an edit. Note that a single edit often alters the positive prompt $x^{+}$ at several points. Moreover, when we apply several edits, the two prompts $x^{+},x^{-}$ may become significantly different from each other. We find that the mean edit-distance for our Python and TypeScript steering pairs is $74$ and $327$ respectively. This is in contrast to steering pairs constructed from sentence templates, that only differ by a small, fixed number of words at a far smaller number of places.

We hypothesize that mispredictions happen because the model fails to identify the inference task as a type prediction task. This likely occurs because our edits shift the prompt outside the distribution of a model’s training data. We thus call our edits distractors, since they distract from the correct task. We leverage activation steering to make the model more robust against such distractors and other semantically irrelevant features of code. For example, distractors often remove important in-context clues like function names which models rely on for prediction. Relying on textual meaning rather than program structure is a limitation of CodeLLMs that we improve with activation steering.

Class balance

The distribution of type annotations in natural data is heavily skewed: builtin types occur far more often than the long tail of user-defined types. To prevent bias towards majority-class types in our steering vectors, we balance every dataset $\mathcal{D}$ to ensure that no target type $t$ occurs more than $25$ times. In our evaluation, we also measure the effectiveness of steering by type (§ 4.2).

	Python		TypeScript
Edit Type	Steering	Held-out	Steering	Held-out
Rename variables	$1,924$	$773$	$1,400$	$539$
Rename types	$1,413$	$484$	$1,094$	$457$
Remove type annotations	$1,952$	$892$	$698$	$297$
Rename variables & types	$2,000$	$1,051$	$2,000$	$811$
Rename variables & remove type annotations	$1,842$	$1,773$	$1,668$	$599$
Rename types & remove type annotations	$2,000$	$781$	$1,619$	$445$
All edits	$2,000$	$1,572$	$2,000$	$762$

Table 1: Overview of StarCoderBase-1B dataset sizes and applied code edits. Note that steering subsets are used for constructing steering vectors, while held out subsets are used to evaluate steering performance on out of distribution data. Dataset sizes vary because we only apply edits to programs where the model correctly predicts types. We then filter to remove overlap in source programs between held-out and steering sets. This accounts for the variation in sizes.

Ablation splits and evaluation sets

We construct several datasets $\mathcal{D}$, where each dataset applies a subset (or all) of the program edits described above to programs in a language. We further split each dataset into a steering subset and a held-out evaluation subset. We ensure that the two subsets have prompts based on different source programs, to test how steering generalizes to unseen programs. Table 1 summarizes the dataset sizes for StarCoderBase-1B. We report the sizes for StarCoderBase-7B in appendix A.

3.2 Constructing Steering Vectors

Given a dataset of steering pairs $(x_{i}^{+},x_{i}^{-},t)\in\mathcal{D}$ and model $M$, we apply a forward pass to every $M(x_{i}^{+}),M(x_{i}^{-})$ to collect model activations on the last token. Specifically, we extract activations from the residual stream of the model, which is the output of a transformer layer $\ell_{i}$ that in turn becomes the input to subsequent layers $\ell_{i+1...n}$. We write $A_{\ell}(x)$ to denote this activation tensor at some layer $\ell$ for prompt $x$. Thus we compute steering vectors $\mathbf{t}_{\ell}$—one for each layer—as the mean difference between positive and negative activations at that layer:

$$\mathbf{t}_{\ell}=\frac{1}{|\mathcal{D}|}\sum_{(x_{i}^{+},x_{i}^{-},t)\in\mathcal{D}}\left(A_{\ell}(x_{i}^{+})-A_{\ell}(x_{i}^{-})\right)$$ (1)

The intuition behind equation 1 is that the distance between positive and negative pairs in activation space encodes the transformation for steering towards the correct type.

Following the FIM format, the last token in all prompts is $\langle\texttt{fim\_middle}\rangle$. The prompt provides the model with context indicating that a type annotation should follow—namely, the code preceding the type annotation site, followed by a colon and the remaining code (e.g. figure 1). We choose to intervene on the $\langle\texttt{fim\_middle}\rangle$ token because we hypothesize that models capable of FIM type prediction build latent representations of types in the residual stream of the last token before prediction. This is in contrast to Rimsky et al. (2023) that calculates steering vectors using the residual stream of the expected token. Conversely, in our experiments we steer a model by adding the steering vector $\mathbf{t}_{\ell}$ to the last token’s residual stream at layer $\ell$. In the next section, we perform steering at the level of both individual layers and sets of adjacent layers simultaneously.

4 Evaluation

4.1 Layer Ablations

Figure 4: Steering accuracy for StarCoderBase-1B on the held-out set of the TypeScript Datasets. Steering vectors are patched onto different subsets of adjacent layers.

We evaluate the accuracy of steering on each edit dataset (table 1). To find the optimal layer for applying steering vectors, we conduct an ablation on TypeScript data with StarCoderBase-1B. We consider patching at single layers as well as intervals of three and five adjacent layers. Figure 4 shows that for code edits where steering is most effective (e.g., renaming types), patching at single versus multiple layers makes little difference. However, for the less effective edits (e.g., removing type annotations), patching multiple layers significantly increases accuracy. We hypothesize that this performance gap exists because task vector refinement occurs over multiple layers, thus needs to be steered over multiple layers. This follows from previous ideas that transformers build predictions incrementally throughout layers (Elhage et al., 2021; Geva et al., 2022). Following this finding, we chose to patch layers 10–14 in StarCoderBase-1B and 19–23 in StarCoderBase-7B.

	Python		TypeScript
Edit Type	Steering	Random	Steering	Random
Rename variables	0.24	0.17	0.29	0.26
Rename types	0.90	0.10	0.75	0.11
Remove type annotations	0.48	0.20	0.41	0.31
Rename types & variables	0.56	0.13	0.49	0.09
Rename types & remove annotations	0.60	0.11	0.63	0.11
Rename variables & remove annotations	0.39	0.18	0.30	0.22
All edits	0.36	0.16	0.51	0.11

(a)

	Python		TypeScript
Edit Type	Steering	Random	Steering	Random
Rename variables	0.20	0.11	0.26	0.09
Rename types	0.84	0.07	0.69	0.07
Remove type annotations	0.49	0.11	0.39	0.18
Rename types & variables	0.56	0.10	0.36	0.05
Rename types & remove annotations	0.59	0.09	0.67	0.05
Rename variables & remove annotations	0.41	0.10	0.27	0.08
All edits	0.50	0.10	0.40	0.04

(b)

Table 2: Steering StarCoderBase-1B and 7B on several datasets of semantics-preserving code edits. For each dataset, we report the accuracy on a held out set of negative prompts, i.e., the model mispredicts all of these types without steering and steering corrects a significant number of mispredictions. The Random column reports accuracy after steering with a randomly initialized steering vector.

4.2 Evaluating Different Edits and Effectiveness By Type

Figure 5: Mean accuracy of steering Python by expected type label frequency. Error bars indicate the interquartile range. Note that we constrain the maximum size of any class to 25.

Having picked the layers to steer, table 2 summarizes the results of steering on each of our models. For StarCoderBase-1B, steering is most effective at fixing distractors from renaming types, achieving an accuracy of 90% in Python and 75% in TypeScript. We report accuracy over the held-out evaluation sets for each edit type. However, we also evaluated steering accuracy on the steering set and found no significant difference in performance between steering splits or held-out splits (appendix B). This is a first indication that our steering vectors generalize well.

To ensure that steering performance is not disproportionately affected by accuracy on a single type, we conduct an analysis of the results of steering per type label. In figure 5, we report steering accuracy on Python factored by target type frequency. Appendix C has a similar result for TypeScript. We find that the performance of steering vectors is well-distributed across types, and not skewed towards the majority class of labels.

4.3 Backup Circuits

A threat to validity for any intervention that involves activation patching is that the patch may not be directly improving performance, but just triggering a backup circuit (McGrath et al., 2023). For example, it could be that patching is just introducing noise into the embedding space that causes other mechanisms to trigger, thus producing the expected outcome. This effect complicates the interpretability of patches as well as steering vectors.

As part of our evaluation, we steer the model with a random steering vector (the Random columns in table 2). We find that steering in this way has non-zero accuracy, though it is generally significantly lower than our computed steering vectors. We posit that the non-zero accuracy of the random vector occurs due to backup circuits. However, the more significant gains from our computed steering vectors indicate that our steering method performs true transformations towards the correct type.

4.4 Steering Versus Finetuning

Fine-tuning is the traditional approach to improving model performance on a downstream task, so we now compare steering to fine-tuning. We fine-tune StarCoderBase-1B on the Python Renamed Types dataset. We fine-tune for five epochs using the AdamW optimizer and a learning rate of $3\times 10^{-5}$ with weight decay of $0.1$. At the end of each epoch, we evaluate on the held-out set.

In figure 6(a), we see that fine-tuning is approximately as effective as steering on the type prediction task. However, figure 6(b) shows that this comes with a caveat: fine-tuning on type-prediction significantly degrades the model’s code completion ability, as measured with HumanEval Chen et al. (2021). In contrast, since the steering vectors for type prediction are a small, localized patch, one can easily toggle them on or off based on the task at hand. On the other hand, fine-tuning for type prediction produces a new specialized model. We have shown that this fine-tuned model is weak at code completion, but in general, other abilities may have been impacted too.

(a) Type prediction accuracy.

(b) HumanEval pass@1.

Figure 6: We fine-tune StarCoderBase-1B on the Python Renamed Types dataset. We find that within 1 epoch, the fine-tuned model achieves the same accuracy on the held-out set as the steered model (figure 6(a)). However, this fine-tuning makes the model much worse at code completion (figure 6(b)).

4.5 Language Transfer in Steering Vectors

Edit Type	$\mathrm{TS}\rightarrow\mathrm{Py}$	$\mathrm{Py}\rightarrow\mathrm{Py}$	$\mathrm{Py}\rightarrow\mathrm{TS}$	$\mathrm{TS}\rightarrow\mathrm{TS}$
Rename variables	0.24	0.24	0.20	0.29
Rename types	0.85	0.90	0.77	0.75
Remove type annotations	0.38	0.48	0.36	0.41
Rename types & variables	0.38	0.56	0.55	0.49
Rename types & remove annotations	0.60	0.60	0.64	0.63
Rename variables & remove annotations	0.26	0.39	0.36	0.30
All edits	0.27	0.36	0.34	0.51

Table 3: Accuracy of steering StarCoderBase-1B. A column labelled $A\rightarrow B$ indicates that the steering vector is computed from language $A$ but evaluated on language $B$.

Syntactically, Python and TypeScript look very different. However, semantically the two languages have a lot in common Politz et al. (2013); Bierman et al. (2014). In particular, both of them use gradual ty** disciplines with many shared semantic features. So, we wondered, could CodeLLMs be learning a representation of types that is shared between these two languages?

To test this hypothesis, we test if steering vectors built on TypeScript data can improve the accuracy of Python type prediction, and vice versa. We conduct this experiment with each of our datasets: we steer StarCoderBase-1B using vectors from language $A$ but evaluate on the corresponding held-out test set from language $B$. Table 3 shows that this is nearly as effective as steering type prediction on the same language.

From this observation, we hypothesize that within the model there exists a representation of program types that is shared across TypeScript and Python. This is in contrast to the scenario where distinct model components are specialized to each programming language. Thus, steering vectors are effective across languages because they intervene on the shared representation between them. Moreover, these results suggest that steering vectors generalize by targeting a high-level representation of types.

5 Conclusion

We investigate activation steering for type prediction by making CodeLLMs more robust against semantically-irrelevant aspects of code. We find that by constructing steering pairs using semantics-preserving code edits we can construct highly effective steering vectors. Our experiments show that steering vectors generalize outside the programs used for steering and outperform naive baselines. We find that our steering vectors are transferable across languages. This suggests the existence of a shared representation of a type within the CodeLLM.

Activation steering can be a powerful technique for improving model performance on tasks where fine-tuning is not advantageous. As CodeLLMs are trained on more programming languages and downstream tasks, activation steering may be used as a lightweight alternative to multiple fine-tuned experts. Rather than fine-tuning several models on different tasks, the same base model could be specialized through steering vectors which are added and removed as needed. This could be especially convenient for resource-constrained applications. Our type prediction vectors, for example, could make a convenient lightweight expert for a coding assistant, useful for applications like type migration.

In future work, we aim to study the underlying causal mechanism in CodeLLMs responsible for type prediction. We further wish to explore how type-prediction steering vectors may generalize to open generation problems like code completion.

6 Ethical Considerations

The purpose of this work is to investigate methods for making CodeLLMs more robust and aligned to user intent. It is our view that interpreting CodeLLMs is necessary for understanding whether models approach programming in a principled way. As models become more integrated into developers’ workflows, model errors could compromise the security of entire systems. For this reason, we focus on making models more robust against distractors.

The datasets used in our investigation are sourced from publicly available code. Our TypeScript dataset is derived from a subset of The Stack v1.2, which contains permissively licensed data with personal identifying information (PII) filtered. The ManyTypes4Py dataset is funded by the European Commission, which follows data privacy laws under the EU General Data Protection Regulation (GDPR). We utilize the StarCoder family of models for their open and publicly available weights and training data.

7 Reproducibility

We commit to making all of our code and datasets public for the purpose of reproducibility and furthering research in the mechanistic interpretability of CodeLLMs.

Appendix A StarCoderBase-7B Steering Dataset Sizes

Table 4 reports the sizes of our steering and evaluation datasets for StarCoderBase-7B.

	Python		TypeScript
Edit Type	Steering	Held-out	Steering	Held-out
Rename variables	$1,094$	465	944	385
Rename types	796	343	773	332
Remove type annotations	$1,528$	706	875	344
Rename variables & types	$1,622$	759	$1,423$	577
Rename variables & remove type annotations	$2,000$	$1,246$	$1,270$	478
Rename types & remove type annotations	$2,000$	977	893	383
All edits	$2,000$	$1,044$	$1,663$	630

Table 4: Overview of StarCoderBase-7B dataset sizes and applied code edits.

Appendix B Steering Accuracy on Construction Set

In § 4.2 we evaluated steering accuracy on held-out evaluation data. In this section, we evaluate accuracy on programs used to construct the steering vectors themselves. Our results in table 5 show that there is no significant difference in accuracy between held-out and steering sets.

	Python		TypeScript
Edit Type	Steering	Held-out	Steering	Held-out
Rename variables	0.23	0.24	0.29	0.29
Rename types	0.89	0.90	0.71	0.75
Remove type annotations	0.44	0.48	0.41	0.41
Rename types & variables	0.56	0.56	0.47	0.49
Rename types & remove annotations	0.60	0.60	0.61	0.63
Rename variables & remove annotations	0.39	0.39	0.33	0.30
All edits	0.42	0.36	0.49	0.51

(a)

	Python		TypeScript
Edit Type	Steering	Held-out	Steering	Held-out
Rename variables	0.19	0.20	0.29	0.26
Rename types	0.82	0.84	0.62	0.69
Remove type annotations	0.46	0.49	0.40	0.39
Rename types & variables	0.50	0.56	0.36	0.36
Rename types & remove annotations	0.52	0.59	0.63	0.67
Rename variables & remove annotations	0.33	0.41	0.29	0.27
All edits	0.46	0.50	0.40	0.40

(b)

Table 5: Steering StarCoderBase-1B and 7B on several datasets of semantics-preserving code edits. For each dataset, we report the accuracy on a held out set of negative prompts. Under the Steering column, we report the accuracy of steering vectors on the set of negative prompts used to construct the vectors.

Appendix C Accuracy by Type

We investigate whether steering vectors are biased towards particular types. Specifically, we look at whether steering vectors have higher accuracy on type labels that appear more frequently in the steering set. Typically, these types are builtin types such as string or str which are most commonly represented in type annotations. We look at the All Edits subsets of Python and Typescript in figure 5 and figure 7, respectively. For each type label frequency ranging from 1 to 25, we plot the mean accuracy of a steering vector on the held-out set. Results show that steering vectors are not biased towards more frequent type labels.

Figure 7: Accuracy by type for the TypeScript All Edits dataset. Error bars represent the interquartile range.