Natural Language Processing Questions

Problem & goals

• Softmax over 1M tokens during inference must be memory- and compute-efficient while producing calibrated probabilities for candidate scoring.

Approaches

1. Sampled softmax (candidate sampling): compute full logits only for a small candidate set (beam or top-k) plus sampled negatives. Good for training; for inference needs careful sampling to avoid bias. Use importance-weighted correction if using sampled negatives.

2. Hierarchical softmax: arrange vocabulary in tree (e.g., binary Huffman-like). Inference cost O(log V) per token. Numerically stable if using log-sum-exp on short paths. Works well if vocabulary has skewed frequencies and tree optimized by frequency.

3. Candidate caching & shortlist: maintain per-context or global shortlists of frequent/cached candidates; compute exact logits for shortlist; backoff to fallback scoring (approximate) for others. For next-token, dynamic shortlists from n-gram/backoff LM or ANN search in embedding space.

4. Approximate nearest neighbor (ANN): precompute token embeddings; for a given context vector, use ANN to find top-k likely tokens rapidly and compute softmax over those. Use re-ranking for top candidates.

Numerical stability

• Compute softmax via log-sum-exp on chosen subset. For hierarchical softmax compute path logits and use stable accumulation.

Trade-offs

• Hierarchical: low latency, exact marginalization along tree but more complex to train and adapt. Sampled/ANN: simpler and scalable, but approximate and may hurt recall of rare tokens.
• Candidate caching + exact softmax on shortlist often practical: small shortlist (e.g., 1k) yields good latency/quality tradeoff.

Recommendation

• For production inference with 1M vocab: use hybrid: shortlist via ANN or frequency+context features, compute exact softmax over shortlist with log-sum-exp stabilization, and fall back to hierarchical or ANN scoring for low-probability tails. Monitor calibration on held-out real queries and tune shortlist size vs latency.