Concept Bottleneck Language Models For Protein Design

## Abstract

Abstract We introduce Concept Bottleneck Protein Language Models (CB-pLM) 1 1 1 Code and models available in the Lbster Repository . , a generative masked language model with a layer where each neuron corresponds to an interpretable concept. Our architecture offers three key benefits: i) Control: We can intervene on concept values to precisely control the properties of generated proteins, achieving a 3 × \times × larger change in desired concept values compared to baselines. ii) Interpretability: A linear mapping between concept values and predicted tokens allows transparent analysis of the model’s decision-making process. iii) Debugging: This transparency facilitates easy debugging of trained models. Our models achieve pre-training perplexity and downstream task performance comparable to traditional masked protein language models, demonstrating that interpretability does not compromise performance. While adaptable to any language model, we focus on masked protein language models due to their importance in drug discovery and the ability to validate our model’s capabilities through real-world experiments and expert knowledge. We scale our CB-pLM from 24 million to 3 billion parameters, making them the largest Concept Bottleneck Models trained and the first capable of generative language modeling.

### 1 Introduction

Protein Language Models (pLMs) have emerged as a prominent framework for protein representation learning, encapsulating millions of years of protein evolution (Hayes et al., ). These models have demonstrated exceptional performance on intricate tasks such as predicting protein structure and function (Lin et al., ; Chen et al., ; Elnaggar et al., ; Xu et al., ). Furthermore, they have opened new frontiers in diverse and critical applications, including healthcare and drug discovery (Hie et al., ; Hayes et al., ). Despite their impressive capabilities, controlling, interpreting, and debugging pLMs remains challenging due to the opaque nature of transformers (Vaswani, ).

Recently, some pLMs have supported conditional generation by adding various protein input representations, functions, and properties as tags (which can be viewed as concepts) and predicting them during training (Madani et al., ; Hayes et al., ; Ruffolo et al., ). However, these models can selectively rely on or ignore different aspects of the input.
Even if these models offer some form of control, there is no effective way to interpret or debug them. Domain experts might want to know which concept the model depends on the most to generate the amino acid “E”, which amino acids are correlated with enzymatic activity, or answer counterfactuals like if we want to decrease hydrophobicity, which amino acid should we replace “E” with, which is not possible with existing models. Debugging language models is notoriously difficult, and it is often unclear what a model has learned or when it will fail. These limitations lead to distrust among domain experts, as humans tend to not trust what they do not understand.

Figure: Figure 1: Changing protein structure using Concept Bottleneck Protein Language Model.
Refer to caption: extracted/5988250/figures/fig1_control_struc.png

Our work relies on the fact that neural networks do not have to be black boxes; since we construct and train these networks, we can design and align them according to our needs. This is a unique advantage of building models from scratch, with control, interpretability, and debugging in mind from the outset (Frey et al., ; Adebayo et al., ). To achieve explainability and alignment, it is essential to use intermediate representations based on well-understood concepts. To this end, we present Concept Bottleneck Language Models, a novel architecture for large language models that relies on human-understandable concepts via a concept bottleneck layer and auxiliary loss terms. This architecture enables concept-controllable generation and explanations, compelling the model to align with human-understandable concepts and providing the means to debug the model and understand its limitations. Our approach builds on the foundation of concept bottleneck generative models (CBGMs) (Ismail et al., ).

In this paper, we focus on protein design, a domain where control (Gruver et al., ) and interpretability (Adebayo et al., ) are critically needed. Protein design benefits from a set of well-defined biophysical and bioinformatics properties that domain experts understand and wish to control. We introduce Concept Bottleneck Protein Language Model (CB-pLM). We summarize our contributions as follows:

- •
We propose a novel architecture, training, and intervention schemes for large language models that incorporate human-understandable concepts.
This results in a controllable, interpretable, and debuggable LLM by design. We focus on masked protein language models, given their critical role in high-stakes applications such as drug development, where the need for control and interpretability is paramount. We train masked protein language models with 24M, 150M, 650M, and 3B parameters using the proposed CB-pLM architecture with over 700 concepts. We do not observe a significant increase in perplexity compared to unconstrained models.
- •
We compare CB-pLM to various conditional pLMs of the same size, all trained on the same dataset, across more than 80 single and multi-property control in silico experiments. Our model demonstrates $3\times$ better control in terms of change in concept magnitude and a $16\%$ improvement in control accuracy compared to other architectures. Additionally, we benchmark CB-pLM against state-of-the-art (SOTA) protein design models explicitly trained to optimize a single concept. Remarkably, our general-purpose CB-pLM, trained to learn over 700 concepts, delivers results comparable to SOTA models, while maintaining the naturalness of the protein.
- •
We demonstrate how CB-pLM allows domain experts to assess what the model has learned, i.e., interpretability, and identify and correct any unwanted behavior in the model, i.e., debuggability by simply visualizing the weights from the model’s decoder layer.

### 2 Concept Bottleneck Language Model

We aim to create a language model that we can easily control, interpret, and debug. To achieve this, we modify the standard masked language model architecture by adding a concept bottleneck layer to explicitly incorporate human-understandable concepts into the language model—termed a concept-bottleneck (masked) language model (CB-LM). The concept layer can then be used to control, interpret, and debug the model. This section discusses the proposed language model architecture, loss functions, training and intervention procedure.

###### Setup

Each sequence $x$ is represented with a set of tokens, i.e., $x=[x_{0},...,x_{s}]$, where $s$ is the maximum sequence length. Each sequence starts with the CLS token i.e. $x_{0}=\texttt{[CLS]}$ followed by a sequence of amino acids ending with [EOS] as well as padding tokens for sequences below maximum length. Additionally, each sequence is annotated with a set of global predefined human-understandable concepts $c=[c_{0},...,c_{k}]$ where $k$ is the number of concepts. Taken together, this results in a training dataset consisting of token-concept pairs $\{(x,c)_{i}\}_{i=1}^{n}$ where $n$ is the number of samples.

#### 2.1 Architecture

Figure: Figure 2: Concept Bottleneck Language Model
Refer to caption: x1.png

The overall architecture is shown in Figure 2.
We start with a standard BERT-based masked language model (MLM) (Devlin et al., ), which is commonly used in many protein language models (Rives et al., ; Lin et al., ; Hayes et al., ; Frey et al., ). A mask token [MASK] is randomly applied to a subset of tokens in the sequence $x$ resulting in as masked version of the sequence . After masking, ${x}$ is passed into a multi-head self-attention transformer encoder, and the output of the encoder is denoted as $\mathbf{H}=\text{TransformerEncoder}({x})$, where $\mathbf{H}=[\mathbf{h}_{0},\dotsc,\mathbf{h}_{s}]$. We introduce three main architectural changes in CB-LM: (a) concept bottleneck module (CB-module), (b) orthogonality network, and (c) linear decoder.

Concept bottleneck module
After ${x}$ passes through the encoder, the hidden representation of the classification [CLS] token, $\mathbf{h_{0}}$, which serves as an aggregate representation of the sequence, is passed to the CB-module shown in Figure 2b. The concept network  (Koh et al., ) learns a function $g$ that maps the input sequence into concept predictions $\hat{c}=g(\mathbf{h}_{0})$. Each concept is also represented with a learnable embedding $e$. The final representation of concept $i$ is $z_{i}=\hat{c_{i}}\times e_{i}$. The output of the CB-module, $z=[z_{1},z_{2},\dots,z_{k}]$, is referred to as the known embedding, as we know exactly which concept each neuron in the embedding corresponds to.

Orthogonality network
This network aims to produce a token representation devoid of any concept information, which we refer to as the unknown token embedding $\widetilde{h}$. This is achieved by enforcing an orthogonality loss, as described in Section 2.2, between its output and the known concepts.

Linear decoder The known and unknown embedding are concatenated and passed into a single linear layer followed by a softmax function, i.e., $\hat{x}_{\text{M}}=\text{softmax}(W^{\text{dec}}\cdot[z,\widetilde{h}_{\text{M
}}])$, where ${{x}}_{\text{M}}$ represents a masked token [MASK]. We use a linear layer as the decoder for interpretability and debugging purposes. A linear layer allows us to calculate the contribution of each concept to the final prediction by simply multiplying the concept activation value with its weights. This enables us to answer questions such as Which concept does the model depend on the most to generate the amino acid“E”? or counterfactual questions like if we want to decrease hydrophobicity, which amino acid should we replace “E” with?. Debugging neural networks is notoriously difficult; however, our setup simplifies this process. Since we have prior knowledge of how challenging a concept is to learn, we can identify and fix bugs by inspecting the concept weights.

Examining the weight matrix $W^{dec}$ can (a) help us understand the correlation between different concepts, aiding in the detection of any spurious correlations learned by the model, (b) elucidate the correlation between concepts and tokens, answering questions like which amino acids are correlated with enzymatic activity?, and (c) identify concepts that we can control and those we cannot, i.e., if the weight for a concept is small, then the model will not be able to control that particular concept.

#### 2.2 Loss Functions and Training

###### Loss functions

Following CBGMs (Ismail et al., ), concept bottleneck language models are trained with three losses: a standard masked language modeling loss $\mathcal{L}_{\text{MLM}}$, concept loss $\mathcal{L}_{\text{Concept}}$ and an orthogonality loss $\mathcal{L}_{\text{orth}}$. The final loss is given by $\mathcal{L}_{\mathrm{total}}=\mathcal{L}_{\text{MLM}}+\alpha\mathcal{L}_{\text
{Concept}}+\beta\mathcal{L}_{\text{orth}}$, where $\alpha$ and $\beta$ are hyperparameters.

- •
Generative masked language modeling loss: $\mathcal{L}_{\text{MLM}}=-\mathbb{E}_{x,m}\left[\frac{1}{m}\sum_{i\in m}\log P
(x_{i}\mid x_{\setminus m})\right]$, where $m$ is the set of positions of the randomly masked tokens, and the model is trained to predict the identity of the masked tokens.
- •
Concept loss: Given the general-purpose nature of large language models, they are expected to handle thousands of concepts with categorical or real values. Often, samples lack values for most concepts. To address this, we normalize all real-valued concepts to [0, 1] and apply mean square error loss on the entire concept embedding: $\mathcal{L}_{\text{Concept}}=\frac{1}{k}\sum_{i=1}^{k}(c_{i}-\hat{c}_{i})^{2}$. Missing values are replaced with default values, and their effect is removed from the loss function by considering only non-missing concepts, achieved by masking the errors before backpropagation.
- •
Orthogonality loss: Similar to Ismail et al. (), we encourage the known and unknown embeddings to encode different information by applying an orthogonality constraint (Ranasinghe et al., ), minimizing the cosine similarity between them. This loss encourages, but does not guarantee, disentanglement.

###### Training

CBM: We used independent training for the CB-layer (Koh et al., ) to reduce concept leakage
(Mahinpei et al., ). Token regularization: Feature attribution is needed for coordinate selection during intervention, as discussed in section 2.3. To ensure attribution is reliable, we
followed Adebayo et al. ()’s recommendation and added Gaussian noise to the token embedding during training. Additional details are
available in Appendix A.2.

#### 2.3 Coordinate selection and Interventions

CB-LM enables concept-level test-time interventions, allowing for global sequence control while maintaining token-level predictions. During training, tokens are randomly masked and predicted, which is not ideal for interventions. Instead, we wish to identify tokens that most significantly affect a concept and intervene on those (Gruver et al., ; Adebayo et al., ). For example, to increase the Aromaticity of a protein, we target amino acids that decrease Aromaticity. This process is called coordinate selection.

A straightforward method to measure a token’s effect on a concept is occlusion, where each token is replaced with a reference value (e.g., [MASK]) one at a time to observe changes in the predicted concept. While effective, this method is slow for large-scale applications.
Instead, we use a gradient-based approximation of occlusion, requiring only a single forward pass. We propose using a single step of integrated gradients (Sundararajan et al., ), which can be seen as a first-order Taylor approximation of occlusion. Since our inputs are discrete, we calculate gradients with respect to the learned token embeddings and sum over the embedding dimension.
Let $T\in\mathbb{R}^{v\times d}$ be the learned token embeddings, where $v$ is the vocabulary size and $d$ is the transformer hidden dimension. Then, let $f_{T}(x)=[T_{x_{0}},T_{x_{1}},\ldots,T_{x_{s}}]$. The attribution $A$ for token $t$ of input $x$ on concept $i$ is:

$$ $A(x,i,t)=(T_{x_{t}}-T_{\texttt{[MASK]}})^{\top}\nabla_{T_{x_{t}}}\hat{c}_{i}(f _{T}(x)),$ $$

where $T_{\texttt{[MASK]}}$ is the embedding for the mask token and $\hat{c}_{i}(f_{T}(x))$ is the concept value predicted by our CB-Layer for concept $i$.

Informally, $A(x,i,t)$ can be seen as quantifying how the choice of token $x_{t}$ influences $c_{i}$ in context of the other tokens $x_{-t}$.
If in the training dataset $x_{t}$ co-occurs with below average values of $c_{i}$ when combined with $x_{-t}$, then we expect $A(x,i,t)$ to be negative.
Similarly if $x_{t}$ co-occurs with above average values of $c_{i}$ in the context $x_{-t}$ then we expect $A(x,i,t)$ to be positive.
Hence if we want to intervene to increase the concept value we greedily choose the top-$k$ coordinates sorted by $A$ in descending order, and if we want to decrease the concept value we greedily choose the bottom-$k$ coordinates.
In Appendix D.1, we show that this attribution method closely matches the performance of occlusion and improves control effectiveness by $2\times$ over random masking.

### 3 Protein Representation Learning and Design with Conditional Language models

One can view a concept bottleneck language model as a special type of conditional language model. Current conditional large protein language models involve appending conditioning [Tags] to the sequence along with the input (Madani et al., ; Shuai et al., ; Hayes et al., ) and then using these tags to steer the output. To compare the effect of different architectures on protein representation learning, and design, we train three types of conditional models on the same training dataset (sequences and concepts) with comparable sizes.

- •
Conditional Protein Language Model (C-pLM): These are traditional conditional models, the input to the model is sequences and concepts (which can be viewed as [Tags]); the model is trained to minimize the generative masked language modeling loss, i.e., $\mathcal{L}_{\mathrm{total}}=\mathcal{L}_{\text{MLM}}$.
- •
Conditional Protein Language Model with Classifier (CC-pLM): This class of models has an extra regression head used to predict the concepts (Hayes et al., ). The input to the model are sequences and concepts; the model is trained to minimize the generative masked language modeling loss and concept loss, i.e., $\mathcal{L}_{\mathrm{total}}=\mathcal{L}_{\text{MLM}}+\alpha\mathcal{L}_{\text
{Concept}}$.
- •
Concept Bottleneck Protein Language Model (CB-pLM) (Ours): A family of protein language models with the architecture shown in Figure 2, $\mathcal{L}_{\mathrm{total}}=\mathcal{L}_{\text{MLM}}+\alpha\mathcal{L}_{\text
{Concept}}+\beta\mathcal{L}_{\text{orth}}$.

CB-pLM is designed to be controllable, interpretable, and debuggable. Other conditional language models offer control through [Tags]; however, the model can ignore these tags during generation and fail to learn the underlying biophysical concepts for proteins. CB-pLM learns concepts as an intermediate step and includes a loss to encourage disentanglement between concepts and unknown embeddings. Since these concepts are inside the model, this forces the model to learn and use these concepts for predictions, leading to better control (Section 4). Furthermore, CB-pLM provides mechanisms for interpretation and debugging (Section 5), which are not available in other conditional models.

Training Data We combined sequences from UniRef50 (Suzek et al., ) and SWISS-PROT (Bairoch & Apweiler, ), removing duplicates. Annotations such as protein clusters, organisms, taxons, biological processes, cellular components, and molecular functions were used as concept annotations. Biopython (Cock et al., ) was used to extract biophysical and bioinformatics sequence-level concepts. Overall, all models support 718 concepts; details are in Appendix B.1.

Architecture and Training  Following Frey et al. (), we optimized the architecture and training for efficiency. We removed biases in attention blocks and intermediate layers, increased the effective batch size, used gradient accumulation, and set the masking rate to 25%. We employed the AdamW, mixed precision training, and gradient clipping for stability. More details are in Appendix B.2.

#### 3.1 Model quality

**Table 1: Validation perplexity.**
| Models | 24M | 150M |
| --- | --- | --- |
| LBSTER (No conditioning) | 12.17 | 9.92 |
| ESM2 (No conditioning) | - | 9.58 |
| C-pLM | 11.65 | 10.57 |
| CC-pLM | 12.37 | 11.24 |
| CB-pLM | 11.57 | 10.49 |

Model quality is evaluated using perplexity, which measures how well the model predicts a token based on its sequence. Perplexity ranges from 1 (perfect model) to the vocabulary size (33 for the pLMs vocabulary (Lin et al., )), with higher values indicating random predictions. We assessed performance on a held-out validation set. Table 1 presents the validation perplexity of different models, where LBSTER (Frey et al., ) and ESM2(Lin et al., ) are standard masked language models without conditioning. C-pLM and CC-pLM use ground-truth concepts as inputs; since we lack actual values for all concepts and can only compute Biopython concepts from the sequence, these models were given Biopython concepts with the rest set as missing values. All models learned the training distribution well, with similar-sized models showing comparable perplexity. Notably, CB-pLM achieves higher perplexity among conditioning models, likely because it predicts concepts rather than requiring additional input information. CB-pLM matches state-of-the-art open-source protein masked language models performance, even surpassing them in the 24M parameter models. These results are promising, as CB-pLM, despite its constraints, shows no significant performance drop compared to unconstrained models.
Scaling: To verify that our proposed architecture can scale to larger models we train a 650M and 3B parameter CB-pLM with perplexities 9.07 and 9.64 respectively.

### 4 Controlling and Steering

In this Section, we show how the CB-pLM enables fine-grained control, which makes it useful for a variety of protein design tasks. Sequence-based protein design requires the ability to modulate the amount of certain bio-physical properties present in a model’s output—a controllable generation problem. For example, we might be interested in controlling a single property, i.e., removing a hydrophobic patch from a protein, or multiple properties, i.e., increasing the expression and binding affinity of an antibody (Gruver et al., ; Tagasovska et al., ).

First, we demonstrate how to use CB-pLMs for single and multi-property control. As an example, we instantiate these property control abilities in a Siltuximab—a monoclonal antibody—case study, where we show Grand AVerage of Hydropathy (GRAVY) reduction via CB-pLMs, which alters the protein’s surface hydrophobicity.

#### 4.1 Comparing conditional language model architectures

Here we assess the effectiveness of the CB-pLM, compared to the alternatives described in Section 3, for both single and multi property concept control.

###### Experimental Setup

We examine 14 biophysical concepts computable from a protein’s sequence (e.g., molecular weight, Gravy, Helix fraction, Turn structure fraction; see Appendix B.1). These concepts allow us to measure the success of positive (increase) and negative (decrease) concept interventions for each conditional generation approach. Using a validation dataset, we selected 10,000 sequences with the lowest and highest concept values for positive and negative interventions. To test generation control, we mask 5% of the input sequence (up to 25 amino acids) and intervene on the concept value. For the C-pLM model, masking is random. For CC-pLM and CB-pLM models, we use feature attribution to select tokens to mask (details in Appendix D.1).

Each masked test sample results in a single generated sequence, and change in concept presence, between the original sequence and the generated sequence can be calculated used as estimate of the intervention effectiveness. We use likelihood from an auxiliary autoregressive causal language model to measure the “naturalness of” or feasibility of the sequence (Bachas et al., ).

##### 4.1.1 Single Concept Interventions

Figure: (a) Intervention accuracy.
Refer to caption: x2.png

Concept Control Accuracy After intervention, we calculate the concept control accuracy as the percentage of generated samples where the change in concept presence is in the direction of the intervention (positive or negative) averaged over a random set of examples (random accuracy would be around 50%).
The aggregate average across all concepts is shown in Figure 3(a) .
We find that CB-pLMs outperform other conditional models by over 15%, reaching near perfect accuracy. We also measure the effectiveness of the intervention as the average change in the desired concept value direction; CB-pLM is also $3\times$ higher than other conditional models shown in Figure 3(b).

Change in Concept Distribution Beyond comparing average changes, we assess how much the distribution of values shifts upon intervention while maintaining protein naturalness. Ideally, the concept distribution should shift in the direction of the intervention while preserving the naturalness of the protein. Figure 4(a) illustrates the change in concept and naturalness distributions for a subset of concepts when we intervene on the sequence once. We observe that, for both positive and negative interventions, the CB-pLM effectively shifts the concept distributions in the desired direction for all concepts while preserving the naturalness of the proteins. This performance is superior to other variations of conditional language models.
Additionally, we examine the models’ ability to iteratively shift a concept’s distribution, where the sequence output from the first intervention serves as the input for the second, and so on. Figure 4(b) shows concept distributions over three iterations for both positive and negative interventions. CB-pLM demonstrates the ability to iteratively shift the concept value with increasing effects, whereas other models fail to achieve this. We refer readers to the Appendix C.1.1 for additional results.

Figure: (a) Shift in protein naturalness/concept value after interventions.
Refer to caption: extracted/5988250/figures/single_intervention_main.jpg

##### 4.1.2 Multi Concept Interventions

Here, we take a step further by demonstrating how CB-pLMs can be used for multi-property optimization—a challenging task in drug discovery (Stanton et al., ; Gruver et al., ).

###### Experimental Setup

Here, we compare the effectiveness of the CB-pLMs to other approaches at enabling control of multiple properties.
Specifically, we focus on the concept Grand AVerage of Hydropathy (GRAVY) and the electric charge at pH 7, both standard protein sequence properties correlated with functional properties, such as solubility, viscosity, and aggregation (Bhandari et al., ; Cock et al., ; Obrezanova et al., ).
We sequentially intervene on one concept at a time.
We test both positive and negative interventions after masking 5% of the sequence as described previously.
Similarly, we measure concept control accuracy as the percentage of generated samples that successfully move toward the intervention direction for both concepts.
We report the average accuracy for both positive and negative interventions, with random accuracy expected around 25%.

###### Result

In Figure 5, we observe that CB-pLM achieves a 94% control accuracy; outperforming the second-best conditional model by 34%.
Controlling via the CB-pLM effectively shifts the distributions towards the Pareto frontier, demonstrating its ability to control multiple properties, as shown in Figure 5 (b).
We refer readers to Appendix C.1.2 for additional discussion and results.

Figure: Figure 5: Multi-concept interventions.
Refer to caption: extracted/5988250/figures/multiconcept_24M.png

#### 4.2 Case Study: Siltuximab Protein Design

Figure: Figure 6: Siltuximab surface hydrophobicity.
Refer to caption: extracted/5988250/figures/siltuximab_gravy.jpg

Siltuximab is an FDA-approved monoclonal antibody used to treat disease in the lymph nodes (Van Rhee et al., ).
Treatment involves intravenous infusion, which can cause patient discomfort and complications.
Redesigning Siltuximab to have more favorable biophysical properties is a promising research direction for exploring more effective delivery systems.
Structural analysis of Siltuximab reveals the presence of a large hydrophobic patch ( red patch shown in Figure 6), which has been linked to antibody developability liabilities associated with solubility and aggregation (Bhandari et al., ; Obrezanova et al., ).
Here we test different the effectiveness of the CB-pLM, as well as other conditioning approaches, at redesigning Siltuximab to lower its GRAVY index values, improving its hydrophobicity, so that variants will likely be more soluble.

##### 4.2.1 Experiments

We train a set of state-of-the-art protein design models on the task of GRAVY reduction, apply the models to Siltuximab, and compare the GRAVY distribution of the generated samples.

Baselines We consider a wide variety of representative approaches to guided and unguided sequence generation:
(a) Unguided Discrete Walk-Jump Sampling (WJS) (Frey et al., ) energy- and score-based model with single-step denoising. (b) Implicitly guided PropEn (Tagasovska et al., ) an encoder-decoder architecture that is implicitly trained to optimize a property of interest.
(c) Explicitly guided LaMBO-2 (Gruver et al., ) a classifier-guided discrete diffusion model. Additional details about baselines and training procedures are available in Appendix C.2.

Results We intervene on the GRAVY concept for generated samples for each candidate protein design model, with the constraint that generated samples should be at most an edit distance of 5 away from the Siltuximab sequence.
In Figure 7(a), we show the distribution of designs from each model type. The CB-pLM comes second to LaMBO-2 in terms of the degree of concept distribution shift induced in the generated output, while preserving the naturalness of the protein (Figure 7(b)).
Our findings suggests that a general-purpose CB-pLM trained to learn over 700 concepts delivers comparable results to state-of-the-art conditioning approaches while preserving the naturalness of the protein.

Figure: (a) GRAVY distributions.
Refer to caption: extracted/5988250/figures/gravy_siltuximab_dist.png

### 5 Interpretability and Debugging

Having established the effectiveness of the CB-pLM model for control and steering, we now demonstrate its ability to offer insight i.e., interpretability and enable model debugging. A key component of the CB-pLM is its final linear layer—a deliberate choice to maintain simplicity. By inspecting the weights of this linear layer, domain experts can understand the impact of specific concepts on the model’s output. Moreover, if the model learns undesirable correlations or if the learning process results in important features that deviate from a domain expert’s prior knowledge, a straightforward inspection of the linear layer’s weights can help identify and address these issues.

#### 5.1 Interpretability

In Figure 8, we show the final layer weights for different concepts and amino acids in our CB-pLM (24M).
We find that the model successfully learns several key biophysical relationships.
We highlight some key insights below (the full weight matrix and more details are available in Appendix D):

- •
For both charge at pH 6 and pH 7, acidic amino acids (D, E) have negative weights, while basic amino acids (R, K) have positive weights. The difference in charge between pH 6 and pH 7 reflects the biophysical properties of Histidine (H), as it contains an imidazole side chain with a pKa of 6.0 and contributes more towards positive charge at lower pH (Figures 8(a), 8(b)).
- •
GRAVY, which defines hydropathy based on the Kyte-Doolittle scale (Kyte & Doolittle, ), is consistent with positive weights assigned to A, I, V, F, C, and M amino acids (Figure 8(c) ).
- •
Aromatic amino acids F, Y, and W have high weights for the aromaticity concept (Figure 8(d)).

Figure: (a) Charge at pH6
Refer to caption: extracted/5988250/figures/charge_at_pH6_bar_chart.png

#### 5.2 Debugging

For large models trained on natural language, it is difficult to identify the inputs on which a model will fail;—exhibit undesirable behavior—a challenge now known as LLM red teaming Perez et al. ().
Such inscrutability is compounded in protein language models because model failure cannot simply detected from inspecting the model’s output.
The linear layer, which maps concept representations to an output embedding, helps provide domain expert with ability to identify model failures.
For example, by inspecting the weights of the linear layer, we can detect how the output of a model can be controlled by a given concept.

To showcase the model debugging capability of a CB-pLM, we trained a Bad CB-pLM 24M, where we change the normalization for the aromaticity, making it very difficult for the model to learn and control it.
We want to see if we can, in fact, identify where this model will fail without the need for intervention experiments.
To do this, we examine the weight for aromaticity-amino acid combinations for the Bad CB-pLM (see Figure 9(a)).
By inspection, we observe that weights are all essentially zero; hence, the model will not be able to control aromaticity.
To confirm the result, in Figure 9(b), we show that intervening on the aromaticity concept does not control the model’s output.

Figure: (a) Outgoing weights from aromaticity for model with bad concept normalization.
Refer to caption: extracted/5988250/figures/aromaticity_bad_24M_bar_chart.png

### 6 Discussion and Conclusion

###### Related Work

Language models are increasingly used across various fields, from natural language processing (NLP) to specialized areas like chemistry and biology, demonstrating exceptional performance on complex tasks. However, their lack of effective interpretability methods limits their application in high-stakes scenarios, causing distrust among domain experts and raising concerns about regulatory compliance, safety, and alignment. Current interpretability research focuses on pre-trained models, either by explaining predictions or analyzing internal circuits through mechanistic interpretability. While recent mechanistic interpretability work with sparse autoencoders (SAEs) Cunningham et al. (); Templeton et al. () has shown promise in uncovering interpretable features, this approach involves training large SAEs and assigning meaning to millions of features, which may not be feasible in all domains. Concept bottleneck language models (CB-LMs) offer a more straightforward solution by inherently encoding each concept, eliminating the need for additional SAEs and post-hoc interpretation.

Concept Bottleneck Models enhance neural network interpretability by incorporating a “concept bottleneck” layer that maps inputs to human-understandable concepts for final predictions (Koh et al., ). Recently, Ismail et al. () extended this approach to generative models (CBGMs), enabling controllable generation and concept-level explanations, particularly in image generation tasks. In this paper, we adapt this approach to generative language models, achieving global control over the entire input while retaining token-level generation. A few very recent works have proposed CBMs for text classification tasks Tan et al. (); Sun et al. (), but this is a much easier and less interesting setting than generative language modeling which we address.

Protein language models are extensively used in biological machine-learning research, trained on vast protein sequences across the evolutionary tree of life. Despite their critical applications in healthcare and drug discovery (Hie et al., ), pLMs currently lack interpretability. Historically, pLM architectures, loss functions, and training setups have closely mirrored the original masked language model setup (Devlin et al., ), differing mainly in vocabulary and dataset. While some protein language models support conditional generation by concatenating different protein functions and properties to the inputs (Shuai et al., ; Madani et al., ; Hayes et al., ), they do not provide mechanisms for interpretability, debuggability, or insights into what the model has learned. See Appendix E.1 for extended related work.

###### Discussion

In this paper, we applied our proposed architecture to masked language models for proteins. This architecture can be adapted to any domain, such as NLP, by simply changing the training dataset, vocabulary, and concepts while keeping the model unchanged. It can also be adapted to other language models, such as autoregressive models, as shown in Figure 24 in Appendix E.2. Our architecture supports the addition of new concepts; given a new dataset, we can fine-tune CB-LM with a new concept by copying weights from the original model, with only the new concept embedding, classifier head, and linear decoder being randomly initialized. The new model is then trained to match the old model’s output for existing concepts while aligning with the ground truth values for the new concept. In addition, CB-LM can be used to generate sequences with unseen concepts or novel combinations, as long as they can be expressed as functions of existing ones. We demonstrate this capability in a toy task detailed in Appendix C.3.2.

###### Limitations

One main limitation of CB-LM, like other concept bottleneck models (Ismail et al., ; Koh et al., ), is the need for the training dataset to be annotated with predefined concepts. Another drawback is the necessity for unknown embeddings. Generally, we believe this is unavoidable, as it is impractical to encode all the information needed for a generative task into a fixed number of predefined concepts; excluding unknown embedding will lead to unacceptable performance losses.  However, reliance on unknown parts can challenge control and interpretability, potentially causing the model to ignore concept outputs. Our method mitigates this with orthogonality loss to improve control, but some concept information can remain in the unknown part, and better disentanglement methods are needed.

###### Conclusion

We present Concept Bottleneck Language Models (CB-LM), a controllable-interpretable-debuggable language model. We have demonstrated the effectiveness of our proposed architecture by applying it to masked protein language models. Our results show that this architecture can scale to large models, as evidenced by training Concept Bottleneck Protein Language Models (CB-pLM) with 24M, 150M, 650M, and 3B parameters to encode over 700 human-understandable concepts while maintaining pre-training perplexity comparable to traditional masked protein language models. Our findings indicate that CB-pLM offers superior control compared to conditional language models, along with interpretability and debugging capabilities. This work represents a significant step towards more interpretable, reliable, and debuggable neural networks.

## Appendix

### Appendix A Additional details on Concept Bottleneck Language Model

#### A.1 Loss Function

Following CBGMs (Ismail et al., ), concept bottleneck language models are trained with three losses: a standard generative masked language modeling loss $\mathcal{L}_{\text{MLM}}$, concept loss $\mathcal{L}_{\text{Concept}}$ and an orthogonality loss $\mathcal{L}_{\text{orth}}$. The final loss is given by

$$ $\mathcal{L}_{\mathrm{total}}=\mathcal{L}_{\text{MLM}}+\alpha\mathcal{L}_{\text {Concept}}+\beta\mathcal{L}_{\text{orth}},$ $$

where $\alpha$ and $\beta$ are hyperparameters.

- •
Generative masked language modeling loss:
$\mathcal{L}_{\text{MLM}}=-\mathbb{E}_{x,m}\left[\frac{1}{m}\sum_{i\in m}\log P
(x_{i}\mid x_{\setminus m})\right],$
where $m$ is the set of positions of the randomly masked tokens, and the model is trained to predict the identity of the masked tokens.
- •
Concept loss: Given the general-purpose nature of large language models, they are expected to handle thousands of concepts with categorical or real values. Often, samples lack values for most concepts. To address this, we normalize all real-valued concepts to [0, 1] and apply mean square error loss on the entire concept embedding:
$\mathcal{L}_{\text{Concept}}=\frac{1}{k}\sum_{i=1}^{k}(c_{i}-\hat{c}_{i})^{2}.$
Missing values are replaced with default values, and their effect is removed from the loss function by considering only non-missing concepts, achieved by masking the errors before backpropagation.
- •
Orthogonality loss: Following Ismail et al. (), to ensure effective control over the model’s output, it is crucial to prevent unknown concepts from being mere transformations of known concepts. This is achieved by enforcing an orthogonality constraint (Ranasinghe et al., ), which minimizes the cosine similarity between the concept context embedding and the unknown context embedding. The orthogonality loss is defined as follows:
$\mathcal{L}_{\mathrm{orth}}=\sum_{j\in B}{\frac{\sum_{i=1}^{i=k}\big{|}\langle
w
_{i}\,,\,w_{k+1}\rangle\big{|}}{\sum_{i=1}^{i=k}1}}$
(1)
where $\langle\cdot\,,\,\cdot\rangle$ is the cosine similarity applied to two embedding, $|\cdot|$ is the absolute value, and $B$ denotes mini-batch size and $j$ denotes each sample in the mini-batch. The cosine similarity in the above equation involves the normalization of features such that $\langle x_{i}\,,\,x_{j}\rangle=\frac{x_{i}\cdot x_{j}}{{\|x_{i}\|}_{2}\cdot{\|
x_{j}\|}_{2}}$, where ${\|\cdot\|}_{2}$ is $l_{2}$ norm.

#### A.2 Training for intepretability

Reducing concept leakage through independent training: One known problem in CBMs is concept leakage (Margeloiu et al., ; Mahinpei et al., ; Havasi et al., ); this happens when soft concept representations encode more information than the concepts themselves. Concept leakage affects both model interpretability and control. Mahinpei et al. () discussed that this leakage happens when CBMs are trained jointly or sequentially. However, leakage is not possible when CBMs are trained independently (i.e., during training, the post-CB-layer part of the network takes the ground-truth concepts themselves rather than the output of the CB-layer). For this reason, we choose to use independent training for CB-LM.

Faithful token attribution through regularization: Apart from concept-level explanations, users might be interested in token-level explanations to identify which tokens were most influential during generation. Token-level explanations can also be used to identify which token to mask (i.e., the token that has the most effect on a particular concept) when intervening on different concepts (Gruver et al., ). Gradient-based explanation methods Baehrens et al. (); Simonyan et al. (); Zeiler & Fergus () are popular approaches for ‘explaining’ models, but many works (Adebayo et al., ; Hooker et al., ) showed that such feature attribution methods are no better than random ranking of the input features. The reliability of feature attributions seems to be correlated to how the model is trained; Shah et al. () showed that gradient-based feature attributions of adversarially robust image classifiers are faithful, unlike their non-robust counterparts; Adebayo et al. (), showed that the addition of Gaussian noise to the input has the same effect as adversarial training for inducing attribution faithfulness. Following Adebayo et al. ()’s recommendation, we add Gaussian noise to the token embedding during training for model regularization and faithful token attribution.

### Appendix B Additional details on Concept Bottleneck Protein Language Model

#### B.1 Training Data

We combined sequences from UniRef50 (Suzek et al., ) and SWISS-PROT (Bairoch & Apweiler, ), removing duplicates. Annotations such as protein clusters, organisms, taxons, biological processes, cellular components, and molecular functions from SWISS-PROT were used as concept annotations. Biopython (Cock et al., ) was used to extract biophysical and bioinformatics sequence-level concepts. We filtered out rare SWISS-PROT concepts; the final list of concepts is shown in 2.

**Table 2: Training concepts.**
| Class | Number of Concepts | Calculated from sequence |
| --- | --- | --- |
| Cluster name | 159 | ✗ |
| Biological process | 140 | ✗ |
| Cellular component | 162 | ✗ |
| Molecular function | 106 | ✗ |
| Organism | 36 | ✗ |
| Taxon | 101 | ✗ |
| Biopython | 14 | ✓ |

###### Biopython

All concepts were computed from the Biopython ProtParam module.

- •
Molecular weight is the mass of the protein.
- •
Charge at pH 6 and pH 7 is the net charge of the protein at pH 6 and pH 7.
- •
Isoelectric point is the pH at which the protein has no net charge.
- •
Aromaticity is the relative frequency of aromatic amino acids (F, W, Y) in the sequence
- •
Instability Index is the sum of the instability weights from sequence decomposition.
- •
Secondary structure fraction is the relative frequency of amino acids in helix: V, I, Y, F, W, L. Amino acids in Turn: N, P, G, S. Amino acids in sheet: E, M, A, L.
- •
Molar extinction coefficient considers the weighted sum of W, Y, and C for protein absorption of 280 nm.
- •
Gravy is the sum of hydropathy values associated with the amino acids in a sequence divided by the total number of amino acids, based on the Kyte and Doolittle hydrophobicity scale.
- •
Protein Scale An amino acid scale is defined by a numerical value assigned to each type of amino acid. The most frequently used scales are the hydrophobicity or hydrophilicity scales, the secondary structure conformational parameters scales and surface accessibility.

#### B.2 Model parameters and configurations at different scale

|  | 24M | 150M | 650M | 3B |
| --- | --- | --- | --- | --- |
| Number of layers | 10 | 27 | 33 | 26 |
| Embedding dim | 408 | 768 | 1280 | 2560 |
| Attention heads | 12 | 12 | 20 | 40 |
| Concept embedding dim | 2 | 2 | 2 | 2 |
| Learning rate | 0.001 | 0.0001 | 0.0001 | 0.0001 |
| Clip norm | 0.5 | 0.5 | 0.5 | 0.5 |
| precision | 16 | 16 | 16 | bf16 |
| Warmup steps | 3000 | 10000 | 30000 | 30000 |
| Effective batch size | 512 | 1024 | 1024 | 1024 |
| Distributed backend | ddp | ddp | ddp | deepspeed_stage_1 |

###### Training:

During training we mask percentage is 25% of the sequence. We then truncate all sequences to a maximum length of 512. Sequences of length less than 512 were padded, but no loss was backpropagated through the network for padding tokens. We use Rotary Position Embedding.

### Appendix C Control Experiments

#### C.1 Comparing conditional language model architectures in silico

###### Experimental Setup

We evaluate control on 14 different concepts that can be calculated from Biopython. We want to see if we can intervene in the models by increasing the concept values (positive interventions) or decreasing the concept values (negative interventions). To accomplish this, we utilized the held-out validation dataset; for each concept, we selected the 10,000 sequences with the lowest concept values and the 10,000 sequences with the highest concept values. We use the lowest value sequences for positive interventions (i.e., we intervene in the model to increase the value of that given concept), and the highest value sequences are used for negative interventions (i.e., we intervene to decrease the concept value). We mask a percentage of the test sequence and intervene on the concept to set it to a minimum or maximum value based on the intervention type. For the conditional model, C-pLM, masking is done randomly; for the conditional classifier CC-pLM and our proposed CB-pLM model we select the tokens that contributed to the given concept the most using feature attribution (Gruver et al., ; Adebayo et al., ) (additional details on how feature attribution is done is available in Appendix D.1) and mask them. For each test sample, we only consider a single generated sample (i.e this can be viewed as single shot generation). After generation, we calculate the new concept values of the generated sequence. We measure accuracy as the percentage of generated samples that successfully moved toward the direction of the intervention. We use likelihoods from an auxiliary autoregressive causal language model to measure the “naturalness of” or feasibility of the sequence (Bachas et al., ).

##### C.1.1 Single concept intervention

In this section, we focus on single-concept intervention; the intervention procedure is as follows:

Figure: Algorithm 1 Single Concept intervention procedure

##### One time intervention

This sections shows results when setting $n=1$ in Algorithm 1.

###### Accuracy

After the intervention, we calculate the accuracy as the percentage of generated samples that successfully moved toward the direction of the intervention for positive and negative interventions, then average them (random accuracy would be around 50%). The average across per concept is shown in Figure 10(b)). Our findings demonstrate that in terms of intervention direction accuracy, CB-pLM outperforms other conditional models on every single concept. Figure 11(b) shows the average change in the desired direction of the concept value of different concepts after our interventions. This is averaged over positive and negative interventions. The concept values are normalized such that the maximum value in training data is 1 and the minimum value is 0. We can see our CB-pLM models clearly outperform baselines for almost all concepts.

Figure: (a) Average intervention accuracy.
Refer to caption: x4.png

Figure: (a) Aggregate intervention effect.
Refer to caption: x6.png

###### Distribution shift

We investigate how the concept distribution shits along with the naturalness of the protein after intervention. Ideally, we would want the concept distribution to shift in the direction of the intervention while maintaining the naturalness of the protein; we plot $\Delta$ Concept/Naturalness distribution (the difference between the original sample and the generated sample for the concept value on $x$-axis and naturalness on $y$-axis).

Figure 12 and Figure 13 shows distribution shifts between different types of 24M and 150M models respectively, with both positive and negative interventions, CB-pLM effectively shifts the concept distributions in the right direction for all concepts while preserving the naturalness of the proteins, unlike other variations of conditional language models. Figure 14 shows distribution shifts between CB-pLM with different sizes. Note that although increasing model size improved the model’s perplexity we did not observe any improvement in terms of control.

Figure: (a) Positive interventions.
Refer to caption: extracted/5988250/figures/dist_single_pos_24M_2.jpg

Figure: (a) Positive interventions.
Refer to caption: extracted/5988250/figures/dist_single_150M_positive.png

Figure: (a) Positive interventions.
Refer to caption: extracted/5988250/figures/positive_size_diff.png

###### Iterative intervention

In this experiment, we examine different models’ ability to shift a concept’s distribution iteratively, i.e.; we vary $n$ in Algorithm 1 from $1\rightarrow 3$. Figures 15 and 16 shows concept distributions for three iterations for both positive and negative interventions. CB-pLM can iteratively shift the concept value with increasing effects while other models fail to do this.

Figure: Figure 15: Iterative positive intervention.
Refer to caption: extracted/5988250/figures/iterative_pos.png

Figure: Figure 16: Iterative negative intervention.
Refer to caption: extracted/5988250/figures/iterative_neg.png

###### Concept intervention correlation

In this experiment, we mask the most influential tokens for each concept as previously described and intervene on the value of that concept. We then calculate the value of all 14 concepts for the generated sequence, and see how intervening on one concept effects all the other concepts. The results are reported in Figures 17(b), 17(c) and 17(d) for different models. In this figure we report a version of accuracy, where the $j$-th element of the $i$-th row represents the (fraction of times concept $j$ changed in the same direction as the intervention when intervening on the $i$-th concept) - (fraction of times concept $j$ changed in the opposite direction of the intervention when intervening on the $i$-th concept). This can be seen as a binary version of correlation. We also compare this with the actual correlation between these concept values on existing proteins in our dataset shown in Figure 17(a). We can see the correlation matrix of interventions on our CBM has many similarities to the correlations between the concept in the real world, especially on pairs with high correlation, reflecting that our concept bottleneck model has picked up on these natural correlations. In contrast, the other models have a quite different correlation pattern.

Figure: (a) Correlation between concept values in our dataset.
Refer to caption: extracted/5988250/figures/gt_correlation_matrix.png

##### C.1.2 Multi-concept Intervention

We show the effects of sequentially intervening on two concepts in Figure 18 with different 150M parameter models below.

Figure: Figure 18: Multi-concept interventions.
Refer to caption: extracted/5988250/figures/multiconcept_150M.png

#### C.2 Protein Design - Siltuximab

###### Baselines

- •
LaMBO-2 (Gruver et al., ) is diffusion optimized sampling, a guidance method for discrete diffusion models that follows gradients in the hidden states of the denoising network. LaMBO-2 was trained on SWISS-PROT ($\sim 2.6$M training samples training) using GRAVY as the optimization objective for explicit guidance.
- •
PropEn (Tagasovska et al., ) is an encoder-decoder architecture that is implicitly trained to optimize a property of interest. PropEn operates by matching training samples so that each training sample is paired with one that has better properties; due to this matching, it is best suited for a low data regime. We train PropEn on a curated antibody dataset of $\sim 1000$ training samples that are close to Siltuximab.
- •
Discrete Walk-Jump Sampling (WJS) (Frey et al., ) is a discrete generative model that learns a smoothed energy function, then samples from the smoothed data manifold with Langevin Markov chain Monte Carlo (MCMC), and projects back to the true data manifold with one-step denoising. WJS was trained on paired observed antibody space (OAS) dataset Olsen et al. (), which is $\sim 120$K paired antibody sequences as done in the original paper.

#### C.3 Additional Control experiments

##### C.3.1 Species Interpolation

Prior work in mechanistic interpretability (Bricken et al., ) demonstrated that sparse autoencoders (SAEs) could uncover interpretable monosemantic features. Recently, Templeton et al. () applied SAEs to Claude 3 Sonnet’s activations, retrieving $\sim 34$M features. They used the autointerpretability approach (Bills et al., ) to assign meanings to these features. One notable feature identified was the Golden Gate Bridge. By clamping this feature to $10\times$ its maximum activation, they influenced the model’s behavior, causing Claude to behave in an out-of-distribution manner, going as far as self-identifying as the Golden Gate Bridge. Here, we show how the concept bottleneck language model demonstrates the same capability, although for CB-LMs there is no need to train an additional SAE and do a post-hoc search for the meaning of different concepts, since here we know what each concept encodes by design.

To demonstrate, we test the CB-pLM’s ability to interpolate between species; we altered the behavior of the CB-pLM by setting a single species concept in the original CB-pLM to 10 times its maximum activation value. We then tested whether the model could generate a sequence similar to that species cluster, regardless of the protein sequences inputted into the model. We conducted this test for two species: (a) rice (species Oryza sativa - referred to as “Rice-CB-pLM”) and (b) human (species Homo sapiens - referred to as “Human-CB-pLM”). We trained a species classifier by finetuning the ESM2(Lin et al., ) model to classify different species (humans, yeast, rice, and bacteria). The model accurately clustered different species together, with an overall accuracy of 98%, as shown in Figure 19(a). We then passed non-rice protein sequences into Rice-CB-pLM and classified the output; most newly generated sequences were now identified as rice (Figure 19(b)). Similarly, passing non-human sequences into Human-CB-pLM resulted in most newly generated sequences being classified as human (Figure 19(c)); note that changing non-human protein sequences into human sequences is a fundamental problem in protein engineering known as “humanization”.

This illustrative example showcases the ability of the concept bottleneck language model to replicate the capabilities of SAEs and autointerpretability with a substantially simpler approach.

Figure: (a) ESM2 specie clusters.
Refer to caption: extracted/5988250/figures/species_umap_before.png

##### C.3.2 Generalization to new concept combination

So far we have shown that our model can successfully control generation of in distribution sequences. However, for many real-world use cases, it is important to be able to generate previously unseen concepts or combinations of concepts.

To test this, we train a model with our architecture on a simpler task based on color-MNIST. To make this task similar to masked language modeling, we train our models by masking 75% of the input, and then train them to reconstruct the remaining input using a masked autoenecoder vision transformer (MAE-ViT) (He et al., ).
Following MAE-ViT, we divide the input into a sequence of 7x7 patches, making our task a masked sequence-to-sequence generative task, just like masked language modeling. Note that unlike our main results, we use a nonlinear decoder to accommodate for pixel-level outputs. For the concepts in this experiment, we use a one-hot encoding of the number represented by the image, together with an RGB encoding of the color of the input digit. It is important to use RGB embeddings instead of one-hot representation of color, as it allows to encode any color in just 3 dimensions, which is useful for testing generalization capabilities.
To test the generalization ability of our CBM architecture, we trained a generative CBM as discussed above on ColorMNIST with 7 discrete training colors: {magenta, grey, cyan, white, red, green, yellow}. On test time, we then intervened on the CBM color outputs to make the model create an unseen color. As shown in Figure 20, our method can accurately regenerate the the input for all 3 unseen colors we tested. These examples were not cherry picked, and our color intervention reliably works on all inputs.

Figure: Figure 20: Results of our generative CBM on color MNIST. We can accurately generate digits in colors completely unseen during training. This highlights that CBMs are able to generate novel outputs.
Refer to caption: extracted/5988250/figures/mnist_ood.png

### Appendix D Interpretability and Debugging

#### D.1 Feature attribution and coordinate selection

The goal of our control experiments is to make small modifications to protein sequences that increase or decrease the value of a certain concept. While during training we randomly select which tokens to mask and predict, this is not an effective strategy for interventions. Instead we wish to identify which amino acids are increasing/decreasing a certain concept the most, and only mask and predict those tokens.

For example if we wish to increase Aromaticity of a protein, we wish to identify the amino acids which contribute the most to decreasing the Aromaticity, and intervene on those tokens. The most straightforward and principled method to measure this is occlusion, where we simply remove a each input token and replace it with a reference value one at a time, and see how our predicted concept changes for each. Since our model was trained with masking, the natural reference value is the mask token, so when removing a certain input we simply replace it with the mask token. Let $T\in\mathbb{R}^{v\times d}$ be the learned token embeddings where $v$ is the vocabulary size and $d$ is transformer hidden dimension. Then let $f_{T}(x)=[T_{x_{0}},T_{x_{1}},\ldots,T_{x_{s}}]$. The occlusion attribution for token $t$ of input $x$ on concept $i$ is then:

OCCLUSION:

$$ $A(x,i,t)=\hat{c}_{i}(f_{T}(x))-\hat{c}_{i}(f_{T}(x_{x_{t}\leftarrow\texttt{[} MASK]}))$ (2) $$

Where $x_{x_{t}\leftarrow\texttt{[}MASK]}$ represents $x$ with the $t$-th token replaced by the mask token, and $\hat{c}_{i}(e(x))$ is the concept value predicted by our CBM model for concept $i$. concept prediction for concept $i$ on input $x$. For positive interventions, we typically then mask the 5% of the tokens (excluding padding) with the smallest attribution value(token being present reduces concept value), while for negative interventions we mask 5% of tokens with the largest attribution value(token being present increases concept value).

While this is a good method, it requires a separate forward pass for every individual token in the input, making it slow to run in large scale. Instead, as is common with feature attributions, we use a gradient based approximation of the occlusion value for up to 512 times speedup compared to occlusion. However, since our inputs are discrete, we cannot directly calculate gradients with respect to them. Instead, we calculate gradients with respect to the learned token embeddings of the model, and sum over these. Then

GRADIENT x (INPUT - MASK):

$$ $A(x,i,t)=\nabla_{f_{T}(x)_{t}}\hat{c}_{i}(f_{T}(x))\cdot(T_{x_{t}}-T_{\texttt{ [}MASK]})$ (3) $$

Where $f_{T}(x)_{t}$ is the $t$-th element of $f_{T}(x)$. This method is essentially a first order Taylor Approximation of the occlusion metric defined above. Alternatively this method could be seen as a version of Integrated Gradients Sundararajan et al. () with only one gradient step and the mask token as the reference value. To see whether this approximation is sufficiently accurate to identify good amino acids to intervene on, we compare it’s performance against occlusion itself in table 3. We also compared against some common feature attribution baselines described below:

GRADIENT x INPUT: Shrikumar et al. ()

$$ $A(x,i,t)=\nabla_{f_{T}(x)_{t}}\hat{c}_{i}(f_{T}(x))\cdot T_{x_{t}}$ (4) $$

GRADIENT: Simonyan et al. ()

$$ $A(x,i,t)=||\nabla_{f_{T}(x)_{t}}\hat{c}_{i}(f_{T}(x))||_{1}$ (5) $$

We compared all these methods, as well as a baseline of randomly choosing tokens to mask. We took a random sample of 50 proteins from the validation dataset, and intervened to increase/decrease each of the 14 biopython concepts one at a time, reporting the average change across concepts when masking 5% of the inputs in table 3. We can see choosing tokens to mask based on occlusion attribution performs the best as expected, however our first order approximation Gradient x (Input - Mask) also performed very well, with average intervention effect only 5% below using occlusion, while being around $100\times$ faster to calculate. On the other hand, simply looking at the magnitude of the gradient performed poorly, even worse than random selection, likely because it doesn’t take into account whether said input is increasing or decreasing the concept value. Overall we can also see that choosing which tokens to mask is very important for controlled generation, and that with good attribution our interventions are more than twice as effective as randomly selecting tokens to intervene on.

**Table 3: Average concept change after interventions using different feature attribution methods on our 24M-CBM model. The average column reports the average change in the desired direction.**
| Method | Positive | Negative | Average |
| --- | --- | --- | --- |
| Random | 0.0312 | -0.0125 | 0.0218 |
| Gradient | 0.0173 | -0.0201 | 0.0187 |
| Gradient x Input | 0.0358 | -0.0269 | 0.0313 |
| Gradient x (Input - Mask) | 0.0480 | -0.0390 | 0.0435 |
| Occlusion | 0.0499 | -0.0419 | 0.0459 |

Based on these results, for all other experiments in the paper we use the Gradient x (Input - Mask) method for our CB-pLM and CC-pLM, while the C-pLM uses random attribution as it does not make concept value prediction $\hat{c}$ which is needed for all other attribution methods.

#### D.2 Understanding concept/output relationships

To understand the model’s ability to infer biophysical properties, we visualized the weights from the final layer for each concept in relation to each amino acid. Our findings indicate that the model successfully learns several key biophysical relationships as defined in the BioPython library (Figures 21, 22 and 23).

Figure: Figure 21: Visualization of the final layer weights in CB-pLM 24M.
Refer to caption: extracted/5988250/figures/weight_vis.png

###### Some key insights include the following:

- •
For both charge at pH 6 and pH 7, acidic amino acids (D, E) have negative weights, while basic amino acids (R, K) have positive weights. The difference in charge between pH 6 and pH 7 reflects the biophysical properties of H , as it contains an imidazole side chain with a pKa of 6.0 and contributes more towards positive charge at lower pH (Figures 22(a), 22(b)).
- •
GRAVY, which defines hydropathy based on the Kyte-Doolittle scale (Kyte & Doolittle, ), is consistent with positive weights assigned to A, I, V, F, C, and M amino acids (Figure 22(c) ).
- •
Aromatic amino acids F, Y, and W have high weights for the aromaticity concept (Figure 22(d)).
- •
Average hydrophilicity uses the Hopp-Wood scale, which differs from the Kyte-Doolittle scale in its empirical definitions, assigns higher values to R, D, E, and K (Hopp & Woods ()). Although the scale also assigns lower values to aromatic residues, such as F, W, Y, we find that our model assigns lower weights for other amino acids which could reflect biases in the dataset (Figure 22(e)).
- •
Average surface accessibility is defined by the Emini Surface fractional probability (Cock et al., ), which is a lookup table describing the likelihood that a residue appears on the protein surface. The scale defines a low value especially for C and for other non-polar amino acids, though to a lesser extent (Figure 22(f)).
- •
Helix fraction, sheet structure fraction, and turn structure fraction are related concepts, are defined by fractions of I, L, F, W, Y for helix, A, L, M, E for sheet and N, S, G, P. While the weights do not closely match that of the BioPython definitions, we note that the weights are nonetheless consistent with amino acids that often have special propensity to be part of helical structures, such as A, M, L, E, K (Leiro et al., ) (Figures 23(a), 23(b) and 23(c)).
- •
Isoelectric point is a property closely associated with charge at pH 7, and likewise, we observe low values for D and E, and high values for R and K (Figure 23(d)).
- •
Molar extinction coefficient, which quantifies absorption of a protein at 280 nm, emphasizes the contributions of W. Likewise, the reduced vs oxidized terms of the molar extinction parameter differ only in how the oxidation state of C is considered, which is also visible in our model weights (Figures 23(f), and 23(e)).
- •
Instability index is defined by a look-up table based on the dideptide decomposition of the sequence (Guruprasad et al., ). We observe that the model attributes high values to P, which often introduces kinks in the protein structure and disrupts secondary structures (Figure 23(g)).
- •
We note that the model weights does not always reflect the concept definitions. For example, W is the heaviest amino acid but is assigned a negative value, which likely reflects biases in the dataset. W is typically one of the least common amino acids, and relative frequency may vary by sequence length. However, the model correctly identifies F as contributing to high molecular weight, and A, and G towards contributing to low molecular weight (Figure 23(h)).

Figure: (a) Charge at pH6
Refer to caption: extracted/5988250/figures/charge_at_pH6_bar_chart.png

Figure: (a) Helix Fraction
Refer to caption: extracted/5988250/figures/helix_fraction_bar_chart.png

### Appendix E Extended Related works and Discussion

#### E.1 Related Work

##### Language models

Language models are increasingly widespread in their usage, spanning from natural language processing (NLP) (Achiam et al., ; Touvron et al., ; Team et al., ) to specialized fields such as chemistry (Bran et al., ) and biology (Cui et al., ; Lin et al., ). These models have demonstrated exceptional performance on intricate tasks, facilitating new frontiers in diverse applications (Brown et al., ; Hayes et al., ). However, the lack of concrete and effective interpretability methods limits their use in many high-stakes applications, engendering distrust among domain experts, and raising concerns regarding regulatory compliance, safety, and alignment (Goodman & Flaxman, ; Gabriel, ).

Current interpretability research for large language models (LLMs) focuses on pre-trained models, either by explaining the model’s predictions (Zhao et al., ) or by analyzing the internal circuits of the network through mechanistic interpretability (Räuker et al., ; Bricken et al., ). Recent work in mechanistic interpretability (Bricken et al., ) demonstrated that sparse autoencoders (SAEs) could uncover interpretable monosemantic features. To apply this on large-scale production models, it involves training very large SAEs, retrieving millions of features, and then trying to assign meaning to these features using an autointerpretability approach (Bills et al., ).

Although there have been promising results in controlling large production models like Claude using such SAEs, it remains unclear how this approach will help when controlling for a concept not found in the SAE or when autointerpretability is not feasible, such as in domains like proteins. Concept bottleneck language model (CB-LM) demonstrates similar capabilities (see Appendix C.3.1 for experiments). However, for CB-LMs, there is no need to train an additional SAE and conduct a post-hoc search for the meaning of different concepts, as each concept’s encoding is known by design.

##### Concept Bottleneck Models

Concept Bottleneck Models (CBMs) (Koh et al., ) incorporate interpretability into neural networks by inserting an interpretable “concept bottleneck” layer into the network and mapping the inputs to a set of human-understandable concepts, which are then used to make the final prediction.

Recently, Ismail et al. () demonstrated that a concept layer can be integrated into generative models (CBGMs). Intervening on this layer allows for controllable generation, and examining the activations provides concept-level explanations. Their work focused on image generation and evaluated the effectiveness of this approach in architectures where the entire input is represented as an embedding, such as GANs, diffusion models, and VAEs.

In this paper, we extend this approach to language models, demonstrating that we can achieve global control over the entire input while retaining token-level generation. To make this adaptation, we inserted a concept bottleneck layer into a language model and created a novel concept module. This differs from CBGMs in several key aspects: (a) We focus on sequential data where inputs are represented as tokens. (b) We employ a unique model architecture tailored to language models. (c) We introduce a new intervention procedure specific to our architecture.

A few very recent works have proposed CBMs for text classification tasks Tan et al. (); Sun et al. (), but this is a much easier and less interesting setting than generative language modeling which we address, as the input and output space of generative models is much larger and cannot be adequately modeled using only known concepts, highlighting the need for an unknown part in the embeddings.

Recent work Yuksekgonul et al. (); Oikarinen et al. (); Yang et al. (); Rao et al. () has also proposed methods to train Concept Bottleneck Models for image classification that do not require labeled concept data by leveraging pretrained multimodal models. While interesting, this approach reduces reliability of concept predictions, and is not feasible in protein setting as no good pretrained general concept predicting models are available.

Similar to ours, some previous work has proposed including an unknown part in the embeddings to improve task performance in the image classification settings Yuksekgonul et al. (); Sawada & Nakamura (), but they employ quite different methods to achieve this in a different setting.

##### Protein Language Models

Protein language models are widely used in biological machine-learning research, trained on extensive protein sequences spanning the evolutionary tree of life. Although pLMs have numerous critical applications in healthcare and drug discovery (Hie et al., ), they currently lack interpretability. Their performance is often explained through intuitive arguments about compression and the learning of co-evolutionary patterns (Rives et al., ; Lin et al., ; Hayes et al., ), but it remains unclear whether these models truly capture meaningful protein concepts.

Historically, pLM architectures, loss functions, and training setups have closely mirrored the original masked language model setup (Devlin et al., ), with the primary differences being the vocabulary and dataset. Our goal is to develop a pLM that offers controllability, interpretability, and debuggability, making it more reliable for design applications. While some protein language models support conditional generation by concatenating different protein functions and properties to the inputs (Shuai et al., ; Madani et al., ; Hayes et al., ), in contrast to our work, they do not provide any mechanisms for interpretability, debuggability, or offer any insights on what a model has learned.

#### E.2 Discussion

##### CB-LM for other domains and architectures

In this paper, we applied our proposed architecture to masked language models for proteins. This versatile architecture can be adapted to any domain, such as NLP, by simply changing the training dataset, vocabulary, and concepts while keeping the model itself unchanged. Additionally, it can be easily adapted to other types of language models, such as autoregressive models, as demonstrated in Figure 24. For autoregressive models, the embedding output from the transformer, typically used to predict the next token, is passed to the CB-layer to extract the known concept embedding and to an orthogonality network to extract the unknown concept embedding. These two embeddings are then concatenated and used to predict the next token.

Figure: Figure 24: Inserting concept bottleneck into an autoregressive language model.
Refer to caption: extracted/5988250/figures/autoregressive.jpg

##### Adding new concepts to a trained CB-LM

An intuitive question that might arise is how we can control for new concepts. Specifically, if we are given a dataset with a new concept, can we fine-tune CB-LM to incorporate it? This can be achieved by creating a new CB-LM with an additional concept and copying the weights from the original model. In this setup, the only randomly initialized weights are for the new concept embedding $e_{k+1}$, the concept classifier head for the new concept, and the linear decoder. The new model can then be trained to match the old model’s concept output for existing concepts while aligning with the ground truth for the newly added concept.

##### Generalization to new concept combination

In many real-world applications, we may need to generate sequences with previously unseen concepts or novel combinations of concepts not present in the training dataset. CB-LM provides this capability, as long as the new concepts can be expressed as functions of the existing ones. We demonstrate this capability in a toy task detailed in Appendix C.3.2, where we train a masked autoencoder vision transformer (MAE-ViT) (He et al., ) with our architecture on a simple color-MNIST generation task. Our results show that the architecture can successfully generate blue-colored MNIST digits, despite the model never having encountered the color blue in the training dataset.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: * Line 245: it appears the time is denoted as a typo 
* Table 1: how was the validation set chosen, is there overlap between the validation set and LBSTER/ESM2 train sets? 
* Section 4.1 
    * Why was the C-pLM baseline tested with random substitutions instead of feature attribution. 
    * Model size does not appear to improve intervention accuracy or change, significantly. 
* Section 4.2
    * What is the motivation for using WSJ here. Why not use ESM2 or an inverse-design approach like ProteinMPNN or ESM-IF, for sequence redesign in the hydrophobic region? 
    * It seems useful to include a non deep-learning baseline here that at minimum randomly re-samples residues in the hydrophobic patch with hydrophilic residues. 
    * Is the hydrophobic patch critical for the function of the antibody? How do you know if the redesigned region has improved or reduced function. 
* Section 5: 
     * Is it not expected that the model should correctly learn or not learn the concepts provided during training? 
     * It is also clear in what application the model debugging of a CB-pLM might be relevant, based on the given example.
soundness: 3
strengths: I believe this CB-pLM is the first model which applies concept bottleneck models to protein design, the authors show with strong evidence that CB-pLM effectively shifts concept distributions better than existing approaches such as C-pLMs and CC-pLMs. The paper is well written, and does a good job at showcasing the in-silico results for conditioning.
summary: The authors apply concept bottleneck language models to protein design. They propose this model has three benefits; control, interpretability, and facilitated debugging. To support this, they demonstrate evidence that the model can optimize for single or double property control, better than existing conditioning approaches. They showcase this by redesigning the hydrophobic patch of an antibody, showing that it responds to the in-silico metrics as expected. Finally, they correlate model weights with trained concepts, and propose these correlations can help model debugging efforts
weaknesses: A limitation of this approach is that biophysical concepts need to be explicitly defined for controlling design. The argument that a CB-pLM is more interpretable because of the training approach is less convincing, as the results appear to be evidence that the model learns the provided concepts correctly, and doesn’t expand to a new interpretation of known biological properties. Concepts used in the model are also straightforward to calculate and interpret the outputs of with bioinformatics tools, alone. The model does not provide additional interpretability on proteins that otherwise exists.  Similarly, while the redesign example appears to currently be tested in a lab, it is not possible to asses if the model improves the solubility and function of  Siltuximab, based on the data provided.

---

confidence: 2
contribution: 2
initial_rating: 5
presentation: 3
questions: Q1 How does compute / training vary across C-pLM, CC-pLM, and CB-pLM. Is total compute/wallclock time controlled (e.g. 24 hr, as the cited cramming paper) or perhaps early stopping on validation is used?

Q2 It is interesting that there is negative scaling happening in figures 10 and 11, with sometimes large decreases in performance as the model becomes larger. Out of curiosity, do authors have any speculations about why this happens?

Q3 In lines 142-144, “This enables us to answer [… q1…], or [q2] counterfactual questions like if we want to decrease hydrophobicity, which amino acid should we replace “E” with?”. I understand q1 is an interpretability question and linear decoder does help with that. For q2 however, I am not sure I fully understood. Is it accurate to rephrase it as “if we intervened on the concept strength of hydrophobicity, which amino acid would the decoder output?” My confusion is on the verb “enables” – I believe the linear layer gives you a more interpretable answer to q2 (say compared to a nonlinear decoder), but you would still be able to answer q2 with a nonlinear decoder,  by intervening on the concept value $\hat{c}_i$ and then running a forward pass of the decoder.

Q4 What’s the spread along the x-axis on figure 4b? Is it naturalness? It might be helpful to also plot the reference lines for no-change (for both naturalness and intervened property) as you do in figure 4a.
soundness: 2
strengths: * Adapting concept-bottleneck generative models to masked language models is an interesting and novel idea.
* Using integrated gradients to find tokens to resample to is an interesting and novel idea.
* Writing is generally easy to follow.
* This work focuses on improving controllability, interpretability, and debuggability of language models (and protein sequence in particular as an example) which is an important topic with large potential impact.
summary: This paper proposes a Transformer Masked Language Model (MLM) variant of concept-bottleneck generative models (CBGM) with three goals in mind: controllable generation, interpretability, and debuggability. This paper adapts the training objectives of CBGM to the masked language model architecture, and proposes a control method specifically for their CB MLM. Evaluation was conducted on the domain of protein sequence modeling, for controllability, interpretability and debuggability. The proposed model leads to improvements over non concept-bottlenecked baselines on controllability. The proposed model demonstrates interpretability and debuggability via inspection of the weights of the linear decoder head. 

Comment: I have experience with language modeling but not with protein domain, so my review will focus more on the method part of the paper (as recommended by AC). Please take this into account when considering my review and recommendation.
weaknesses: Method

(considering a broader application of the proposed method to general language/sequence modeling, as authors mentioned in the abstract) 

M.1 The model proposed is only a conditional generative model given a partial sequence rather than a full generative model of an entire sequence. Classical masked language models (like BERT) do not model the joint distribution over sequences [1]. There are ways to derive joints from MLMs by making their conditionals consistent with one another (see, e.g. [1] and the related works cited there), or by modifying the training objective so that the conditionals are encouraged to be consistent during training (e.g. [2] [3]). Other protein models based on MLM (e.g. [3], which the authors cited) have done the latter kind of training modifications to enable sampling full sequences. I'm not sure how this limitation matters for protein modeling, could authors give more discussions about this? 

M.2 Naturalness of the controllably generated sequence is not considered during intervention, but only used as a measurement afterwards. See point below for consequence.

M.3 The proposed controllability technique only proposes changes local to a particular known sequence x, by changing n% (5 in experiments) of x at a time. Again, I cannot comment on the significance of this limitation for protein modeling. For other modalities like images and language, it can be difficult to change sequence-level concepts via chaining a sequence of such local changes. Even if it is possible, since naturalness is not considered as a factor during intervention but only afterwards, applying a sequence of interventions could lead to going outside of the training distribution (e.g. the interventions produced in Figure 20 goes outside of the color-MNIST training distribution, where all digits are uniformly colored). Empirically though, (Fig 4b 15 16) the proposed method seems to outperform baselines so maybe this is less of an issue in practice, for the concepts in this paper and/or for the protein domain.

Experiments

E.1 The model is trained with 25% masking, but is used to generate with 5% masking in the experiments. This could make the model worse because it’s a different input distribution. Is there a particular justification for masking 5%?

E.2 When evaluating controllability against baselines (and more specifically CC-pLM, since it's closer to CB-pLM), it’s not clear whether the gain comes from intervening at better positions (due to applying integrated gradients to a different architecture), better sensitivity of the model to interventions (due to the proposed concept bottleneck), or a combination, making it hard to evaluate the limitations / benefits of the architecture versus the intervention technique.

[1] Hennigen and Kim, 2023. Deriving Language Models from Masked Language Models.

[2] Germain et al., 2015. MADE: Masked Autoencoder for Distribution Estimation.

[3] Hayes et al., 2024. Simulating 500 million years of evolution with a language model.

---

confidence: 2
contribution: 2
initial_rating: 3
presentation: 1
questions: 1. Line 76, what are “human understandable concepts”? This is not very well defined in the paper and it is left vague. Similarly, what does debuggable model mean? Even if the proposed model leverages a bottleneck to force some representations to align to a set of pre-specified “concepts” the whole model before the bottleneck is a complete black-box. Does debuggable mean that the user can inspect the black box before the bottleneck too? 
2. In the abstract and intro it is mentioned that the proposed model has been scaled to 3B parameters, however there are no empirical results showcasing models larger than 150M (most of the results are reported on the small 24M model). Furthermore, in Table 1 the comparison with other State of The Art models suggests that the proposed model regresses in perplexity over similarly sized non-interpretable models. Can the authors provide an apples to apples comparison with other public models at scales larger than 150M?
3. Is the current method supposed to scale to a large number of protein sequences? Will it generalize in the case of mislabelled or unlabelled samples? Given annotation cost, a model capable of leveraging partially annotated/unlabeled samples is more likely to scale.
4. What is the maximum context length (length of the protein sequence) that the model supports?
5. It is not clear how the concepts in the bottleneck layer should be defined and why the model needs to learn to associate specific concepts to “disentangled” representations in one-to-one relation with the concepts bank. From the paper it appears that the concepts are annotated and hence the “factor of variations” are known at training time. This could become problematic since scaling on annotated data is not always possible. Does the proposed method allow one to not directly observe the factors of variations and still learn some minimal and sufficient concept bank, similarly to the method proposed in [1]? Furthermore, when annotations are used, are the learned concepts compositional? 
6. Experiments in section 4.1.1 and 4.1.2 have been conducted on ablations of the proposed method. Why is it not compared with other State of The Art methods of similar sizes?


[1] M. Fumero et al., “Leveraging sparse and shared feature activations for disentangled representation learning”, NeurIPs 2023.

Minor:
- Line 121, typo.
- Line 125, labels in match with Figure 2 even if they are not related.
- Line 226 repeated.
- Line 245, typo
soundness: 2
strengths: Developing and testing new models for protein modeling is an important problem. In particular, having methods that can be interpretable and that can allow better control for drug design is a high impact project.
summary: The paper introduces a novel protein language model designed to provide predictions that can be interpreted by human experts without compromising performance. The proposed model, dubbed Concept Bottleneck Protein Language Model (CB-pLM), empirically proves to be easier to control and interpret over previously proposed State of the Art Protein Language Models.
weaknesses: The paper contains many typos, it is meandering and unfocused, and generally is very hard to read (with many repetitions and vaguely defined concepts).  

For example, what does “debuggable model“ mean? This is stated as one of the main contributions of the proposed architecture, hence the reader expects it to be clearly defined and thoroughly assessed through empirical validation. In the current version of the manuscript this is not properly defined and hence not convincingly validated with experiments. 

What are the novel contributions specific to the Protein Language domain? Can the authors separate the specific novelty introduced for their protein language model? What are the main differences with the previously proposed concept Bottleneck Generative Models? If the main contribution of the paper is to test whether CB-LMs work on the Protein domain then it would be good to devote more space in the paper to the empirical validation with a thorough comparison with other State of the Art models at different scales.

The discussion section briefly mentions some properties of the model that have not been tested in the main empirical section of the paper. It would be good either move the strong empirical evidence supporting concept composition and generalization to novel concepts (added with fine-tuning) in the main or, if results are not convincing enough, remove the claim. Similarly, in the conclusions (and abstract) it is mentioned that the current model has been tested up to the 3B scale and compared with other public models. However I did not manage to find evidence of experiments on scales larger than 150M (see for example Table 1 in which the proposed method does not perform comparably with others).

---

confidence: 3
contribution: 4
initial_rating: 8
presentation: 3
questions: There are some minor typos e.g. "00:01 01:15" in Section 3.1, as well as a missing line break before \textbf{Scaling} in the same subsection.
soundness: 3
strengths: The controllability experiments are fairly convincing, though the stated application on debugging is fairly cursory and could be improved (the intervention to test debugging is fairly blunt, and also the results are just hand inspection and thus not reproducible). However, for me the controllability experiments are enough to be a solid contribution.

The perplexity results were moderately convincing, though more discussion of comparisons to LBSTER and ESM2 would be useful -- it was not clear to me for instance if your method was given extra information (the tags) that LBSTER / ESM2 did not have. It was also strange that the largest model had worse perplexity.

The writing is overall good and easy to follow, but I felt that the level of detail in some experiment descriptions was not enough to be fully reproducible from the text.
summary: This work applies concept bottleneck models to protein language models. The approach uses fairly standard ideas in concept bottleneck models, but they are well-executed and lead to novel applications in e.g. controllability.
weaknesses: See above (strengths and weaknesses are merged)

---

confidence: 3
contribution: 4
initial_rating: 8
presentation: 4
questions: If even a small wet lab experiment is possible, I would be super curious to see the result.
soundness: 4
strengths: Human-interpretable creative control is a super key in generative AI broadly, but especially in protein design.  As such, the proposed approach based on the concept bottleneck is quite compelling.  The computational experiments provide evidence that the proposed architecture is effective.  The paper is written in a clear way.
summary: This paper takes the approach called concept bottleneck language modeling and it applies it in the protein language model context.  This provides creative control in terms of human-interpretable biological concepts, as well as the ability to debug models.  The proposed approach appears to work better than other generic approaches to controllable generation such as ProGen/CTRL, as demonstrated in experimental comparison with other models trained by the researchers themselves.  A case study with redesign of siltuximab shows the ability to quite precisely control certain properties, though no wet lab experiments are shown to verify.
weaknesses: No wet lab experiments were done to verify results.  This would strengthen the paper considerably.

Comparisons to other architectural approaches are based on the authors' own implementations, rather than by external research groups that may be motivated to optimize the other architectures more.  Notwithstanding, the performance differences are quite considerable.

No discussion on biosecurity risks is given.

Meta-review: This paper introduces "Concept Bottleneck Protein Language Models" (CB-pLM), a novel approach that integrates a concept bottleneck layer into protein language models. This architecture aims to enhance control, interpretability, and debugging capabilities without compromising performance. The authors demonstrate that CB-pLM can effectively manipulate protein properties, offering significant improvements in controllability and interpretability over existing models, and that they can be scaled to the billion-parameter-sized regime.

The main strengths of this paper are:
* Innovative Architecture: The integration of a concept bottleneck layer into protein language models is novel and offers significant improvements in control and interpretability.
* Empirical Validation of Controllability: The paper provides strong empirical evidence that CB-pLM can manipulate protein properties more effectively than existing models.
* Scalability: The model is shown to scale up to 3 billion parameters, demonstrating its robustness and potential for large-scale applications.
* Practical Applications: The ability to control and interpret protein properties has significant implications for drug discovery and other bio applications.

The main weaknesses, in turn, are:
* Need for concept specification: a typical downside of concept bottleneck layers that applies to CB-pLM too is that the "concepts" need to be explicitly specified. This can be done either with prior knowledge or, as suggested by the authors, using a post-hoc interpretability method.
* Limited/preliminary experimental results for "debugging": multiple reviewers raised concerns about the replicability/significance of this part of the experimental validation.
* Limited applicability and generalizability of the approach beyond protein sequences.  

All things considered, this paper provides a modest but compelling contribution on an important open research problem, namely interpretability in generative AI for biology, which is of high potential impact. Hence, I recommend acceptance.

Decision: Accept (Poster)