1 of 15

OAK: Enriching Doc Representations using Auxiliary Knowledge for XC

Shikhar Mohan, Deepak Saini, Anshul Mittal, Sayak Ray Chowdhury, Bhawna Paliwal, Jian Jiao, Manish Gupta, Manik Varma.

​

ICML 2024

gmanish@microsoft.com

1

07-Jun-24

2 of 15

What is XC, and what was missing in XC methods?

  • Predicting the most relevant subset of labels for a data point from an extremely large set of labels.
  • Applications
    • Product recommendation
    • Document tagging
    • Search and ads
    • Query recommendation
  • Sparse representations of short-text XC datasets is bad!
    • User queries or products/webpage titles
    • LF-WikiTitles500K: 3-5 words

​

OAK: Enriching Document Representations using Auxiliary Knowledge for Extreme Classification. Shikhar Mohan, Deepak Saini, Anshul Mittal, Sayak Ray Chowdhury, Bhawna Paliwal, Jian Jiao, Manish Gupta, Manik Varma. ICML. 2024.

  • Why not use auxiliary info?
    • Frequently clicked webpages for search queries in sponsored search
    • Previously searched queries for web search query auto-completion
  • WikiTitle 🡪 Related (SeeAlso) Wiki pages
    • AK: relevant categories

gmanish@microsoft.com

2

07-Jun-24

3 of 15

How do you use auxiliary knowledge pieces (AKPs)?

  •  

OAK: Enriching Document Representations using Auxiliary Knowledge for Extreme Classification. Shikhar Mohan, Deepak Saini, Anshul Mittal, Sayak Ray Chowdhury, Bhawna Paliwal, Jian Jiao, Manish Gupta, Manik Varma. ICML. 2024.

  • Approach 2
    • Use GNN methods like GraphFormers and GraphSAGE to encode the document-AKPs linkage info.
    • Problems
      • High storage and computational costs
      • Cannot leverage the auxiliary data sourced from disparate tasks effectively.

gmanish@microsoft.com

3

07-Jun-24

4 of 15

How does OAK (Online Auxiliary Knowledge) work?

  • 3 stage OAK training
    • Linker training
      • Link documents with relevant AKPs using an XC method
    • OAK pretraining
      • Combine the document embeddings with relevant trainable AKP embeddings via attention-based pooling to get an enriched doc rep.
      • Train augmentation block in a Siamese fashion.
      • Encoder regularization
      • Calibration regularization term inspired by DPO.
    • OAK finetuning
      • Freeze augmentation block parameters.
      • Learn a per-label refinement vector to fine-tune the label embeddings to obtain the final label classifiers.

OAK: Enriching Document Representations using Auxiliary Knowledge for Extreme Classification. Shikhar Mohan, Deepak Saini, Anshul Mittal, Sayak Ray Chowdhury, Bhawna Paliwal, Jian Jiao, Manish Gupta, Manik Varma. ICML. 2024.

gmanish@microsoft.com

4

07-Jun-24

5 of 15

What is OAK’s Architecture?

  •  

OAK: Enriching Document Representations using Auxiliary Knowledge for Extreme Classification. Shikhar Mohan, Deepak Saini, Anshul Mittal, Sayak Ray Chowdhury, Bhawna Paliwal, Jian Jiao, Manish Gupta, Manik Varma. ICML. 2024.

gmanish@microsoft.com

5

07-Jun-24

6 of 15

Training the augmentation block

  •  

OAK: Enriching Document Representations using Auxiliary Knowledge for Extreme Classification. Shikhar Mohan, Deepak Saini, Anshul Mittal, Sayak Ray Chowdhury, Bhawna Paliwal, Jian Jiao, Manish Gupta, Manik Varma. ICML. 2024.

gmanish@microsoft.com

6

07-Jun-24

7 of 15

Training the augmentation block

  •  

OAK: Enriching Document Representations using Auxiliary Knowledge for Extreme Classification. Shikhar Mohan, Deepak Saini, Anshul Mittal, Sayak Ray Chowdhury, Bhawna Paliwal, Jian Jiao, Manish Gupta, Manik Varma. ICML. 2024.

gmanish@microsoft.com

7

07-Jun-24

8 of 15

Inference using OAK

  •  

OAK: Enriching Document Representations using Auxiliary Knowledge for Extreme Classification. Shikhar Mohan, Deepak Saini, Anshul Mittal, Sayak Ray Chowdhury, Bhawna Paliwal, Jian Jiao, Manish Gupta, Manik Varma. ICML. 2024.

gmanish@microsoft.com

8

07-Jun-24

9 of 15

Results

  • OAK offers higher P@1 on standard XC benchmark datasets.
  • ∼2% in P@1 over NGAME
  • Compared to GraphFormer, OAK demonstrates substantial gains in accuracy, over 15-20% and compared to GraphSAGE, 5-7% higher P@1 across all relevant datasets.

​

OAK: Enriching Document Representations using Auxiliary Knowledge for Extreme Classification. Shikhar Mohan, Deepak Saini, Anshul Mittal, Sayak Ray Chowdhury, Bhawna Paliwal, Jian Jiao, Manish Gupta, Manik Varma. ICML. 2024.

  • OAK offers 5% higher P@1 on LF-ORCAS-800K.

gmanish@microsoft.com

9

07-Jun-24

10 of 15

Ablations

Ablation for AKP representation

OAK: Enriching Document Representations using Auxiliary Knowledge for Extreme Classification. Shikhar Mohan, Deepak Saini, Anshul Mittal, Sayak Ray Chowdhury, Bhawna Paliwal, Jian Jiao, Manish Gupta, Manik Varma. ICML. 2024.

gmanish@microsoft.com

10

07-Jun-24

Ablation for combiner

Ablation for regularisation loss

Ablation for Early concatenation vs late attentive fusion (LF-WikiSeeAlsoTitles-320K dataset).

11 of 15

Performance on Tail

  • LF-WikiSeeAlsoTitles-320K
  • Increasing document/label frequency quantiles

OAK: Enriching Document Representations using Auxiliary Knowledge for Extreme Classification. Shikhar Mohan, Deepak Saini, Anshul Mittal, Sayak Ray Chowdhury, Bhawna Paliwal, Jian Jiao, Manish Gupta, Manik Varma. ICML. 2024.

  • OAK consistently outperforms NGAME and GraphFormer, and is able to predict precisely for rare documents as well as rare labels.

gmanish@microsoft.com

11

07-Jun-24

12 of 15

Results using Oracle Linker AKPs

  • For short text datasets, results are better when using ground truth AKPs.

OAK: Enriching Document Representations using Auxiliary Knowledge for Extreme Classification. Shikhar Mohan, Deepak Saini, Anshul Mittal, Sayak Ray Chowdhury, Bhawna Paliwal, Jian Jiao, Manish Gupta, Manik Varma. ICML. 2024.

  • For full text datasets, where trained linker quality is better than for short-text datasets, results with ground truth AKPs are indeed lower compared to ones with predicted AKPs.

gmanish@microsoft.com

12

07-Jun-24

13 of 15

Application to Sponsored search

  • Match queries to billions of bid keywords.
  • Auxiliary data: webpages clicked.

​

​

​

​

  • Offline results: OAK categorically performs better than proprietary variations of leading dense retrieval algorithms deployed in production.

​

​

OAK: Enriching Document Representations using Auxiliary Knowledge for Extreme Classification. Shikhar Mohan, Deepak Saini, Anshul Mittal, Sayak Ray Chowdhury, Bhawna Paliwal, Jian Jiao, Manish Gupta, Manik Varma. ICML. 2024.

  • Online
    • CTR: Clicks/impressions
    • Impression Yield (IY) and Click Yield (CY): #ad impressions and clicks per user query search.
    • Keyword Density (KD): Fraction of predicted keywords passing relevance filters.
  • +0.84% CTR, +2.7% KD

gmanish@microsoft.com

13

07-Jun-24

14 of 15

Sample predictions from OAK, NGAME and GraphFormer

OAK: Enriching Document Representations using Auxiliary Knowledge for Extreme Classification. Shikhar Mohan, Deepak Saini, Anshul Mittal, Sayak Ray Chowdhury, Bhawna Paliwal, Jian Jiao, Manish Gupta, Manik Varma. ICML. 2024.

gmanish@microsoft.com

14

07-Jun-24

15 of 15

Summary

  • Use of AKPs for diverse and accurate XC.
  • Modular OAK architecture
    • Joint training of AKP reps on top of any encoder
    • 3-stage training scheme: linker training, OAK pre-training and OAK finetuning
    • Mutual information calibration loss
  • SOTA results on several XC tasks
    • Ad keyword prediction for user queries
    • Wikipedia categories and “See Also” prediction
    • Webpage prediction
  • OAK’s inference takes <10ms enabling it to be deployed for a sponsored search query to keyword prediction task.

OAK: Enriching Document Representations using Auxiliary Knowledge for Extreme Classification. Shikhar Mohan, Deepak Saini, Anshul Mittal, Sayak Ray Chowdhury, Bhawna Paliwal, Jian Jiao, Manish Gupta, Manik Varma. ICML. 2024.

​

gmanish@microsoft.com

15

07-Jun-24