[Paper Review] Cascaded Fast and Slow Models for Efficient Semantic Code Search

Akhilesh Deepak Gotmare, Junnan Li+2 more|arXiv (Cornell University)|Oct 15, 2021

Software Engineering Research24 references7 citations

TL;DR

This paper proposes CasCode, a cascaded fast and slow transformer-based framework for efficient and accurate semantic code search. It first uses a fast encoder to retrieve top-K code candidates via dense vector similarity, then re-ranks them with a slow, joint classifier model for improved accuracy, achieving a new state-of-the-art MRR of 0.7795 on CodeSearchNet across six languages.

ABSTRACT

The goal of natural language semantic code search is to retrieve a semantically relevant code snippet from a fixed set of candidates using a natural language query. Existing approaches are neither effective nor efficient enough towards a practical semantic code search system. In this paper, we propose an efficient and accurate semantic code search framework with cascaded fast and slow models, in which a fast transformer encoder model is learned to optimize a scalable index for fast retrieval followed by learning a slow classification-based re-ranking model to improve the performance of the top K results from the fast retrieval. To further reduce the high memory cost of deploying two separate models in practice, we propose to jointly train the fast and slow model based on a single transformer encoder with shared parameters. The proposed cascaded approach is not only efficient and scalable, but also achieves state-of-the-art results with an average mean reciprocal ranking (MRR) score of 0.7795 (across 6 programming languages) as opposed to the previous state-of-the-art result of 0.713 MRR on the CodeSearchNet benchmark.

Motivation & Objective

• Address the inefficiency and suboptimal performance of existing semantic code search systems in practical deployment scenarios.
• Overcome the high computational cost of joint classification models when applied to large-scale code repositories.
• Balance speed and accuracy in semantic code search by combining fast retrieval with high-accuracy re-ranking.
• Reduce memory overhead in production by jointly training fast and slow components using shared parameters.
• Improve retrieval performance on the CodeSearchNet benchmark beyond prior state-of-the-art methods.

Proposed method

• Use a fast transformer encoder to generate dense semantic embeddings for natural language queries and code snippets, enabling efficient similarity-based retrieval via a precomputed vector index.
• Apply a separate, slow binary classifier model that jointly encodes query-code pairs to predict matching probabilities, improving accuracy over simple similarity metrics.
• Implement a cascaded inference pipeline where the top-K candidates from the fast retrieval stage are re-ranked by the slow classifier.
• Propose a shared-parameter variant where a single transformer encoder serves both the fast and slow stages through multi-task training with contrastive (infoNCE) and binary cross-entropy objectives.
• Construct a precomputed vector index (PL index) offline to accelerate online inference, excluding this cost from runtime measurements.
• Process the top-K candidates from the fast stage in a batched manner during the second stage to optimize inference efficiency.

Experimental results

Research questions

• RQ1 Can a cascaded architecture combining fast retrieval and slow re-ranking improve both efficiency and accuracy in semantic code search?
• RQ2 To what extent can shared parameter training reduce memory costs without sacrificing performance in a dual-stage code search system?
• RQ3 How does the performance of a joint classifier model compare to a simple embedding-based retrieval model in terms of MRR and inference speed?
• RQ4 What is the trade-off between inference speed and retrieval accuracy when varying the number of top candidates (K) processed in the re-ranking stage?
• RQ5 Why does multi-task training of the shared-parameter model underperform compared to training separate fast and slow models?

Key findings

• CasCode achieves a new state-of-the-art mean reciprocal ranking (MRR) score of 0.7795 on the CodeSearchNet benchmark across six programming languages, significantly outperforming the previous SOTA of 0.713.
• The separate-model variant of CasCode achieves an MRR of 0.7825 with an inference duration of 0.2883 seconds per query, processing 3.46 queries per second on a single A100 GPU.
• The shared-parameter variant reduces model size by half while maintaining strong performance, achieving an MRR of 0.7686 with an inference duration of 0.2956 seconds per query.
• Reducing the number of candidates processed in the re-ranking stage from K=100 to K=10 improves inference speed by nearly 3x (from 3.46 to 9.78 queries per second) with only a minor drop in MRR (0.7825 to 0.7724).
• The shared-parameter model, despite using half the parameters, incurs nearly identical inference cost to the separate model, demonstrating efficient memory utilization.
• Multi-task training of the shared model leads to slightly worse performance than separately fine-tuned models, indicating a challenge in balancing the joint optimization of contrastive and classification objectives.