Top-k selection retains the k highest-scoring candidates. In sparse mixture-of-experts routing, a learned projection scores each token against experts, selects the strongest experts, and normalizes their scores with softmax. section 3.5.1
A practical top-k routing pipeline typically performs these steps:
- Compute candidate scores with a linear projection. section 3.5.1
- Select the k largest scores and record their expert indices. section 3.5.1
- Apply softmax only to the selected scores to obtain routing weights. section 3.5.1
- Add capacity limits, load-balancing losses, or noisy routing when avoiding expert collapse matters. section 3.5.1
The same selection pattern appears in retrieval-augmented generation: retrieve the top-k matching documents, then optionally rerank or compress them before passing context to the language model. Increasing k can provide broader coverage, while reranking or compression filters irrelevant material. section 5.4.3
Sparse gating
A PyTorch gate calls torch.topk(gate, k=top_k, dim=-1, largest=True), reshapes the selected values, and applies softmax across the selected experts. A GShard-style variant fixes k at two and adds randomized routing, capacity constraints, and load balancing. section 3.5.1
Listing 3.8 Top-k gating function for sparse MoE
class NativeGate(BaseGate): #1
def __init__(self, d_model, num_expert, world_size, top_k=2, gate_bias=True):
super().__init__(num_expert, world_size) #2
self.gate = nn.Linear(d_model, self.tot_expert, bias=gate_bias)
self.top_k = top_k
def forward(self, inp, return_all_scores=False):
gate = self.gate(inp) #3
gate_top_k_val, gate_top_k_idx = torch.topk( #4
gate, k=self.top_k, dim=-1, largest=True, sorted=False
)
gate_top_k_val = gate_top_k_val.view(-1, self.top_k) #5
gate_score = F.softmax(gate_top_k_val, dim=-1)
self.set_loss(torch.zeros(1, requires_grad=True).to(inp.device))
if return_all_scores: #6
return gate_top_k_idx, gate_score, gate
return gate_top_k_idx, gate_score
Inspect the annotated sparse-MoE gate implementation, including score computation, top-k indexing, and selected-score normalization. section 3.5.1
Figure 5.2 Overview of a RAG system. In the preproduction phase, documents from a knowledge base are chunked and embedded into a vector database. During the production phase, a user query is embedded and used to retrieve top-matching documents based on similarity. An optional reranker refines these results to improve relevance. The final ranked context is then passed, together with the original query, to the language model for response generation.
Trace how retrieval top-k feeds reranking, contextual compression, and final language-model context selection. section 5.4.3
