Editing Antibody Design by AI (section)

== Optimizing antibody sequences with gradient ascent and a neural network ensemble ==

Given that direct optimization on one-hot discrete space is difficult and inefficient, we used the gradients back-propagated from the neural network output layer to the input layer to guide the improvement of an input sequence. We relaxed the one-hot constraint during optimization to fully take advantage of gradient methods, periodically projecting the current continuous representation back to a discrete one-hot input by taking the arg-max of each amino acid position. By tracking the improvement and checking the convergence in one-hot space, we made sure that the optimization over one-hot space was happening along with the continuous space. By choosing different step sizes λ and projection intervals k, we were able to control the level of divergence between proposed sequences and seed sequences. The gradient ascent approach can easily be applied to batches of seed sequences to fully utilize the parallel computing power of GPUs, and thus is efficient.

Because of the deterministic characteristic of gradient ascent, a single network can only produce limited number of sequences following the gradients determined by its local function landscape. We used an ensemble of 24 networks trained with different initialization and different subset of sequences, such that each of them learns a landscape with different “perspectives”. This allowed us to optimize a single seed sequence with multiple potential directions and thus produced much more diverse candidates than a single network. We also varied the optimization hyperparameters, namely the step size (λ=1e-4, 5e-4, 1e-3, 5e-3, 1e-2) and projection intervals (k=10, 20, 30, 40, 50, 60, 70, 80), in order to produce diverse sequences 36 of varying distances to their corresponding seed. In total, we used 960 optimization methods to generate a pool of candidates.

A single neural network may overfit a training set given that a large portion of sequence space is not observed in our training data. We used a two-step voting-thresholding strategy with an ensemble to increase the robustness of our optimization results to tackle this issue. First, we used all 960 optimization methods to vote on optimal sequences. Each method voted for each sequence it produced, and the sequences with the top number of votes were selected for next step. This voting step prioritizes sequences that were optimized with parameters that are neither too conservative nor too aggressive, and sequences that were constantly proposed by networks trained from different perspectives. On the following step, we used an ensemble consisting of all 18 regression neural networks trained as before to evaluate each candidate. Different networks tend to disagree in the regions where model uncertainty is high, and thus can be used to increase the robustness of enrichment prediction for unseen sequences. We threshold the candidates by the 5% confidence lower bound of the scores obtained from the 18 neural nets, and chose the sequences with non-negative lower bound as the positive candidates that will be used for comparison with the seeds. We used sequences with negative scores as negative controls to validate the performance of computational models.