SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Table 3: OpenCell Validation Set Image Prediction Accuracy after Finetuning
Fine-Tuned	Threshold Image Encoder	Nucleus Proportion MAPE	Image MAE	PDF MAE	SSIM	FID	IS
No	HPA	0.0181 ± 0.0168	0.4154 ± 0.0594	0.3887 ± 0.1270	0.1250 ± 0.1149	3.9509	2.1739 ± 0.1255
No	OpenCell	0.0161 ± 0.0148	0.4953 ± 0.0064	0.3620 ± 0.1168	0.1220 ± 0.1188	1.5844	2.6069 ± 0.1175
Yes	HPA	0.0166 ± 0.0151	0.3776 ± 0.0834	0.3477 ± 0.1268	0.1869 ± 0.1503	17.4075	2.9113 ± 0.1199
Yes	OpenCell	0.0159 ± 0.0156	0.4996 ± 0.0006	0.3506 ± 0.1208	0.1574 ± 0.1372	2.5026	2.7168 ± 0.1137
Yes	HPA Finetuned	0.0170 ± 0.0160	0.3449 ± 0.1305	0.3487 ± 0.1340	0.1881 ± 0.1541	19.2683	3.6083 ± 0.2013
5.3 Finetuning
While the HPA dataset contains information about a wide variety of proteins, the model does not innately perform as well on the OpenCell data. We considered the potential of utilizing an HPA-trained model and finetuning on the OpenCell data, thereby introducing a wider protein context than what is found in the OpenCell data alone while adapting to the imaging conditions and cell type found within the new dataset. We experimented with different finetuning strategies for CELL-E2 on the OpenCell dataset. We used the pre-trained HPA checkpoint as the starting point for all finetuned models, continuing training on the OpenCell train set. We also evaluated the pre-trained HPA and OpenCell checkpoints without any finetuning as baselines. The finetuned models differed in how they updated the image encoders:
- HPA Finetuned (HPA VQGAN): we kept the original VQGAN image encoders from the HPA checkpoint.
- HPA Finetuned (OpenCell VQGAN): we replaced the image encoders with VQGANs trained only on OpenCell data.
- HPA Finetuned (Finetuned HPA VQGAN): we finetuned the HPA image encoders while keeping the rest of the model frozen, then froze the image encoders and update the transformer weights.
Fig. S8 shows image predictions on an OpenCell validation protein for models with hidden size  $= 480$ . Surprisingly, the pretrained HPA model already achieved strong performance on the OpenCell dataset without any finetuning (see Table S8).
The best results were obtained by utilizing a pre-trained HPA checkpoint. We first finetuned both VQGAN image encoders while freezing the rest of the model. We then froze the VQGAN weights and allowed the base transformer to update (see Table 3). We attribute the  $1.81\%$  improvement in MAE, along with the improvements in FID and IS, to the finetuning of both the VQGANs, as it improved the consistency of image patch tokens. This provided the checkpoint with more reliable image patches to generate from. However, swapping the HPA VQGAN with an OpenCell one led to a similar losses of distribution information seen in Fig. S7. This could be because the model overfits before being able to learn probabilities across tokens. The learning obstacle comes from the possibility that images patches within the finetuned OpenCell VQGAN have sufficient (or even more) pixel consistency with the images, but the patch positional indices are misaligned with those of the HPA VQGAN. These findings are consistent with those found in analogous text-to-image works utilizing diffusion models. We did not find that finetuning improved the model's sequence reconstruction ability (see Table S9).
6.1 CELL-E Comparison
In Table S10 and Table S11, we compare the performance of image localization prediction from scratch for CELL-E 2 and CELL-E. On the OpenCell validation set, CELL-E under-performs CELL-E 2 both before and after finetuning with regards to Nucleus Proportion MAPE. CELL-E 2 achieves worse Image and PDF MAE metrics before finetuning, however after finetuning CELL-E 2 achieves a  $2.2\%$  improvement for Image MAE and  $1.7\%$  for PDF MAE. On the contrary, CELL-E performs better with respect to image fidelity metrics SSIM and FID.
With respect to generation time, we found that the CELL-E 2 with hidden size of 480 was able to generate a prediction  $65 \times$  faster than the CELL-E model. This is a result of CELL-E 2's capability
to generate a prediction in a single step (.2784 seconds), which enables the advent of large-scale in silico mutagenesis studies.
6.2 De novo NLS Design
CELL-E 2's bidirectional integration of sequence and image information allows for an entirely novel image-based approach to de novo protein design. We applied CELL-E 2 to generate NLSs, which is a short amino acid sequence motif that can relocate a target protein into the cell nucleus when append to the target protein. In this case, our choice of the target protein is the Green Fluorescent Protein (GFP), a common protein engineering target [62-64] that is non-native to the human proteome and absent in the datasets. NLSs are usually identified by experimental mutagenesis studies or in silico screens that search for frequent sequences in nuclear proteins [51, 65]. However, these methods may yield candidates that are highly similar to known ones or not specific to the target protein. A more recent approach uses machine learning on sequence identity to augment featurization and statistical priors [17], but it is limited by the distribution of training samples due to the scarcity of experimentally verified NLSs. CELL-E 2 overcomes these limitations because it does not rely on explicit labels, and can therefore leverage significantly more unlabelled image data.
We generated a list of 255 novel NLS sequences for GFP using the procedure described in Appendix D.2. Briefly, we insert mask tokens of set length in a GFP sequence and task model with best sequence in-filling performance (0penCell1_2560) to fill in the masked amino acids, conditioned on a threshold image generated from the nucleus image (via Cellpose segmentation [66]). To verify the accuracy of the prediction, we pass the predicted sequence through the best performing image model (HPA Finteuned (Finetuned HPA VQGAN)_480), and quantify the proportion of signal intensity within the nucleus of the predicted threshold image (Fig. S9). The NLS sequences were then ranked based on sequence and embedding similarity with known NLSs (see Appendix D.2). The list of candidates can be found in Appendix D.3. We found several NLS candidates with high predicted signal in the nucleus, but which were fairly dissimilar from any protein found within NLSdb [65].
Classical NLSs are characterized by having regions of basic, positively charged amino acids arginine (R) and lysine (K) [67, 68], and are categorized as "monopartite" or "bipartite", either having a single cluster of basic amino acids or two clusters separated by a linker, respectively [69]. We observed a positive correlation between percentage of R and K residues in our predicted NLSs and sequence homology with known NLSs (Table S12). The number of clusters per sequence followed a similar trend, with sequences with relatively low sequence homology (Max ID % ≤ 33) having at most 2 clusters in 88% of predictions (Fig. S10). The remaining predictions, if correct, represent non-classical NLSs.
To further verify our predicted sequences, we passed the predicted NLS appended to GFP through Deeploc 2.0 [32], a leading sequence-to-class protein localization model, which predicted  $89\%$  of generated sequences were nuclear localizing and  $91\%$  contained a potential nuclear localizing signal.
Similar to CELL-E, we observed high attention weights on documented localization sequences correlated with positive protein signal within the threshold image (Fig. S11). For sequences with high predicted nucleus proportion intensities, we observed high activation across the entire sequence (novel NLS and GFP residues), with some NLS weights being an order of magnitude higher than others across the GFP sequences (Fig. 3). On the contrary, predicted sequences with comparatively less predicted intensity within the nucleus had low activation across the sequence, with little to none in the proposed NLS. We observed similar amounts of attention placed on the nucleus image patches, which largely corresponded to the location of the predicted threshold patches.
7 Conclusion & Future Work
In this paper, we have presented CELL-E 2, a novel bidirectional NAR model for protein design and engineering. CELL-E 2 can generate both image and sequence predictions, handle multimodal inputs and outputs, and run significantly faster than the state-of-the-art. By pre-training on a large HPA dataset and fine-tuning on OpenCell, CELL-E 2 can achieve competitive or superior performance on image and sequence reconstruction tasks. However, one limitation of CELL-E 2 is its output resolution, which is currently  $(256\times 256)$ . This resolution may not capture the fine details of microscopy images, which may limit applications in real-world use where Megapixel images are ac
Relative Attention Weights for Image Prediction
Figure 3: Attention weights associated with positive signal within the predicted image. Tokens with higher attention weight associated with background patches (low signal) are not highlighted. See Appendix D.3 for more information about the visualization process. We show 3 sequences with the highest (left column) and lowest (right column, not included in Table S13) predicted nucleus proportion intensity. The NLS+GFP sequences are shown with the designed NLS boxed in red.
quired. Increasing the output resolution of CELL-E 2 is one direction for future work. Furthermore, the sequence prediction struggles with the prediction of large stretches of amino acids as opposed to singular masked positions. Within this work, we encountered a trade-off between sequence prediction quality and prediction speed which may be overcome by reformulating the masking strategy. Similar findings were seen when compared with CELL-E, where we found accuracy measurements to improve with CELL-E 2 at the detriment of image quality metrics. The in-order prediction sequence we utilized in this paper may serve as a bottleneck for protein engineering applications despite the speed advantages gained from using a NAR architecture.
Another direction for future work is to incorporate structural information into the sequence embeddings. CELL-E 2 can generate novel NLS sequences with similar properties to GFP but low homology to existing sequences. However, the current sequence embeddings are based on a language model that may not capture all the structural features of the proteins. These features may affect the image appearance and vice versa.
We believe that CELL-E 2 is a promising model for protein design and engineering. We hope that our work will inspire more research on bidirectional NAR models for this domain and other domains that involve multimodal data.
8 Acknowledgments.
This research is supported in part by NIH grants R01GM131641 (BH) and R35-GM134922 (YSS). B.H. is a Chan Zuckerberg Biohub Investigator. A.A. contributed to visualizations and demos. E.K. was responsible for exploration and code.