Original: Knowledge distillation for fast and accurate DNA sequence correction
Revised: Knowledge distillation for fast and accurate DNA sequence correction

Original: Anonymous Author(s)

Original: AfﬁliationAddress

Original: Abstract
Revised: Abstract

Original: Accurate genome sequencing can improve our understanding of biology and thegenetic basis of disease.
Revised: Accurate genome sequencing can improve our understanding of biology and the genetic basis of disease.

Original: The standard approach for generating DNA sequencesfrom PacBio instruments relies on HMM-based models.
Revised: The standard approach for generating DNA sequences from PacBio instruments relies on HMM-based models.

Original: Here, we introduce
Revised: Here, we introduce

Original: Distilled DeepConsensus - a distilled transformer–encoder model for sequencecorrection, which improves upon the HMM-based methods with runtime constraintsin mind.
Revised: Distilled DeepConsensus - a distilled transformer–encoder model for sequence correction, which improves upon the HMM-based methods with runtime constraints in mind.

Original: Distilled DeepConsensus is 1.3x faster and 1.5x smaller than its largercounterpart while improving the yield of high quality reads (Q30) over the
Revised: Distilled DeepConsensus is 1.3x faster and 1.5x smaller than its larger counterpart while improving the yield of high quality reads (Q30) over the

Original: HMM-based method by 1.69x (vs. 1.73x for larger model) .
Revised: HMM-based method by 1.69x (vs. 1.73x for larger model) .

Original: With improvedaccuracy of genomic sequences, Distilled DeepConsensus improves downstreamapplications of genomic sequence analysis such as reducing variant calling errorsby 39% (34% for larger model) and improving genome assembly quality by3.8% (4.2% for larger model) .
Revised: With improved accuracy of genomic sequences, Distilled DeepConsensus improves downstream applications of genomic sequence analysis such as reducing variant calling errors by 39% (34% for larger model) and improving genome assembly quality by 3.8% (4.2% for larger model) .

Original: We show that the representations learned by
Revised: We show that the representations learned by

Original: Distilled DeepConsensus are similar between faster and slower models.
Revised: Distilled DeepConsensus are similar between faster and slower models.

Original:  Introduction
Revised:  Introduction

Original: Genome sequencing is widely used to advance our understanding of biological processes, the geneticbasis of diseases, and is increasingly used for personalized medicine directly in the clinic (1).
Revised: Genome sequencing is widely used to advance our understanding of biological processes, the genetic basis of diseases, and is increasingly used for personalized medicine directly in the clinic (1).

Original: Thus,obtaining accurate genome sequences is of paramount importance.
Revised: Thus, obtaining accurate genome sequences is of paramount importance.

Original: Over the years, sequencingplatforms have been developed by Illumina, Paciﬁc Biosciences (PacBio) and Oxford Nanopore, thathave had either short and accurate reads or long and noisy reads.
Revised: Over the years, sequencing platforms have been developed by Illumina, Paciﬁc Biosciences (PacBio) and Oxford Nanopore, that have had either short and accurate reads or long and noisy reads.

Original: Recently, PacBio developed a novellong read HiFi sequencing technology (2) that has led to new state-of-the-art performance in variantcalling (3) and the ﬁrst telomere-to-telomere human assembly (4).
Revised: Recently, PacBio developed a novel long read HiFi sequencing technology (2) that has led to new state-of-the-art performance in variant calling (3) and the ﬁrst telomere-to-telomere human assembly (4).

Original: While highly accurate, some errorsstill exist and here we focus on whether we can use machine learning, and in particular transformers,to further improve the accuracy of the genomic sequences with on-device runtime constraints in mind.
Revised: While highly accurate, some errors still exist and here we focus on whether we can use machine learning, and in particular transformers, to further improve the accuracy of the genomic sequences with on-device runtime constraints in mind.

Original: PacBio HiFi sequencing works by circularizing DNA, allowing for multiple sequencing passes overthe same molecule.
Revised: PacBio HiFi sequencing works by circularizing DNA, allowing for multiple sequencing passes over the same molecule.

Original: This produces a set of noisy sequences/subreads of the same region.
Revised: This produces a set of noisy sequences/subreads of the same region.

Original: The subreadsare then combined into a consensus sequence/read using a hidden Markov model (HMM) (2).
Revised: The subreads are then combined into a consensus sequence/read using a hidden Markov model (HMM) (2).

Revised: In

Original: Inorder to further improve upon the accuracy of the HMM-based method for correct DNA sequencegeneration, DeepConsensus - an encoder-only transformer model has been developed (5).
Revised: These order to further improve upon the accuracy of the HMM-based method for correct DNA sequence generation, DeepConsensus - an encoder-only transformer model has been developed (5).

Original: However,the accuracy and yield improvements of DeepConsensus come with latency and model-size costs.
Revised: However, the accuracy and yield improvements of DeepConsensus come with latency and model-size costs.

Original: For example, unlike HMMs, transformer models are too large to be run on sequencing platforms withresource constraints.
Revised: For example, unlike HMMs, transformer models are too large to be run on sequencing platforms with resource constraints.

Original: As sequencing throughput increases, the problem compounds.
Revised: As sequencing throughput increases, the problem compounds.

Original: In this work, we use Knowledge Distillation (KD) (6) to improve the runtime of the transformer-based
Revised: In this work, we use Knowledge Distillation (KD) (6) to improve the runtime of the transformer-based

Original: DeepConsensus model, using techniques similar to (7; 8; 9).
Revised: DeepConsensus model, using techniques similar to (7; 8; 9).

Original: Our model is much smaller and fasterthan the original model, with slight performance degradation.
Revised: Our model is much smaller and faster than the original model, with slight performance degradation.

Original: Importantly, these runtime boosts arecritical for enabling real time on-sequencing-device inference.
Revised: Importantly, these runtime boosts are critical for enabling real time on-sequencing-device inference.

Original: In other domains of biology, as theamount of data increases, such as sequencing millions of single cells, and model size increases (10), [Figure] M i [Figure]
Revised: In other domains of biology, as the amount of data increases, such as sequencing millions of single cells, and model size increases (10), the runtime considerations may become increasingly important.

Original: m a t c h e

Original: [Figure]

Original: N o n H P

Original: i n s e r t i o n s

Original: N o n H P d e l e t i o n s

Original: H P i n s e r t i o n s H P

Original: [Figure]

Original: e

Original: [Figure]

Original: e t i o n s

Original: [Figure]

Original: the runtime considerations may become increasingly important.

Original: This work explores initial steps inthis direction and our contributions are as follows:
Revised: This work explores initial steps in this direction and our contributions are as follows:

Original: We show that knowledge distillation can improve model accuracy, resulting in a faster modelwithout a signiﬁcant degradation in performance as compared to a large model.
Revised: We show that knowledge distillation can improve model accuracy, resulting in a faster model without a signiﬁcant degradation in performance as compared to a large model.

Original: We demonstrate relative performance parity between larger and smaller models on downstream tasks like variant calling and de novo assembly generation.3.
Revised: We demonstrate relative performance parity between larger and smaller models on downstream tasks like variant calling and de novo assembly generation.

Original: We show that the representations learned by larger and smaller models are similar.
Revised: We show that the representations learned by larger and smaller models are similar.

Original:  Model
Revised:  Model

Original: DeepConsensus (5) takes as input a tensor of length 100bp containing information about the subreadssuch as the noisy subread sequences, strand direction, pulse width, interpulse duration as well asgeneral information such as signal-to-noise ratio of each nucleotide and the consensus sequence givenby pbccs HMM-based algorithm (Fig. 1a).
Revised: DeepConsensus (5) takes as input a tensor of length 100bp containing information about the subreads such as the noisy subread sequences, strand direction, pulse width, interpulse duration as well as general information such as signal-to-noise ratio of each nucleotide and the consensus sequence given by pbccs HMM-based algorithm (Fig. 1a).

Original: The goal is to predict the correct consensus sequencefrom the noisy subread sequences and the other auxiliary information.
Revised: The goal is to predict the correct consensus sequence from the noisy subread sequences and the other auxiliary information.

Original: Since this problem is similarto sequence-to-sequence tasks, Transformers (11) are leveraged for this problem.
Revised: Since this problem is similar to sequence-to-sequence tasks, Transformers (11) are leveraged for this problem.

Original: Both the teacher and student models are encoder-only transformers.
Revised: Both the teacher and student models are encoder-only transformers.

Original: The teacher has six encoderstacks.
Revised: The teacher has six encoder stacks.

Original: We vary the number of encoder stacks from 3 to 5 to obtain smaller (student) Distilled
Revised: We vary the number of encoder stacks from 3 to 5 to obtain smaller (student) Distilled

Original: DeepConsensus models of different sizes.
Revised: DeepConsensus models of different sizes.

Original: Each encoder stack is composed of a self-attention layerwith hidden dimension of 280 and 2 attention heads and a feedforward network with a ﬁlter size of2048.
Revised: Each encoder stack is composed of a self-attention layer with hidden dimension of 280 and 2 attention heads and a feedforward network with a ﬁlter size of 2048.

Original: The output of the entire encoder block, which represents the contextual embedding of eachbase, is decoded with the feedforward layer followed by a softmax function, that for each base outputsthe probability of A, C, T, G or a gap.
Revised: The output of the entire encoder block, which represents the contextual embedding of each base, is decoded with the feedforward layer followed by a softmax function, that for each base outputs the probability of A, C, T, G or a gap.

Original: We use this simpler decoder to save computational resources.
Revised: We use this simpler decoder to save computational resources.

Original: For our labels, we leverage a high quality assembly of the human genome (T2T-CHM13), recentlypublished by the Telomere-to-Telomere consortium (4), as part of the effort to obtain the completesequence of the human genome and use data (2 SMRT Cells) from a CHM13 sample with 24-kbinsert sizes from (5). 
Revised: For our labels, we leverage a high quality assembly of the human genome (T2T-CHM13), recently published by the Telomere-to-Telomere consortium (4), as part of the effort to obtain the complete sequence of the human genome and use data (2 SMRT Cells) from a CHM13 sample with 24-kb insert sizes from (5). [Figure]

Original: We use an alignment-based loss function, which is based on differentiable dynamic programming (12),
Revised: We use an alignment-based loss function, which is based on differentiable dynamic programming (12),

Original: L ( y, p ) to train the teacher model as in (5), where y =
Revised: L ( y, p ) to train the teacher model as in (5), where y =

Original: ( y 1 , ..., y m ) is the true sequence andp =
Revised: ( y 1 , ..., y m ) is the true sequence and p =

Original: ( p 1 , ..., p n ) are the predicted probabilities with p i representing the probability over A, C, T, G ora gap classes.
Revised: ( p 1 , ..., p n ) are the predicted probabilities with p i representing the probability over A, C, T, G or a gap classes.

Original: In order to train smaller student models from the teacher, in addition to the alignment-based lossfunction, we use a typical knowledge distillation loss (6), to update the model weights of the student:
Revised: In order to train smaller student models from the teacher, in addition to the alignment-based loss function, we use a typical knowledge distillation loss (6), to update the model weights of the student:

Original: [Equation]
Revised: [Equation]

Original: where p student and p teacher are the probabilities output by the student and the teacher models, respectively, and α is determined using a hyper-parameter search.
Revised: where p student and p teacher are the probabilities output by the student and the teacher models, respectively, and α is determined using a hyper-parameter search.

Original: We explore both the Kullback–Leibler (KL) divergence and the mean squared error (computed per base and averaged across positions) for
Revised: We explore both the Kullback–Leibler (KL) divergence and the mean squared error (computed per base and averaged across positions) for

Original: L KD .
Revised: L KD .

Original: In addition, we initialize the student layers from the teacher layers, since this was found to bebeneﬁcial in (7).
Revised: In addition, we initialize the student layers from the teacher layers, since this was found to be beneﬁcial in (7).

Original:  Results 
Revised:  Results 3.1 Performance vs. runtime trade-off

Original: 3.1 Performance vs. runtime trade-off

Original: We measure the model performance by comparing the yield at empirical Q30 (number of bases inreads with Q30 or higher quality score, i.e. identity of 0.999 or higher) produced by DeepConsensusmodels versus pbccs.
Revised: We measure the model performance by comparing the yield at empirical Q30 (number of bases in reads with Q30 or higher quality score, i.e. identity of 0.999 or higher) produced by DeepConsensus models versus pbccs.

Original: This metric (yield at empirical Q30 over CCS) indicates how many more highquality bases we obtain with a particular model as compared to pbccs.
Revised: This metric (yield at empirical Q30 over CCS) indicates how many more high quality bases we obtain with a particular model as compared to pbccs.

Original: The models are trained onthe data collected from CHM13 cell line (chromosomes 1-18, 2 SMRT Cells), hyper-parameters areoptimized using CHM13 cell line (chromosomes 21-22, 2 SMRT Cells), and the models are tested ona completely held-out genome HG002 (chromosome 20, 1 SMRT Cell).
Revised: The models are trained on the data collected from CHM13 cell line (chromosomes 1-18, 2 SMRT Cells), hyper-parameters are optimized using CHM13 cell line (chromosomes 21-22, 2 SMRT Cells), and the models are tested on a completely held-out genome HG002 (chromosome 20, 1 SMRT Cell).

Original: Here, we use the KL lossfor matching teacher and student probabilities and initialize the student model from top layers of theteacher.
Revised: Here, we use the KL loss for matching teacher and student probabilities and initialize the student model from top layers of the teacher.

Original: Model runtime is measured using number of seconds on a Google Cloud VM with Intel[-
-]
Revised: Model runtime is measured using number of seconds on a Google Cloud VM with Intel[-
-]

Original: CPU cores (n1-standard-32).
Revised: CPU cores (n1-standard-32).

Original: In order to assess whether distillation is useful, we also train student models of the same size butwithout the distillation loss, only with L ( y, p student ) .
Revised: In order to assess whether distillation is useful, we also train student models of the same size but without the distillation loss, only with L ( y, p student ) .

Original: We show that distillation consistently resultsin improved yield at empirical Q30 over CCS across different student architectures (3, 4, 5 layers)and replicates (Fig. 1b).
Revised: We show that distillation consistently results in improved yield at empirical Q30 over CCS across different student architectures (3, 4, 5 layers) and replicates (Fig. 1b).

Original: Some student models, for example the 5 layer student model, achieve a similar yield at empirical
Revised: Some student models, for example the 5 layer student model, achieve a similar yield at empirical

Original: Q30 over CCS as the teacher while taking 90% of teacher’s runtime .
Revised: Q30 over CCS as the teacher while taking 90% of teacher’s runtime .

Original: For the remaining analysis,we focus on the student model with 4 layers, 1.3x faster and 1.5x smaller (in terms of # parameters)than the teacher, and slight performance degradation (1.69 vs 1.73) compared to the teacher.
Revised: For the remaining analysis, we focus on the student model with 4 layers, 1.3x faster and 1.5x smaller (in terms of # parameters) than the teacher, and slight performance degradation (1.69 vs 1.73) compared to the teacher.

Original: 3.2 Teacher versus student error proﬁles
Revised: 3.2 Teacher versus student error proﬁles

Original: We analyze whether the teacher and the student make similar kinds of errors.
Revised: We analyze whether the teacher and the student make similar kinds of errors.

Original: We quantify the numberof base pairs with mismatches, insertions and deletions per read and additionally subdivide errors bytheir sequence context.
Revised: We quantify the number of base pairs with mismatches, insertions and deletions per read and additionally subdivide errors by their sequence context.

Original: It has been shown that PacBio HiFi sequencing has higher error rates withinhomopolymer (HP) regions (2), which are regions with consecutive repeating bases (e. g. AAAA),and thus we split indel errors by HP and non-HP context.
Revised: It has been shown that PacBio HiFi sequencing has higher error rates within homopolymer (HP) regions (2), which are regions with consecutive repeating bases (e. g. AAAA), [Figure] [Table]

Revised: , DeepConsensus and Distilled DeepConsensus reads from 2 SMRT Cells.

Revised: and thus we split indel errors by HP and non-HP context.

Original: Fig.
Revised: Fig.

Original: 1c shows that the error types made bythe teacher and the student are similar with HP indels and in particular HP insertions for the studentbeing the major error type.
Revised: 1c shows that the error types made by the teacher and the student are similar with HP indels and in particular HP insertions for the student being the major error type.

Original: 3.3 Performance on downstream tasks: variant calling and assembly
Revised: 3.3 Performance on downstream tasks: variant calling and assembly

Original: [Figure] Variant calling We perform variant calling analysis using a custom DeepVariant (13) model (warm- [Figure] started from DeepVariant 1.4) trained on sequences from 2 SMRT Cells generated by DeepConsensusteacher and student models respectively and evaluated on held-out chromosome 20.
Revised: Variant calling We perform variant calling analysis using a custom DeepVariant (13) model (warmstarted from DeepVariant 1.4) trained on sequences from 2 SMRT Cells generated by DeepConsensus teacher and student models respectively and evaluated on held-out chromosome 20.

Original: For sequencesoutput by pbccs, we use DeepVariant 1.4 (trained on pbccs output).
Revised: For sequences output by pbccs, we use DeepVariant 1.4 (trained on pbccs output).

Original: Fig.
Revised: Fig.

Original: 2 shows that the studentmodel has 39% error reduction in variant calling in comparison to pbccs with the student beingslightly better than the teacher on this task (34% reduction for larger model).[Figure] [Table]
Revised: 2 shows that the student model has 39% error reduction in variant calling in comparison to pbccs with the student being slightly better than the teacher on this task (34% reduction for larger model).

Original: , and false duplication rates of the genome assemblies generated with pbccs, DeepConsensus and Distilled DeepConsensus reads from 2 SMRT Cells.

Original: Assembly analysis Given increased yield and higher quality reads generated by the DeepConsensusstudent model and the larger teacher model
Revised: Assembly analysis Given increased yield and higher quality reads generated by the DeepConsensus student model and the larger teacher model

Original: , we compare the student and the teacher models in termsof de novo assembly generation.
Revised: , we compare the student and the teacher models in terms of de novo assembly generation.

Original: Phased assemblies were generated with hiﬁasm assembler usingreads (from 2 SMRT Cells) output by pbccs, and DeepConsensus student and teacher models, similarto (5).
Revised: Phased assemblies were generated with hiﬁasm assembler using reads (from 2 SMRT Cells) output by pbccs, and DeepConsensus student and teacher models, similar to (5).

Original: In Table 1, we report the estimated qualities, using YAK (14), of the assemblies and show thatthe mean quality value (across haplotypes) is similar between the teacher and the student modelswith both being better than pbccs (3.8% and 4.2% for student and teacher, respectively).
Revised: In Table 1, we report the estimated qualities, using YAK (14), of the assemblies and show that the mean quality value (across haplotypes) is similar between the teacher and the student models with both being better than pbccs (3.8% and 4.2% for student and teacher, respectively).

Original: Additionally,both the teacher and the student models increase assembly gene completeness and decrease falsegene duplication rate.
Revised: Additionally, both the teacher and the student models increase assembly gene completeness and decrease false gene duplication rate.

Original:  Analysis
Revised:  Analysis

Original: We performed a study of distillation training options to evaluate which factors contributed the most togood performance in the student models.
Revised: We performed a study of distillation training options to evaluate which factors contributed the most to good performance in the student models.

Original: We experimented with three distillation losses: alignmentto ground truth and Kullback Leibler (KL) divergence with the teacher probabilities, alignment toground truth and MSE with the teacher probabilities, and just KL with the teacher probabilities.
Revised: We experimented with three distillation losses: alignment to ground truth and Kullback Leibler (KL) divergence with the teacher probabilities, alignment to ground truth and MSE with the teacher probabilities, and just KL with the teacher probabilities.

Original: Additionally, we experimented with four types of initialization schemes.
Revised: Additionally, we experimented with four types of initialization schemes.

Original: The relative beneﬁts of eachcan be seen in Fig.
Revised: The relative beneﬁts of each can be seen in Fig.

Original: 3a and b.
Revised: 3a and b.

Original: We found that MSE outperformed KL divergence, and that initializingthe student network to "alternating" layers (1, 2, 4, 6 of the teacher) performed best among layerinitialization schemes.
Revised: We found that MSE outperformed KL divergence, and that initializing the student network to "alternating" layers (1, 2, 4, 6 of the teacher) performed best among layer initialization schemes.

Original: Fig.
Revised: Fig.

Original: 3c shows a per-example agreement analysis (15) to better understand whether models with betteroverall performances are just correctly predicting more examples, or if they’re correctly predictingdifferent examples.
Revised: 3c shows a per-example agreement analysis (15) to better understand whether models with better overall performances are just correctly predicting more examples, or if they’re correctly predicting different examples.

Original: The results show that the better models are more accurate only slightly morethan the worse models when there is a disagreement.
Revised: The results show that the better models are more accurate only slightly more than the worse models when there is a disagreement.

Original: This means that the set of examples correctlypredicted by the better models is a larger and different set than the set of examples correctly predictedby the worse models, but neither is as a superset of the other.
Revised: This means that the set of examples correctly predicted by the better models is a larger and different set than the set of examples correctly predicted by the worse models, but neither is as a superset of the other.

Original: To test whether teachers and students learn the same representations, we apply linear Centered Kernel

Original: Alignment (CKA) between pairs of layers across networks, following the methodology of (16).

Original: CKAcomputes a [0, 1]-valued similarity between two Gram matrices computed from two representationsover the same sample of input examples (17).

Original: We see from the student vs. student plot (Fig.

Original: 3d, left)that the internal representations of the student model are very stable across different replicates.

Original: In theteacher vs. student plot (Fig.

Original: 3d, right), we ﬁnd that layers in the student network correspond wellwith particular layers in the teacher.

Original: Interestingly, we ﬁnd that despite initializing with top layersfrom the teacher ( 1 → 1 , 2 → 2 , 4 → 3 , 6 → 4 ), the student layers after training have a differentcorrespondence ( 1 → 1 , 2 or 4 → 2 , 4 → 3 , 6 → 4 ).

Original: This representational difference might alsoexplain why the multiheaded attention of the student is different than the teacher: for example forlayer 3, the student attention head looks an average of 1.8 base pairs different than the teacher forlayer 4.

Original: [Figure]
Revised: [Figure]

Original: r

Original: [Figure]

Original: (b) t e a c

Original: [Figure]

Original: e r [Figure] t u d
Revised: e r s t u d

Original: [Figure]
Revised: [Figure]

Original: n t
Revised: n t

Original: [Figure]
Revised: [Figure]

Original: t
Revised: t

Original: [Figure]
Revised: [Figure]

Original: d

Original: [Figure]

Original: n t r e
Revised: n t r e

Original: [Figure]
Revised: [Figure]

Original: lic
Revised: li

Original: [Figure]
Revised: [Figure]

Original: t e s t
Revised: a t

Original: [Figure]
Revised: [Figure]

Original: d
Revised: d

Original: [Figure]
Revised: [Figure]

Original: nt ( s u
Revised: n

Original: [Figure]
Revised: [Figure]

Original: p t i
Revised: p t i

Original: [Figure]
Revised: [Figure]

Original: l )
Revised: y (

Original: [Figure]
Revised: [Figure]

Original: [Table] 50(c)    4 Student (replicate) S t u d e n t 1 2 3 4 5 6 Teacher    0.
Revised: example replicates different student from Fig.

Original: C K A s i m il a r i t y
Revised: 3b not To test whether teachers and students learn the same representations, we apply linear Centered Kernel Alignment (CKA) between pairs of layers across networks, following the methodology of (16). CKA computes a [0, 1]-valued similarity between two Gram matrices computed from two representations over the same sample of input examples (17).

Original: (d)
Revised: We see from the student vs. student plot (Fig. 3d, left) that the internal representations of the student model are very stable across different replicates. In the teacher vs. student plot (Fig. 3d, right), we ﬁnd that layers in the student network correspond well with particular layers in the teacher. Interestingly, we ﬁnd that despite initializing with top layers from the teacher ( 1 → 1 , 2 → 2 , 4 → 3 , 6 → 4 ), the student layers after training have a different correspondence ( 1 → 1 , 2 or 4 → 2 , 4 → 3 , 6 → 4 ). This representational difference might also explain why the multiheaded attention of the student is different than the teacher: for example for layer 3, the student attention head looks an average of 1.8 base pairs different than the teacher for layer 4.