Folding a Monomer with Structural Context from Other Monomers

[architvasan]

Hi, thanks for the help with the other issue!

I have another question similar to that one.

General description:
I want to fold a single monomer while incorporating structural information from other monomers in a complex. However, I do not want the model to fold the other monomers—just use them as a reference/template to guide the folding of the target monomer.

I am unable to fold the entire complex due to token and memory limitations. Including structural context from the other monomers—without explicitly folding them—could help bypass these restrictions while still capturing important interactions.

Questions:

Does the model support folding a monomer while keeping other monomers fixed as a structural reference?
If not, is there a way to provide structural constraints from the known multimeric complex?
What would be the best approach to ensure the folded monomer correctly interacts with the rest of the complex?
Use Case:

I have an experimental structure for subunit A and need to fold a subunit B while keeping subunit B static as a guide.
Due to memory/token constraints, folding the entire complex is not really an option.

I would be open to suggestions involving incorporating protein-protein docking with CHAI-1 if this is too complicated for the current/near future package.

Any guidance on how to achieve this would be greatly appreciated! Thanks!

[arogozhnikov]

Hi Archit,

If I understand correctly, you already know some interactions of monomer you want to fold with other sequences.
If so, you can try removing aminoacids from other proteins that are not in direct contact with a monomer of interest.

To keep model informed about the rest of protein (that you removed) without modelling it, you'll need to provide

• ESMs (embed full sequence, then subselect AAs of interest)
• MSAs (same, you'll need to subsample to selected tokens)
• templates (this will inform model about relative positions of AAs)
• optionally, restraints. Restraints can help to orient complex the right way - where each protein goes.

This will be efficient in terms of compute (fewer tokens), but that will be quite an exercise, as you need to first build a full AllAtomFeatureContext, then subsample it to leave only 'necessary tokens'.

[architvasan]

Hi, thank you for the very prompt response!

So, I actually tried the idea of removing amino acids for subunit B not in direct contact of A but I found it was difficult to choose the distance cutoff for truncating B. When I chose 15 A cutoff, this included too much of the subunit B and it was not computationally efficient but produced a much better cutoff than a 10 A cutoff.

I would be very interested in your second point and I am willing to build the code for this! Could you please point me to the code I would need to modify in order to do this?

So could I compute the ESM embeddings using this kinda code:

import torch
from esm import pretrained

# Load a pretrained ESM model (example: ESM-2)
model, alphabet = pretrained.esm2_t33_650M_UR50D()
batch_converter = alphabet.get_batch_converter()

# Input sequence
sequence = "MKTLLLTLVVVTIVCLDLGYT"  # Replace with your full protein sequence
data = [("protein1", sequence)]
batch_labels, batch_strs, batch_tokens = batch_converter(data)

# Disable gradient calculations (for inference)
with torch.no_grad():
    results = model(batch_tokens, repr_layers=[33], return_contacts=False)

# Extract embeddings (layer 33 is usually best for structure-relevant features)
embeddings = results["representations"][33]

and then sample the amino acids I want from here?

How would I go about creating a structural template here using your available code base? I see a chai_lab.data.dataset.templates.context file. Which classes/functions should I play with here to accomplish this?

Where would I go about creating an AllAtomFeatureContext for the system of interest?

[arogozhnikov]

You'll need to understand what goes into AllAtomFeatureContext first.

Trick would be to first construct a large AllAtomFeatureContext with all proteins of interest (you'll need to remove checks for n_tokens), serialize it, but not pass it to run_folding_on_context.

And then to carefully construct a smaller AllAtomFeatureContext by subselecting only necessary tokens (in all of input fields) and pass it to run_folding_on_context.

[architvasan]

Okay, this sounds straightforward. I see AllAtomFeatureContext here: https://github.com/chaidiscovery/chai-lab/blob/main/chai_lab/data/dataset/all_atom_feature_context.py#L25
I will go back to the AllAtomStructureContext, TemplateContext, EmbeddingContext, and RestraintContext classes to see if I can implement the large AllAtomFeatureContext first for folding the protein of interest. Then it sounds not too difficult to perform the subselection procedure. Please let me know if I am missing anything here.

For more context, I am doing this for a design workflow where most of the system remains static and a small region is mutated several thousand times. This approach seems like it would mitigate the computational bottlenecks here once it would work!

[architvasan]

Another two possibly silly questions are regarding using TemplateContext and AllAtomStructureContext.

1. What is the difference between the AllAtomStructureContext and the TemplateContext in your code? Are they both imparting information about the template structure?

2. What is the workflow for constructing a TemplateContext? I see these arguments for TemplateContext:

template_restype: Int[Tensor, "n_templates n_tokens"]
template_pseudo_beta_mask: Bool[Tensor, "n_templates n_tokens"]
template_backbone_frame_mask: Bool[Tensor, "n_templates n_tokens"]
template_distances: Float[Tensor, "n_templates n_tokens n_tokens"]
template_unit_vector: Float[Tensor, "n_templates n_tokens n_tokens 3"]

Would I need to construct these arguments myself from the pdb or is there logic in another class that would construct these from an input pdb?

Thanks!

[arogozhnikov]

TemplateContext - we describe how one can construct it in #101 (you'll need to prepare those yourself).

the AllAtomStructureContext and the TemplateContext in your code? Are they both imparting information about the template structure?

No, only templateContext contains this info. AllAtom contains templateContext as one of its fields.

This approach seems like it would mitigate the computational bottlenecks here once it would work!

yes, this sounds directionally correct. As a side thought: we never trained with template information that connects different proteins in a complex (only different AAs within same chain), but this may still work, so I recommend trying that.

[wukevin]

Just a heads up -- we've just added beta support for templates after merging #300