Concept Bottleneck Language Models For Protein Design

Thank you for your positive feedback, we are glad you found our approach compelling. Below, we address your concerns point-by-point.

• Re: No wet lab experiments were done to verify results. This would strengthen the paper considerably.
  Please find our response in the general response R2.

• Re: Comparisons to other architectural approaches are based on the authors' own implementations
  We have implemented baselines ourselves because we are unaware of open-source models that offer control for the same concepts we are testing.

• Re: No discussion on biosecurity risks is given.
  Thank you for pointing this out; we have included biosecurity risks in the discussion section in the updated manuscript.

Thank you for your positive feedback, we are glad you found our work well executed and having novel applications. Below, we address concerns point-by-point.

• Re: Perplexity
  Our model, CB-pLM, does not require tags during inference. Reviewer ey43 pointed out that our validation dataset might have been included in the training dataset for LBSTER/ESM2. To address this, we repeated the experiments using the Mason dataset [1] and updated the manuscript with the new values. We have also clarified which models require tags and which do not. Regarding the perplexity results, we did not observe significant performance gains beyond the 650M model, noting that our 650M model already outperforms ESM2 3B. Please find our reasoning for this observation in the general response R1.

• Re: Typos
  Thank you for catching this. We have corrected this in the updated manuscript.

References

[1] Mason, D. M., et al. "Optimization of therapeutic antibodies by predicting antigen specificity from antibody sequence via deep learning."

Re: Questions

• Q1: What are “human understandable concepts”? What does debuggable model mean?
  A human understandable concept is an input feature or transformation of that input feature that a human (typically a domain expert) can attach meaning to. For debuggable model, please find our response in the general response R3.

• Q2: Can the authors provide an apples to apples comparison with other public models at scales larger than 150M?
  Please see W4.

• Q3: Is the current method supposed to scale to a large number of protein sequences?
  Since CB-pLM can scale to a large number of parameters, we do not foresee any issues with supporting more training data. Complete annotations for all protein sequences are not required; a subset of annotated data is often sufficient to learn the desired concepts. In fact, during training, we combined two datasets, UniRef50 and Swissport; the UniRef50 dataset lacks concept annotations found in the Swissport dataset, resulting in less than 3% of the data being fully annotated. Details on how we handle missing concepts can be found in Appendix A.1.

• Q4: Is the current method supposed to scale to a large number of protein sequences?
  The maximum context length is 512, similar to [1]. Note that a typical antibody is 249 amino acids. See model details in Appendix B.2.

• Q5: It is not clear how the concepts in the bottleneck layer should be defined.
  Very good questions! We’ll address these in parts.

  • Defining concepts: The focus of our work is scientific applications where there is often ample domain expertise and prior knowledge about the underlying factors of variation. However, the exact functional relationship between such factors and the output might be unknown—hence why machine learning is being applied. One such setting is drug discovery where certain key biophysical properties are known, and have been experimentally confirmed, to affect a protein’s function. Given such knowledge, domain experts often want to assess the extent to which a generative model’s representations encode these properties. We’ll address your point about unknown concepts soon, but we contend that, in scientific applications, defining concepts amounts to encoding these factors as concepts in the model. Across various applications, we appeal to domain expertise to define concepts.
  • Why should the model learn to associate specific concepts to “disentangled" representations: To start, one of our key constraints is: we enforce orthogonality between the known and unknown concepts, but not ‘between’/’among’ known concepts. As we previously discussed, we rely on a domain expert to specify known concepts. Consequently, we are relying on expert knowledge prior to prune the concept set to a minimally sufficient list that consists of factors that can be expected to be disentangled. Further, our proposal includes a loss function that requires the unknown concepts to be orthogonal to the unknown ones. A key point here is that if the list of known concepts is not disentangled, then we cannot expect the representations learned by the CB layer to also be disentangled. The key constraint we enforce is orthogonality between known concept and unknown concept sets. We expect future work to address limitations of our current proposal.
  • Scaling Concept Annotation: Complete annotations for the entire dataset are not required; a subset of annotated data is often sufficient to learn the desired concepts. Details on how we handle missing concepts can be found in Appendix A.1.
  • Unsupervised Concept Discovery: Without any concept labels, domain specific assumptions, and/or constraints to the learning process, automatically learning the factors of variation—unsupervised—from data alone is not identifiable [2]. Another key challenge is that without supervision or constraints, the unsupervised concepts that a model learns might be difficult to interpret for humans and domain experts. Despite these challenges, unsupervised concept discovery with concept bottleneck layers is feasible, and can be directly applied here as well.
  • Compositionality: We do not explicitly address compositionality beyond assumptions that go into the domain expert’s concept selection process. Our proposal can be extended in straightforward ways to incorporate concept compositionality [3].

  We thank the reviewer for these important questions, and hope our answers help contextualize the scope of our contributions.

• Q6: Experiments in section 4.1.1 and 4.1.2 are not compared with other State of The Art methods of similar sizes?
  CB-pLM offers novel capabilities not available in current open source pLMs, i.e., SOTA pLMs do not offer control over the concepts we are testing. Please let us know if you have any particular models in mind.

Thank you for your review. Below, we address the weaknesses and respond to your questions point-by-point.

Re: Weaknesses

• W1: Paper contains many typos, it is meandering and unfocused
  We have addressed the specific points in our revised manuscript which is hopefully easier to follow. If you have additional suggestions on parts you still find unclear please let us know.

• W2: What does “debuggable model“ mean?
  Please find our response in the general response R3.

• W2: What are the novel contributions?
  Our work introduces the first concept-bottleneck language model. The original concept bottleneck generative model (CBGM) paper focused on diffusion models, VAEs, and GANs, which have inherently different architectures and properties compared to language models. Adapting CBGMs to language models is not straightforward. The main difference is that in language models, generation is done tokenwise while concepts are defined over the entire sequence, while for diffusion, VAE, and GANs, the entire output is generated at once. To adapt concept bottleneck to language models, we have made the following changes:

  • Architecture:
    • In CB-LMs, the concepts are extracted from the CLS token; in CBGMs, the concepts are extracted from the entire representation.
    • In CB-LM, we add an orthogonality network to remove concept information for every token. This was not done in CBGMs since it was not needed.
    • CBGMs use a CEM for the CB-layer i.e., for each concept, there is a network that predicts a concept embedding from the input representation; this did not scale well for LLMs. So, in CB-LMs, for each concept, we learned a concept representation (this representation is learned but is not dependent on the input as in CEMs); in CB-LMs, the concept embedding is the concept representation multiplied by the concept value as described in section 2.1. We find this significantly improves control and reduces the number of added parameters.
  • Training:
    • We added noise regularization to the input embedding to help with faithful attribution as this is a requirement for successful interventions in our setup and again not in CBGMs.
  • Interventions:
    • We introduced an intervention scheme specific to language models, again this is not present in CBGMs, the intervention scheme is described in Section 2.3.

  Overall, the novelty of our work is not just in the domain but rather the architecture, training scheme, and intervention scheme as well as the domain. Regarding empirical validation, we compared our model against state-of-the-art open-source protein language models in terms of perplexity. However, our model offers unique control, interpretability, and debuggability capabilities that current protein language models do not, making direct benchmarking in these aspects impossible.

• W3: The discussion section briefly mentions some properties of the model that have not been tested in the main empirical section of the paper.
  Thank you for pointing this out, we have updated this section in the manuscript to be titled "discussion and future work” and rewrote the section.

• W4: Scale experiments
  The results for the 650M and 3B models are in Appendix C. The reason they are in the appendix and not in the main text is that (a) we do not have 650M and 3B baseline models to compare with, and it will be too expensive and non-environmental-friendly to train baselines of this size. (b) we did not observe much gains in performance with scaling (Please find our response in the general response R1).
  We trained these size models in the first place to show this architecture can, in fact, be scaled up easily. We have repeated the perplexity experiments to address concerns raised by ey43 and added 650M and 3B ESM2 baselines to address your concerns. Our model significantly outperforms ESM2 in different scales.

Re: Typos

Thank you for pointing this out. We have corrected the typos in lines 121 and 245. Regarding line 226, the repetition is intentional to emphasize and clarify the novel capabilities offered by CB-pLM.

References

[1] Frey et al. "Cramming Protein Language Model Training in 24 GPU Hours."

[2] Locatello et al. "Challenging common assumptions in the unsupervised learning of disentangled representations."

[3] Stein et al. "Towards compositionality in concept learning."

Dear Reviewer e4sF,

Given that the end of the discussion period is approaching, we would like to ask if you have any further concerns or questions, particularly as a follow-up to our response? If we have addressed all your concerns, we hope you will consider increasing your score.

Thank you in advance!

Thank you for the detailed review and feedback. We are pleased that you found our adaptation of concept bottleneck models to MLMs and the use of integrated gradients for resampling both interesting and novel. We are happy to clarify any design choices we made, especially those influenced by the specific requirements of the protein domain rather than NLP. Below, we address the weaknesses and respond to your questions point-by-point.

Re: Weaknesses

• M.1: Generating full sequences from MLMs.
  Our approach can be extended to autoregressive models by making similar architectural changes, as described in Figure 27. For generating the full sequence using the current MLM, one can employ the approach described by ESM3 [1]; starting from a fully or partially masked context, tokens can be sampled one at a time or in parallel, in any order, until all positions are fully unmasked. However, many studies have noted that full sequence generation is not as useful or practical for antibody protein design and therapeutics [2,3,4]. Therefore, we do not focus on this aspect in the context of this paper.

• M.2: Naturalness of CB-pLM designs
  To address your concern regarding the naturalness of the proposed designs, we ran additional experiments to verify that our interventions do not compromise the integrity of the sequences. For the Siltuximab experiment, we folded designs from all models using ABodyBuilder2 [5] and calculated various protein surface properties using the Therapeutic Antibody Profiler (TAP) [6], which employs physics-based computations. Ideally, we aim for minimal changes from the original antibody. The results, plotted in Figure 20, show that CB-pLM designs are the least disruptive compared to other methods, where the dotted line in the figure represents the original antibody. The TAP metrics confirm that our CB-pLM designs maintain naturalness more effectively than other approaches.

• M.3: The proposed controllability technique only proposes changes locally.
  In protein design, the goal is generally to change a specific concept with minimal alterations to the sequence, as larger changes can introduce various liabilities or cause manufacturing difficulties. Prior work [2,3,4,7] has shown that significant changes to a protein's properties can be achieved with very few edits. Moreover, [8] has demonstrated that the functional effects of changes to an antibody sequence are not only non-additive but can exhibit state-dependent higher-order interactions, a phenomenon known as epistasis [9]. Overall, it is well established in the literature that even single-point mutations can have drastic effects on sequence-level concepts. Therefore, in our evaluation, we focused on making localized changes, which are well-suited for protein modeling and help maintain the naturalness of the sequences while achieving the desired modifications.
  For the color MNIST, the experiment's goal is generating concepts outside the training distribution, showing the intervention generalizes successfully beyond the training distribution. This is unrelated to the masking percentage since, in this experiment, the percentage of masking in training and intervention is the same.

• E.1: The model is trained with 25% masking but is used to generate with 5% masking in the experiments.
  Please refer to our response to M3 for a detailed justification. To demonstrate that our model maintains consistent performance with different masking percentages, we repeated the intervention experiments with a 25% masking setup. The results, found in Figure 17 of the updated manuscript, show that the outcomes are consistent with those obtained using 5% masking. In both cases, CB-pLM successfully shifts the concept in the correct direction and significantly outperforms C-pLM and CC-pLM.

• E.2: Difficult to evaluate if performance gains are due to architecture changes or intervention technique.
  In Table 3 of Appendix D, we show that even when masking is applied randomly with the CB-pLM model, we achieve an average intervention effect of 0.0218. This performance is superior to the average change of 0.013 observed with CC-pLM using targeted masking. This demonstrates that CB-pLM is inherently more sensitive to interventions. Moreover, Table 3 illustrates that better localization, achieved through targeted interventions, further enhances our method, increasing the intervention magnitude by approximately 2x. This indicates that while the architecture of CB-pLM contributes significantly to its improved performance, the effectiveness of our intervention technique also plays a crucial role.

Dear Reviewer QrUQ,

Given that the end of the discussion period is approaching, we would like to ask if you have any further concerns or questions, particularly as a follow-up to our response? If we have addressed all your concerns, we hope you will consider increasing your score.

Thank you in advance!

Thank you for the detailed review and feedback. First, we address the weaknesses and then respond to your questions point-by-point.

Re: Weaknesses

• Concepts need to be explicitly defined for controlling design.
  To understand and control the concepts used by a model for generation, one can either explicitly add these concepts to the model, as we do in our paper, or use post-hoc methods (e.g., probes, SAEs, PCA) to identify them. Both approaches have tradeoffs. Our method requires defining concepts and some data annotation beforehand, while post-hoc methods involve a significant effort to analyze the model's representations after training. Post-hoc methods face challenges such as explanation faithfulness and concept disentanglement. These methods might capture elements in the model’s representations that the model doesn't actually use to determine its output. Even if a concept can be predicted from the representations, it may not be important to the model. Moreover, identifying a concept in the representations doesn't ensure reliable control over it, as the model may represent that concept alongside other correlated factors---disentanglement. By explicitly defining the concepts we want to control ahead of time, we avoid these issues. This approach requires more upfront work during training.
  Please note that we do not need the entire dataset to be fully annotated; a subset is often sufficient to learn the desired concepts. For example, when training, we add all concepts in SWISS-PROT; however, UniRef50 data do not have these concepts annotated in them, so we only have less than 3% of the data fully annotated.

• The argument that a CB-pLM is more interpretable is not convincing. CB-pLM doesn’t expand to a new interpretation of known biological properties.
  There are two types of interpretability offered by CB-pLM: Local Interpretability: This means that when generating a sequence, the model can: a) Indicate which concepts it relies on for its predictions. b) Accurately identify the input tokens that influence its output. We achieve the first point through our concept-bottleneck module and associated losses, and the second point through our proposed feature masking and regulation scheme. Global Interpretability: This involves insights into the relationship learned by the model between predefined concepts and various tokens. This is achieved through the linear decoder; for instance, our model can localize which amino acids are crucial for protein hydrophobicity. We have clarified this point in section 5.1 of the updated manuscript. Interpretations beyond the predefined concepts are outside the scope of this work.

• It is not possible to assess if the model improves the solubility and function of Siltuximab, based on the data provided.
  Please find our response in the general response R2.

Dear Reviewer ey43,

Given that the end of the discussion period is approaching, we would like to ask if you have any further concerns or questions, particularly as a follow-up to our response? If we have addressed all your concerns, we hope you will consider increasing your score.

Thank you in advance!

Thank you for addressing my questions. Specifically, I regrettably missed that only 3% of the data was annotated during training and can appreciate that it is quite impressive that the model can learn the correct and translatable concepts with so few annotations. I also appreciate the thoughtful response to the rest of my questions. I still am not convinced that the CB framework helps with model debugging. However, after reading all the reviews and your responses I have increased my score accordingly and have no follow up questions.

Changes in Manuscript

We would like to highlight the following changes to the manuscript to address reviewers' concerns:

• We have repeated the perplexity experiment in section 3.1 using a publicly available antibodies dataset Mason [14] as a test dataset to ensure that there is no cross-contamination between the test dataset and the training dataset of open-source models (Reviewer ey43).
• Added ESM2-650M and ESM2-3B as baselines in section 3.1 (Reviewer e4sF).
• Added ESM2 and Hydrophilic Resample as baselines for the Siltuximab experiment in section 4.2 (Reviewer ey43).
• Added TAP metrics as an additional validation for protein naturalness (Reviewer QrUQ).
• Added Hydrophobicity SAP/TAP metrics for Siltuximab redesigns as an additional proxy to wet-lab results (Reviewers ey43, AGed).
• Repeated the control experiments with 25% masking (Reviewer QrUQ).

References

[1] Li, Francesca-Zhoufan, et al. "Feature reuse and scaling: Understanding transfer learning with protein language models."

[2] Frey, Nathan C., et al. "Cramming Protein Language Model Training in 24 GPU Hours."

[3] Fournier, Quentin, et al. "Protein Language Models: Is Scaling Necessary?"

[4] Elnaggar, Ahmed, et al. "Ankh: Optimized protein language model unlocks general-purpose modelling."

[5] D'Amour, Alexander, et al. "Underspecification presents challenges for credibility in modern machine learning, 2020."

[6] Stanton, Samuel, et al. "Accelerating bayesian optimization for biological sequence design with denoising autoencoders." International Conference on Machine Learning.

[7] Gruver, Nate, et al. "Protein design with guided discrete diffusion."

[8] Frey, Nathan C., et al. "Protein discovery with discrete walk-jump sampling."

[9] Alamdari, S., et al. "Protein generation with evolutionary diffusion: Sequence is all you need."

[10] Tagasovska, Nataša, et al. "Implicitly Guided Design with PropEn: Match your Data to Follow the Gradient."

[11] Hewitt and Liang. "Designing and interpreting probes with control tasks."

[12] Kumar, Tan, and Sharma. "Probing classifiers are unreliable for concept removal and detection."

[13] D’Amour et al. "Underspecification presents challenges for credibility in modern machine learning."

[14] Mason, D. M., et al. "Optimization of therapeutic antibodies by predicting antigen specificity from antibody sequence via deep learning."