Efficiency of Learned Indexes on Genome Spectra

Data structures on $k$-mers (fixed-length strings) originating from genomes form a crucial component of many bioinformatics tools [22]. Their space usage is now a major bottleneck at the forefront of biological discovery. For example, the Logan project needed two petabytes to index the 31-mers from the Sequence Read Archive [7]. Existing data structures often meet the space challenge by exploiting the non-uniform structure of genomic $k$-mer sets [6]. In doing so, they circumvent [23, 28, 17, 2] the theoretical worst- and average-case lower bounds, which offer a more pessimistic perspective [25].

A popular approach is based on learned indexes, a recent research direction that takes advantage of the internal patterns in the rank curve of a dataset [5, 18, 32, 10]. The rank of an element $x$ in a multiset $S$ is defined as the position of the first occurrence of $x$ in the sorted list of $S$. A learned index is constructed from $S$ and, given a query $k$-mer, returns an error-bounded prediction for its rank. Initial approaches used machine learning to construct the index [13, 15, 14], but it later turned out that in this setting, it is more effective to model the rank curve using a piecewise linear approximation [17, 2, 4], as depicted in Figure 1.

Learned indexes based on piecewise linear approximations (PLAs) of the rank curve have resulted in practically groundbreaking time/space tradeoffs, both on genomic [17, 8, 2] and other kinds of data [11, 12, 8, 16, 21, 31, 20, 19]. However, existing theoretical worst-case analyses only show improved performance under assumptions that are unrealistic for genomic datasets, such as the independence and identical distribution of gaps between $k$-mers [9]. This leaves a gap between the theoretical analysis of $k$-mer-based learned indexes and their much better performance on real data [17, 2]. By better understanding and characterizing the structural properties of genomic $k$-mer sets, we can not only better predict the performance of learned indexes but also drive the development of new PLA-based tools that maximally exploit the properties of real data.

Our approach in this paper is to develop a measure of piecewise linear (PL) approximability which can be used to parameterize the theoretical analysis of PLA-based methods. Parameterizing string complexity is an active area of research [26], but none of the existing measures capture the PL-approximability of a genome’s rank curve. Previous papers have proposed modeling the PL-approximability of $S$ using a power-law [11, 9], and have implicitly used the two parameters of a power-law fit as proxies for the PL-approximability of $S$. However, the interpretation of such parameters is limited when the rank curve strays from a power-law, as (we will show) occurs in real genomic data.

We develop a novel measure of PL-approximability, which we call CaPLa (Canonical PL-approximability) (Section 2). It builds on the idea of a power-law fit but adds the uncertainty of the fit as part of the measure. We prove basic properties of CaPLa, including its existence, and give an algorithm to compute it (Section 3). The algorithm is exact under certain conditions of the data, which we empirically show are nearly always met in practice. We demonstrate how CaPLa can be applied to derive space bounds that accurately reflect properties of real data (Section 4).

Finally, we apply CaPLa to analyze genome spectra, where a spectrum is the multiset of all $k$-mers appearing in a string.¹¹1Our tool for computing CaPLa is available at https://github.com/medvedevgroup/CaPLa. We use a dataset of more than 500 genomes and find that CaPLa varies greatly across the tree of life and even within individual genomes (Section 5). We also show that CaPLa is a good predictor of the space usage of a learned index data structure (the PLA-index [2]). We use controlled experiments to elucidate what makes genome spectra different from random $k$-mer multisets. Finally, we verify that CaPLa has some of the desired properties of a measure of genome complexity, namely robustness to variation in genome or $k$-mer size. Our work helps reduce the gap between theory and practice and aims to deepen our understanding of when and how to best apply PLA-based indexes to real genomic data.

In this section, we present the key definitions of PL-approximability and the motivation behind them. We start with a multiset $S$ of elements from an integer universe $[u]=\{0,1,\mathrm{\dots },u-1\}$. In the case of $k$-mers, $u=4^k$, though our definitions work for any $u$. We will use the notation that $n$ is the number distinct and $N$ is the total number of elements in $S$. The function we will be approximating is defined as ${\text{rank}}_S(x)=|\{y\in S\mid y\le x\}|$, for all $x\in [u]$. The key tool in designing PLA-based learned data structures [12, 11, 17, 2, 8] is the following:

Figure 1 shows an example of a PLA. The usefulness of a PLA comes from the fact that a data structure can avoid storing the full rank function and instead just store the PLA, as long as it has a way to handle the uncertainty due to the error of the approximation. To this end, the PLA with the smallest number of segments is the most space-efficient representation. Such a PLA can also be computed efficiently [27]. As a result, minimum-sized PLAs have formed the basis of several data structures [11, 8, 2, 21, 31, 19]. We formalize it with the following definition.

In this paper, we introduce a closely related measure, but one that factors out the effect of the data size.

The PL-approximability of $S$ can be interpreted as the longest average number of elements spanned by a segment in a PLA of $S$, so it directly quantifies the efficiency of a PLA in capturing the structure and distribution of the underlying data.

The space and time performance of learned data structures based on PLAs are often bounded in terms of the PLA-size [11, 2, 8]. Typically, smaller values of $\epsilon$ lead to faster query times, but require a larger number of segments $b(\epsilon )$ and thus increased space. The PLA-size and PL-approximability can then be tabulated for a given dataset and plugged into these space- and time-bounds to assess the resulting space-time trade-off. However, this approach provides little insight into broader trends and differences among data distributions. Ideally, the PL-approximability could be modeled as a parametrized family of functions, such that fitting this family to $S$ yields parameters that fully capture the PL-approximability of $S$.

A key difficulty is that the shape of the PLA-size function varies significantly depending on the structure and distribution of the underlying data. In general, the PLA-size is bounded between $b(\epsilon )=1$ and $b(\epsilon )=N/(2\epsilon )$ for all $\epsilon$ [11]. For sets where the gaps between elements are random and independently drawn, the PLA-size has been shown to be $b(\epsilon )=𝒪(N/{\epsilon }^2)$ [9]. For real data, the PLA-size rarely admits a closed-form expression, but power-law relations of the form $b(\epsilon )=n/(\beta {\epsilon }^{\alpha })$ have been empirically observed [9, 11]. Here, $\beta$ and $\alpha$ are dataset-specific constants found through fitting to the tabulated PLA-size. This power-law modeling provides a degree of generality, but when the PLA-size does not exactly follow a power-law, it has several drawbacks. Figure 2 demonstrates that this type of modeling 1) has a high variability in the error, 2) is not robust to the choice of fitting algorithm, and 3) is unpredictable in whether it over- or under-estimates the true PLA-size. Such drawbacks make it challenging to incorporate this type of model into downstream analyses of data structures, especially since its predictions are neither worst- nor average-case bounds.

To address this challenge, we propose bounding the PLA-size using two power laws instead of relying on a single fitted curve. We first fix a finite set $ℰ\subset {\mathbb{Z}}^{+}$ of $\epsilon$ values for which the PL-approximability is evaluated. This reflects practical applications, where $\epsilon$ is chosen from a set of values according to the desired space-time trade-off.

When ${\beta }_{\mathrm{low}}={\beta }_{\mathrm{high}}$, this is equivalent to a perfect power-law fit. The role of $\alpha$ is to capture the rate of segment growth as a function of $\epsilon$. If ${\beta }_{\mathrm{low}}={\beta }_{\mathrm{high}}$, then when $\epsilon$ doubles, the average segment length increases by a factor of $2^{\alpha }$. There is similarly an intuitive interpretation of ${\beta }_{\mathrm{low}}$ and ${\beta }_{\mathrm{high}}$: when $\epsilon =1$, the average segment length is between ${\beta }_{\mathrm{low}}$ and ${\beta }_{\mathrm{high}}$.

Definition 4 captures a full spectrum of data distributions and their PLA-size, from the best case $b(\epsilon )=1$ (with $\alpha =0$ and ${\beta }_{\mathrm{low}}={\beta }_{\mathrm{high}}=n$), to the worst case $b(\epsilon )=N/(2\epsilon )$ from [11] (with $\alpha =1$ and ${\beta }_{\mathrm{low}}={\beta }_{\mathrm{high}}=2n/N$), to the case of random gaps from [9] (with $\alpha =2$ and ${\beta }_{\mathrm{low}}={\beta }_{\mathrm{high}}=cn/N$ for a constant $c>0$). More importantly, it enables bounding the PL-approximability of real datasets where an exact power-law fitting is not possible.

The PL-approximability of a single dataset can be power-law bounded for infinitely many values of $\alpha ,{\beta }_{\mathrm{low}}$, and ${\beta }_{\mathrm{high}}$. To capture a single parameterization that can be used as a proxy for the PL-approximability of a dataset, we introduce the following notion of canonical PL-approximability.

We note that if $S$ has CaPLa $({\alpha }^{\ast },{\beta }_{\mathrm{low}}^{\ast },{\beta }_{\mathrm{high}}^{\ast })$, then Equation 1 immediately implies that $S$ is power-law bounded with parameters $({\alpha }^{\ast },{\beta }_{\mathrm{low}}^{\ast },{\beta }_{\mathrm{high}}^{\ast })$.

In this section, we show under which conditions the canonical PL-approximability exists and how it can be efficiently computed. We do so by introducing a geometric object, the twisted ribbon, which captures the entire family of valid power-law bounds for $S$. We then study the structural properties of this ribbon to identify a pinch point, which yields the canonical PL-approximability (Definition 5). Finally, we describe a two-phase algorithm that efficiently locates a pinch point based on the twisted ribbon’s geometry.

We first observe that, for any given $\alpha$, the tightest possible power-law bound on the PL-approximability of $S$ is given by choosing ${\beta }_{\mathrm{low}}$ and ${\beta }_{\mathrm{high}}$ as $L(\alpha )$ and $H(\alpha )$, respectively. The proof follows easily from the definitions of $L$ and $H$.

We can plot the functions $H$ and $L$ on a Cartesian plane with $\alpha$ along the horizontal axis and $\beta$ along the vertical axis. The region enclosed between these two curves forms the aforementioned twisted ribbon, depicted in Figure 3 and formally defined as follows.

Observe that choosing ${\beta }_{\mathrm{low}}$ or ${\beta }_{\mathrm{high}}$ inside of the twisted ribbon (that is, setting ${\beta }_{\mathrm{low}}>L(\alpha )$ or for some $\alpha$) would violate Equation 1, and therefore would not yield a valid bound on the PL-approximability of $S$. This reduces the space of relevant power-law bounds to that of ${\{(\alpha ,L(\alpha ),H(\alpha ))\}}_{\alpha \ge 0}$.

We further aim to summarize this entire family of bounds into a single representative instance. Intuitively, this means identifying a value of $\alpha$ that yields the tightest bound. This corresponds to a value ${\alpha }^{\ast }$ that minimizes the twisted ribbon width $W(\alpha )=H(\alpha )-L(\alpha )$, i.e., what we call a pinch point (Definition 5).

To study the existence and computation of a pinch point, we first show some structural properties of the twisted ribbon. For technical convenience, we assume that the (nonincreasing integer) function $b(\epsilon )$ has been extended to a continuous, nonincreasing, and differentiable function on the reals, e.g., via monotone cubic interpolation.

Note that the lemma holds for any dataset, as its conditions are related to a parameter ($ℰ$) chosen by the user. The following result shows that a pinch point exists and is confined between (or at) the flattening points mentioned in Lemma 8 and depicted in Figure 3.

We now focus on finding a pinch point, that is, on solving the problem ${argmin}_{\alpha \ge 0}W(\alpha )$. Observe $W$ is non-differentiable due to the presence of $\min$ and $\max$ operations. Moreover, we assume that $W$ is unimodal, namely, it has a unique minimum ${\alpha }^{\ast }$, is strictly decreasing for all , and strictly increasing for $\alpha >{\alpha }^{\ast }$. While we do not have any general theoretical guarantee on the unimodality of $W$ (except in simple cases such as those shown in the left and right plots of Figure 3), our empirical observation suggests that unimodality often holds in practice (see Section 5.5). If this is not the case, our algorithm finds a local minimum instead. We also assume $b(\epsilon )$ has been precomputed for all $\epsilon \in ℰ$, allowing $W(\alpha )$ to be evaluated in $𝒪(|ℰ|)$ time via a direct computation based on Equation 2. This precomputation can be done in $𝒪(N|ℰ|)$ time [27].

The algorithm works in two phases. In the first phase, we identify the flattening points ${\alpha }_L$ and ${\alpha }_H$ by observing from Lemma 8 that ${\alpha }_L$ is the smallest value of $\alpha$ such that $\frac{n}{{\epsilon }^{\alpha }b(\epsilon )}=\frac{n}{b(1)}$ for some $\epsilon \in ℰ\setminus \{1\}$. This equality simplifies to $\alpha ={\log }_{\epsilon }\frac{b(1)}{b(\epsilon )}$, so we obtain ${\alpha }_L={\min }_{\epsilon \in ℰ\setminus \{1\}}{\log }_{\epsilon }\frac{b(1)}{b(\epsilon )}$, which can be computed in $𝒪(|ℰ|)$ time. A similar derivation with $\max$ instead of $\min$ yields ${\alpha }_H$.

In the second phase, we run a derivative-free minimization algorithm (golden-section search) within the interval $[{\alpha }_L,{\alpha }_H]$, which is guaranteed to contain ${\alpha }^{\ast }$ due to Theorem 9. It converges in $\lceil {\log }_{\varphi }\frac{{\alpha }_H-{\alpha }_L}{\tau }\rceil$ iterations, where $\tau$ is the desired tolerance, and $\varphi \approx 1.618$ is the golden ratio [29]. Since the cost of each iteration is dominated by the evaluation of $W(\alpha )$, the overall time complexity of finding the pinch point is $𝒪(|ℰ|\log \frac{{\alpha }_H-{\alpha }_L}{\tau })$. This gives the following theorem.

Section 5.1 will show that the runtime of this algorithm is on the order of a few milliseconds in practice.

There are a number of PLA-based learned data structures whose space-time performance is bounded in terms of the PLA-size, such as the PGM-index [11, Theorem 1], LeMonHash [8, Theorem 2], and the PLA-index [2, Theorem 1]. In this section, we showcase how our PL-approximability framework can be used to reinterpret such bounds, focusing on the PLA-index as a concrete example particularly relevant in genomic applications. The PLA-index has been applied to substantially improve the run time of suffix array search and decrease the space of a read aligner and direct-access lookup tables [2].

The PLA-index uses a piecewise linear $\epsilon$-approximation to a multiset of $k$-mers $S$ to estimate the rank of a query $k$-mer in $S$ [2]. The tuning parameter $\epsilon$ represents a time/space trade-off, with smaller $\epsilon$ values resulting in faster index query times but more segments. Theorem 1 in [2] describes the size of the PLA-index as a function of the number of segments used by the PLA index (denoted by $s_{\epsilon }$), $N$, and $\epsilon$:

As $s_{\epsilon }$ can be high in the worst case, Equation 3 does not capture the fact that the space is very low in practice.

Here, we show that if $S$ has a power-law bounded PL-approximability, then Equation 3 can be turned into a useful predictor of PLA-index space usage. One of the technical difficulties is that PL-approximability gives bounds on the PLA size $b(\epsilon )$, whereas Equation 3 is stated in terms of $s_{\epsilon }$. Though the PLA-index uses a segmentation based on the optimal one, it requires consecutive segments to share an endpoint and in some rare cases breaks segments. This dependence of $s_{\epsilon }$ on $b(\epsilon )$ is difficult to quantify in terms of $S$. Instead, we isolate the dependence by expressing our results in terms of ${\gamma }_{\epsilon }\triangleq \frac{s_{\epsilon }}{b(\epsilon )}$. Though ${\gamma }_{\epsilon }$ is theoretically unknown (except that ${\gamma }_{\epsilon }\ge 1$), empirical results show that it is close to one on genome spectra. For example, for the 21-spectrum of chromosome 1 of the human genome, ${\gamma }_{\epsilon }$ values are below 1.21 for all tested $\epsilon$ up to $\epsilon =1023$.

We have isolated and combined some terms into a variable $C$, in order to convey that they can be thought of as constants in practice. First, $N/n$ is in practice smaller than 2 when avoiding small values of $k$, for which a suffix array would not make sense as an application anyway. Second, as we mentioned earlier, ${\gamma }_{\epsilon }$ is in practice a number close to 1.

Interpreting Eq. 4 is helped by the observation that given the type of values we see in practice, the dominating term is the leading one, i.e. $1/({\beta }_1{\epsilon }^{\alpha })$. This intuition is confirmed by Figure 4AB, which plots Eq. 4 using real numbers from human chromosome 1. As suggested by the leading term, increasing the $\alpha$ value of a dataset by $\delta$ results in the leading term decreasing by a factor of ${\epsilon }^{\delta }$, with the plots confirming this type of exponential decay for Eq. 4 overall. Note that for $\epsilon =63$ (Panel A), the decay is slower than for $\epsilon =1023$ (Panel B), as expected from the leading term. Figure 4CD shows how the gap between ${\beta }_{\mathrm{low}}$ and ${\beta }_{\mathrm{high}}$ affects the gap between the upper and lower bounds of Theorem 11 and gives the scale of the gap at which our bounds become meaningless.

The goal of our experiments is threefold. First, we want to understand how much PL-approximability varies among the tree of life. Second, we want to understand the properties of genome spectra that make them different from random $k$-mer multisets. Finally, we want to evaluate the suitability of CaPLa as a measure of PL-approximability. Our tool for computing CaPLa, as well as reproducibility information, is available through GitHub at https://github.com/medvedevgroup/CaPLa.

In this paper, we introduced the notion of PL-approximability of a multiset and a corresponding measure called CaPLa (Section 2). Our theoretical results show that CaPLa is well-defined and can be computed efficiently, under reasonable assumptions which we demonstrate hold in almost all tested data (Section 3). Our experimental results show that CaPLa varies substantially not only across genomes and kingdoms, but also within individual genomes, and this variability can lead to differences of up to 50% in the space required by a PLA-based data structure (Sections 5.2 and 5.3). Within our CaPLa framework, the key factors distinguishing a genome spectrum from a random $k$-mer multiset are the presence of repeats and the non-uniform distribution of $k$-mers. The combined effect of these two factors substantially reduces PL-approximability (Section 5.4). CaPLa has several desirable properties, namely, it is well-defined and efficient to compute across our experimented genomes, it is not trivially correlated with genome length thus indicating it captures the intrinsic properties of the genome, and it is robust with respect to $k$-mer size when $k\ge 21$ (Section 5.5).

Our result has a concrete application to parameter selection in PLA-based indexes such as the PLA-index. When a target space is specified, an approximate estimate of the dataset’s CaPLa can be used with Theorem 11 to compute a suitable value for the tuning parameter $\epsilon$. This is particularly advantageous when the CaPLa is known for a similar dataset (e.g., from another individual of the same species), as it avoids the need to recompute the full set of $(\epsilon ,b(\epsilon ))$ pairs. In the absence of prior CaPLa information, future work could also explore efficient estimation of CaPLa over a small sample of the dataset itself (e.g., a contiguous region of the suffix array).

Our paper presents several future directions. From an application perspective, it will be interesting to connect the variability of PL-approximability to biological properties. Our preliminary results show that PL-approximability is higher in fungi and lower in repeat-rich chromosomes like the Y, but this raises more questions than it answers. From a theory perspective, it would be interesting to find the conditions under which the pinch point is unique. In our experiments, we could not find any instances of non-unique pinch points, suggesting that a theoretical proof may be possible. Finally, Theorem 9 requires that $1\in ℰ$ in order to guarantee the existence of a pinch point. This condition is necessary, as in fact one can prove that a pinch point does not exist if $1\notin ℰ$. This raises an interesting open question if there is some alternative notion of CaPLa that can be used when $1\notin ℰ$.