Trainable Semantic Indexes (TASTI)

• Trainable Semantic Indexes (TASTI) are embedding-driven frameworks that encode unstructured data into a shared latent space for efficient ML-style querying.
• The approach integrates embedding learning, representative clustering, and index storage to drastically reduce annotation overhead and accelerate query response.
• TASTI provides theoretical guarantees on query accuracy and scalability, validated across diverse applications like video, text, and audio analytics.

Trainable Semantic Indexes (TASTI) are a class of machine learning infrastructure that enables efficient ML-centric querying over unstructured data—such as video, text, and audio—by replacing per-query proxy models with a reusable, embedding-driven semantic index. TASTI frameworks leverage the concept that semantically similar instances, as defined by user-driven metrics or learning objectives, can be encoded into a shared vector or code-based latent space and clustered to facilitate efficient retrieval and large-scale analytics. By jointly optimizing representation learning, storage, and inference objectives, TASTI methods enable substantial cost and latency reductions relative to prior proxy-based approaches, advance the scalability of ML-based querying, and provide theoretical guarantees on downstream query performance (Kang et al., 2020, Jin et al., 2023, Jiang et al., 2022).

1. Motivation and Problem Context

Unstructured data pipelines such as those for surveillance video analysis, large-scale text mining, and speech processing frequently depend on computationally costly target labelers—e.g., high-capacity neural networks or human annotators—to extract semantic attributes. Classical acceleration approaches (e.g., BlazeIt, NoScope, SUPG) rely on query-specific proxy models trained to approximate these target labelers' outputs, incurring high annotation costs, retraining overhead for each new query predicate, and poor cross-query computation reuse.

TASTI generalizes across queries by constructing a single embedding-based index that efficiently supports selection predicates, approximate aggregation, and rank/limit queries. It leverages semantic similarity within the dataset so that, after initial indexing and minimal labeling, any number of subsequent ML-style queries can be answered with no further proxy model training or large-scale annotation (Kang et al., 2020).

2. Index Construction and Embedding Architecture

TASTI’s canonical instantiation involves three principal steps:

1. Embedding Learning: Each data record $x_i$ is embedded into a $d$-dimensional latent space $\varphi(x_i)$, typically via a pretrained DNN backbone (e.g., ResNet-18 for images, BERT for text, PANNs for audio). Embeddings are optionally finetuned using a triplet-loss head, enforcing that semantically close pairs (according to a user-supplied proximity $d$ in label space) are closer in representation than far pairs.
2. Representative Selection and Clustering: Furthest Point First (FPF) selection identifies $N_2$ cluster representatives in the embedding space, providing a 2-approximation to minimum max intra-cluster radius.
3. Index Storage: Each record stores distances to its $k$ nearest cluster representatives. The index comprises:
   • A table mapping each cluster representative $r_j$ to cached target labeler output $f(r_j)$,
   • A matrix of $\mathrm{MinKDistances}[i] \in \mathbb{R}^k$ for each record.

Index construction cost is $O(C\,c_T + L\,c_E + N\,c_E + N C D\,c_D)$, where $d$0 is the annotation budget, $d$1 the cost of invoking the target labeler, and $d$2, $d$3 the costs of embedding and distance computations, respectively (Kang et al., 2020).

Alternative TASTI instantiations use discrete semantic IDs generated by sequential neural coding (e.g., LMIndexer). Here, a transformer backbone produces per-document codes via progressive, autoregressive selection from codebooks, with self-supervised reconstruction losses and hierarchical coarse-to-fine semantics enforced during training (Jin et al., 2023).

3. Query-Time Scoring and Proxy Score Propagation

At inference, TASTI eschews per-query model training. Instead, proxy scores $d$4 for each record $d$5 are generated by propagating the cached outputs of the nearest cluster representatives:

• Numeric scores:

$d$6

where $d$7 are the distances to the nearest $d$8 clusters.

• Categorical or binary scores apply weighted voting across the cluster representatives.

These proxy scores are consumable by existing downstream algorithms designed for selection, aggregation, or limit/rank queries (e.g., control-variate sampling in BlazeIt, recall estimation in SUPG). Query time is $d$9 per query, with no need for further annotation or retraining unless index refinement (“cracking”) is desired (Kang et al., 2020).

4. Theoretical Analysis and Guarantees

TASTI provides theoretical guarantees relating embedding quality to query accuracy. Under compactness of the data domain and Lipschitz continuity of queries, the expected error between proxy outputs $\varphi(x_i)$0 (using the nearest cluster) and the true target labeler output $\varphi(x_i)$1 is controlled by the cluster density and the embedding triplet loss:

• Zero-Loss Case: If $\varphi(x_i)$2 (perfect embedding), then

$\varphi(x_i)$3

where $\varphi(x_i)$4 is the semantic cluster radius and $\varphi(x_i)$5 the Lipschitz constant of loss $\varphi(x_i)$6.

• Nonzero-Loss Case: If $\varphi(x_i)$7, then the error admits a similar additive bound involving $\varphi(x_i)$8, $\varphi(x_i)$9, and the supremum of cluster sizes.

Conditions for these bounds include embedding respect for the user metric $d$0 (low triplet loss), dense clustering, and query Lipschitzness. No specific distributional assumptions are required (Kang et al., 2020).

5. Empirical Evaluation and Performance

TASTI has been evaluated on video (e.g., Mask R-CNN targets, ResNet-18 embeddings), text (WikiSQL, BERT), and speech datasets (Common Voice, PANNs). The following outcomes have been established:

• Indexing Cost: Up to $d$1 reduction in target labeler invocations compared to BlazeIt's proxy model construction, leading to $d$2 faster index construction (Kang et al., 2020).
• Query Acceleration: Aggregation queries use half as many labels to meet equivalent error, selection queries demonstrate up to $d$3 lower false positive rates versus per-query proxies, and limit/rank queries require $d$4 fewer target labeler calls.
• Semantic ID Quality: In LMIndexer, semantic ID quality (AMI vs. ground-truth) reaches $d$5–$d$6 compared to $d$7–$d$8 for two-stage baselines; document retrieval recall@10 outperforms dual-encoder baselines on large-scale retrieval tasks (Jin et al., 2023).

Critical ablations confirm that both embedding finetuning (triplet loss, reconstruction objectives) and the clustering method (FPF, codebook warmup) are essential for optimal performance, especially for rare-event or limit queries.

6. Advanced Indexing: End-to-End and Structured Codes

Recent extensions to TASTI encompass end-to-end differentiable indexing via product quantization (PQ) with learnable rotations. Block Givens coordinate descent (GCD-R, GCD-G, and GCD-S) replaces standard SVD-based OPQ/PQ index rotation learning, enabling parallelizable, stable rotation learning on GPU. In this setting, the quantized index layer $d$9 and the embedding network are co-trained, with straight-through estimators handling non-differentiable PQ steps. This avoids $N_2$0 recall loss from post-training quantization and allows seamless scaling to billions of items using IVF+PQ hybrids (Jiang et al., 2022).

In the discrete coding variant, TASTI is realized as a sequential token indexer (e.g., LMIndexer), generating interpretable, hierarchical semantic IDs using self-supervised document reconstruction, contrastive, and commitment losses. The code assigns each instance a sequence of discrete symbols with coarse-to-fine granularity, facilitating efficient retrieval and interpretability (Jin et al., 2023).

7. Limitations and Future Directions

TASTI's performance depends on the representational capacity of the embedding function and the semantic suitability of the user-supplied closeness metric. Quality guarantees assume queries are Lipschitz and the embedding is high fidelity with respect to the semantic metric. Triplet-loss and cluster selection methods are pivotal for rare-event queries. Hyperparameter sensitivity to cluster count, neighbor count $N_2$1, and embedding dimension is modest but non-negligible.

Potential extensions include adaptive ID length per record, dynamic codebook sizing or tree pruning, multimodal indexes combining text and image features, and integration of supervised or contrastive signals from behavioral logs. Application domains span large-scale recommendation, retrieval, and analytics for any unstructured data store requiring semantic access patterns without incurring per-query training costs (Jin et al., 2023, Kang et al., 2020).