todo

[lucidrains]

• modules

  • SmoothLDDTLoss (thanks to @joseph-c-kim)
  • WeightedRigidAlign (@engelberger)
  • ExpressCoordinatesInFrame (@engelberger)
  • ComputeAlignmentError (@engelberger)
  • CentreRandomAugmentation (@engelberger)

• miscellaneous

  • f_tokenbond embedding to pairwise init (default to one single chain for starters if not passed in)
  • take care of normalization and unnormalization of atomic coordinates
  • distance labels should be derived from atom positions if not given
  • weighted rigid align module needs to account for atom_mask (variable number of atoms per batch sample)
  • sample without replacement in MSAModule
  • make sure diffusion loss accounts for mask of nucleic acid / ligand + bond loss during fine tuning
  • return the entire loss breakdown for logging in eventual trainer
  • hook up the centre random augmentation

• @lucidrains take care of

  • packed atom representation
    • given atom lengths and a sequence, do an average pool based on those lengths - atom -> token
    • given atom lengths and a sequence, expand sequence to consecutives, for token -> atom
  • fix packed atom representation when going from token level -> atom level pairwise repr
  • packed repr - make sure repeating pairwise is done in one specialized function, also take care of curtailing or padding the mask through some kwarg
  • able to pass in residue indices for only protein training, everything else derived
  • atom transformer attention bias needs to be calculated efficiently in the Alphafold3 module, use asserts to make sure shape is correct within local_attn fn
  • take care of residue identities / indices -> atom feats + atom bonds + attention biasing for atom transformers
  • allow for additional modifier embeddings to each molecule with optional scale
  • add a colab link or simple web server that visualizes all the molecules under life.py enumerated
  • allow for atom resolution confidence heads
  • a layer or two of atom attention in atom resolution confidence heads
  • allow for atom masking per residue, and use a missing residue null token if all the atoms are masked out
  • allow for atom resolution for modified aa and nucleotides
  • auto-detect distogram or token centre atom to be among missing atom and set to -1
  • computation of model selection score needs to be redone in atom resolution

• training

  • validation and test dataset
  • add config driven training with pydantic validation for constructing trainer and base model
  • saving and loading for both base alphafold3 model as well as trainer + optimizer states
  • add trainer orchestrator config that contains many training configs and one model
  • able to reconstitute the entire training history

• dataset classes for handling

  • single protein input
  • multimer input
  • multimer + nucleic acid(s) input
  • multimer + ligand input
  • for atom positions, create another dataclass that breaks it down by biomolecule type, and order + validate it automatically against what is given in Alphafold3Input
  • handle modifications to residues + nucleotides (phosphory, n-glycans, methylation)
  • figure out whether disulfide bonds are provided at any time

• dataset pipelines

  • handle atom level caching with a single decorator @lucidrains
  • proteins - pdb
  • nucleic acids
  • msa
  • templates

• improvisations

  • add attn logit soft capping, validated in gemma 2
  • add register tokens
  • improve atom transformer with some linear attention + other efficient attention tricks
  • frame averaging in place of their random aug
  • rectified flow instead of diffusion, using this repo
  • add layer sharing
  • instead of all the attention biasing complexity in atom transformer, alternate between GNN (with the sparse bonds) + flash attention
  • additional conditioning on diffusion module
  • use conditionally routed attention for atom encoder and decoder

• cleanup

  • remove unpacked representation