I improved TikTok Music Search!

During my internship with TikTok's Search Algorithms team, I worked on improving music search relevance.

Background

There was an existing search relevance model in production, however that model was trained using documents from several verticals, while music-specific examples were relatively sparse and contained only a few textual fields. This resulted in a few categories of cases where the model failed to yield relevant results.

Music queries contain structure that does not appear as frequently in other search domains. Users search not only for song and artist names, but also for versions, lyrics, moods, years, genres and increasingly vague descriptions of songs.

For example, queries could look like:

sad taylor swift song"
"as it was acoustic"
"espresso sped up"
"slowed version"
"2016 remix songs

Benchmarking the baseline model

Before we start to figure out how to train a model better than the current one, I needed to first analyse systematically where the current model fails, and how can I improve the training data sampling and distribution to mitigate these shortcomings.

A few patterns emerged from this. (not saying everything, just the more interesting ones here)

1. Version attributes were interpreted as ordinary keywords— the model didn't pay attention to specific version attributes like sped up, live, instrumental.
2. Semantic intent combined with specific attributes weren't yielding good results— examples include sad taylor swift song.

How the training data looks like

The training data is usually just a query + $X$ number of documents, where $X$ is between 10 to 50.

Training data curation

Generation

For a task like this, training data distribution is often the most imperative aspect which can make or break model performance. This is a supervised training task and the relevance_score acts as the labels, from a discrete scale of 0-3.

To create the training data, I gathered the past ~7 days of search logs from our HDFS cluster, formatted it nicely and created a pipeline to get a frontier LLM to label all these documents. This was a challenge as we had to first decide the requirements for the different labels and to also formulate a prompt that the LLM will adhere closely to. We did strict evaluations on sample data before sending the whole dataset to the LLM for labelling.

Data Augmentation

Additionally, we did synthetic data augmentation— getting an LLM to rewrite some queries according to the categories of bad cases we've seen earlier. Rather than asking the LLM to produce arbitrary queries, mutations were organised around specific categories of search behaviour. While doing this, we also had to ensure the training data distribution remains consistent (i.e. every category of query augmentation should be represented equally). However, not every augmentation is reasonable or realistic for every query and LLMs have their own generation biases.

I therefore used query click entropy as a coarse estimate of query ambiguity. A low-entropy query typically concentrates user clicks on a small number of results and therefore behaves like a specific song or artist search. A high-entropy query has a more diffuse click distribution and is more likely to represent broad semantic intent such as genre, mood or vibe. This provided a useful routing signal for synthetic augmentation.

The pipeline also tracked the distribution of generated categories in real time.

Suppose the target distribution contains: $C_1,C_2,\ldots,C_k$ and one category becomes underrepresented. The generation prompt is then steered towards that category on subsequent calls. Conceptually: $P(C_i \mid \text{next generation}) \propto \text{target}_i-\text{observed}_i$ within each entropy tier.

Other sampling strategies

The initial training data gathered from the search logs was dominated by mid-length queries (3–6 tokens), underrepresenting two important query types with distinct failure patterns:

• Short queries (fewer than 3 words): Short and often generic queries (e.g. "house music", "music", single artist names). These were oversampled to improve the model's ability to handle ambiguous, low-specificity queries.
• Long-tail queries (more than 7 tokens): Highly specific queries often containing artist name, song title, version modifier, lyrical search, and format tokens simultaneously (e.g. "harry styles as it was acoustic version full song").

Disagreement-based sampling was also done. This was based on the search logs; relevance is only 1 of 3 targets used to determine rank, one other target used is Click-Through-Rate prediction.

• High rank, low relevance disagreement: Documents ranked highly in search results (top-5 positions) but assigned low relevance scores by the baseline model.
• Low rank, high relevance disagreement: Documents ranked lowly in search results but assigned high relevance scores by the baseline model.

Training

The model was trained as a cross-encoder. Instead of separately encoding the query and document, all textual fields were concatenated into one sequence.

The full sequence is processed jointly: $h= \operatorname{Encoder}(q,d)$ and the classification representation (the CLS token) is mapped to a scalar: $s(q,d) = \sigma(Wh_{\text{CLS}}+b)$.

Training procedure

Step 1: Fine-tuning

First, we fine-tuned a 24-layer BERT cross encoder model using a pretrained checkpoint. The output of this model will be a continuous scalar, between 0 and 1, representing how relevant the document is to the query. We used shrinkage-weighted MSE loss to optimise for the score and listwise pairwise margin loss to optimise for ranking. This resulted in a 20+% improvement in offline NDCG compared to the previous model (the one in production).

Step 2: Knowledge Distillation

To cut serving latency in production and reduce computational power required, I embarked on a project to train up a 6-layer model that can match the 24-layer model I've already trained up. I used offline knowledge distillation to achieve this.

The original dataset contained hard, discrete relevance labels from 0 to 3. Knowledge distillation provides a richer supervision signal by training the student on the teacher’s continuous relevance scores and, more importantly, the relative score margins between documents for the same query. This captures information that is lost in the original coarse labels. It also allows us to use a lot more training examples to train up the 6-layer model.

Step 3: Mixture-Of-Experts (MoE) Training

We needed an MoE mechanism as the BERT model we're training is a ranking model intended for different types of documents, so a single shared set of feed-forward weights has to represent scoring behaviour for content categories that likely look similar on the surface.

Stage 2, using distillation, makes a dense 6-layer model. Stage 3 turns it into an MoE via sparse upcycling: rather than initialising experts randomly (notoriously unstable), I copied the proven dense FFN into every expert slot.

Training at scale

Training ran as multi-node, multi-GPU distributed data-parallel (DDP), on 8 NVIDIA A100 chips. Each GPU holds a full copy of the model and processes a different slice of the global batch, so the effective batch size is local_batch × world_size. After each backward pass, every GPU has different gradients (due to each GPU handling different data), so they will be averaged across all GPUs before the optimiser step (all-reduce).

One of the main bottlenecks of training at scale is usually getting the data fast enough to the nodes.

$\text{overall throughput} = \min( \text{how fast GPUs can train}, \text{how fast batches arrive} )$

Some techniques we employed include:

Compiled preprocessing

A typical HuggingFace-style pipeline might perform tokenisation and feature construction in ordinary Python, for example through Dataset.map() or a Python collate_fn. At scale, that introduces interpreter overhead, GIL contention and object-allocation overhead. Hence, CPU operations were mostly moved out of interpreted Python and JIT-compiled with MatxScript. This reduced Python/GIL overhead and helped prevent GPUs from idling while waiting for batches.

Parallel data loading and prefetching

Multiple CPU workers prepared future batches in parallel while the GPU trained on the current batch: $\text{GPU: batch }N; \text{CPU: batch }N+1,N+2$. This overlaps preprocessing with GPU computation.

Sharded datasets

Training data was partitioned across DDP workers so that each rank read a different subset of examples, reducing duplicated I/O and unnecessary contention.

Overlap communication with backward computation

DDP can begin reducing gradients for earlier layers while later layers are still performing backpropagation, hiding part of the all-reduce cost behind useful computation.

Evaluation

We used offline testing using evaluation datasets with NDCG as a metric, as well as online A/B testing. Through the evaluation process, we systematically analysed positive and negative cases of the newly trained model, relative to the one in production. We further improved model performance by doing hard negative mining and including it in the training set.

One example of a hard negative we found during the initial training runs.

We got an LLM to include this as one of the categories for query rewriting, and synthetically added some training examples of this particular category.