This is a collection of my notes which I refer to on a regular basis. Hope it is also helpful for others stumbling by.

I think of these notes as mountaineering pegs. Often at the time of studying a particular paper or topic, the concepts are clear. Over time, however, only a vague impression remains, and I can no longer tussle with the issues. So these notes serve as pegs that I hope to use to re-scale an old hill, or at least scale it with less effort than starting from scratch.

Generally, I highlight main points like so, and put in-line code or numbers like so.

Built with mdBook.

LightGBM Memory

TLDR: Solutions for memory issues during training of a LightGBM model:

1. Cast numeric values into np.float32 to save data space
2. Keep num_leaves <= 100 (or some reasonable number)
3. If feature dimension is large (e.g. M >= 1000), try colsample_bytree = 0.1, although this might not help too much if the bottleneck is during bin histogram construction (rather than the actual training)
4. If number of rows and features are both large (e.g. N >= 1_000_000 and M >= 1000, i.e. >= 4 GB) then the data itself is taking up a lot of memory. It would be worthwhile to put the data on disk and use lgb.Dataset by providing the file path as the data argument instead. Then, we should set two_round=True for the train method params. The explanation for two round is rather unclear, but it should help with memory when Dataset is loading from disk (rather than from a numpy.array in memory). For this option, I had some trouble getting it to work with categorical columns.

For more details can refer to the experiments below.

Experiments

I often run into memory issues running LightGBM. So here are some experiments to measure memory usage and understand how hyperparameters can affect memory usage.

The function of interest is the fit method for the learn to rank task.

import lightGBM as lgb
def f():
    model = lgb.LGBMRanker(**params, objective="lambdarank")
    model.fit(
        X=data,
        y=y,
        group=groups,
    )

The memory usage is measured using the memory_profiler module, which checks the memory usage at .1 second intervals. The maximum is then taken to represent the maximum memory usage of the fit function. We also take note of the size of the data itself (using data.nbytes) and subtract that away to get closer to the LightGBM memory usage. Do note that this memory profiling is not very rigorous, so the results are best for relative comparison within each experiment rather than across experiments.

from memory_profiler import memory_usage
def run(params):
    mem_usage = memory_usage(f)
    return max(mem_usage) / 1000 # GB

We set the default parameters as follows and generate the data this way. For the experiments below, the default parameters are used unless specified otherwise.

DEFAULT_PARAMS = {
    "N": 200000, # number of instances
    "M": 500, # feature dimension
    "n_estimators": 100,
    "num_leaves": 100,
    "histogram_pool_size": -1,
}
data = np.random.randn(DEFAULT_PARAMS["N"], DEFAULT_PARAMS["M"])
groups = [20] * int(N / 20) # assume each session has 20 rows
y = np.random.randint(2, size=N) # randomly choose 0 or 1

Large num_leaves can get very memory intensive. We should not need too many leaves, so generally using num_leaves <= 100 and increasing the number of estimators seems sensible to Gme.

• num_leaves: 10, Maximum memory usage: 2.28 GB - 0.80 GB = 1.48 GB
• num_leaves: 100, Maximum memory usage: 2.52 GB - 0.80 GB = 1.72 GB
• num_leaves: 1000, Maximum memory usage: 4.04 GB - 0.80 GB = 3.24 GB

Increasing n_estimators doesn't seem to raise memory much, but increases run time because each tree is fitted sequentially on the residual errors, so it cannot be parallelized.

• n_estimators: 10, Maximum memory usage: 2.28 GB - 0.80 GB = 1.48 GB
• n_estimators: 100, Maximum memory usage: 2.53 GB - 0.80 GB = 1.73 GB
• n_estimators: 1000, Maximum memory usage: 2.69 GB - 0.80 GB = 1.89 GB

Increasing N increases memory sublinearly. It seems that the data size itself will be more of a problem than the increase in LightGBM memory usage as N increases. For extremely large N, we can also set the subsample parameter to use only a fraction of the training instances for each step (i.e. stochastic rather than full gradient descent). By default subsample=1.0.

• N: 1,000, Maximum memory usage: 0.38 GB - 0.00 GB = 0.38 GB
• N: 10,000, Maximum memory usage: 0.45 GB - 0.04 GB = 0.41 GB
• N: 100,000, Maximum memory usage: 1.46 GB - 0.40 GB = 1.06 GB
• N: 1,000,000, Maximum memory usage: 6.12 GB - 4.00 GB = 2.12 GB
• N: 2,000,000, Maximum memory usage: 10.48 GB - 8.00 GB = 2.48 GB

In contrast to N, memory usage is quite sensitive to M, seems to increase linearly when M gets large. M=10,000 blows up my memory. I suppose this could be mitigated by setting colsample_bytree or colsample_bynode to sample a smaller subset.

• M: 100, Maximum memory usage: 2.08 GB - 0.16 GB = 1.92 GB
• M: 1000, Maximum memory usage: 4.92 GB - 1.60 GB = 3.32 GB
• M: 2000, Maximum memory usage: 9.69 GB - 3.20 GB = 6.49 GB
• M: 3000, Maximum memory usage: 14.35 GB - 4.80 GB = 9.55 GB

To deal with the high memory usage of large M, we can set colsample_bytree which samples a subset of columns before training each tree. This will help to mitigate the memory usage. For this experiment, we set M=2000 to simulate data with high number of dimensions.

• colsample_bytree: 0.1, Maximum memory usage: 8.60 GB - 3.20 GB = 5.40 GB
• colsample_bytree: 0.2, Maximum memory usage: 9.58 GB - 3.20 GB = 6.38 GB
• colsample_bytree: 0.4, Maximum memory usage: 10.06 GB - 3.20 GB = 6.86 GB
• colsample_bytree: 0.6, Maximum memory usage: 10.07 GB - 3.20 GB = 6.87 GB
• colsample_bytree: 0.8, Maximum memory usage: 10.46 GB - 3.20 GB = 7.26 GB

In contrast, setting colsample_bynode does not help memory usage at all. Not too sure why, but I suppose since multiple nodes for the same tree can be split at the same time, the full feature set still has to be kept in memory.

• colsample_bynode: 0.1, Maximum memory usage: 10.49 GB - 3.20 GB = 7.29 GB
• colsample_bynode: 0.2, Maximum memory usage: 10.49 GB - 3.20 GB = 7.29 GB
• colsample_bynode: 0.4, Maximum memory usage: 10.49 GB - 3.20 GB = 7.29 GB
• colsample_bynode: 0.6, Maximum memory usage: 10.49 GB - 3.20 GB = 7.29 GB
• colsample_bynode: 0.8, Maximum memory usage: 10.48 GB - 3.20 GB = 7.28 GB

Tweaking boosting and data_sample_strategy don't seem to affect memory usage too much. Using dart seems to require a bit more memory than the traditional gbdt.

• data_sample_strategy: bagging, boosting: gbdt, Maximum memory usage: 8.90 GB - 3.20 GB = 5.70 GB
• data_sample_strategy: goss, boosting: gbdt, Maximum memory usage: 9.58 GB - 3.20 GB = 6.38 GB
• data_sample_strategy: bagging, boosting: dart, Maximum memory usage: 9.81 GB - 3.20 GB = 6.61 GB
• data_sample_strategy: goss, boosting: dart, Maximum memory usage: 9.80 GB - 3.20 GB = 6.60 GB

Another bottleneck we can tackle is to realize that LightGBM is a two-stage algorithm. In the first stage, LightGBM uses the full dataset to construct bins for each numeric variable (controlled by the max_bins argument) based on the optimal splits. In the second stage, these discretized bins are then used to map and split the numeric variables during the actual training process to contruct trees. From my understanding, the first stage cannot be chunked as it requires the full dataset, but the second stage can be chunked (as per any stochastic gradient descent algorithm) where a fraction of the dataset is loaded at each time. Hence, the real bottleneck appears to be the first stage, when the bins are constructed.

According to this thread, we can separate the memory usage between the two stages by using lgb.Dataset. First, we initialize the Dataset object and make sure to set free_raw_data=True (this tells it to free the original data array after the binning is done). Then, we trigger the actual dataset construction using dataset.construct(). Thereafter, we are free to delete the original data array to free up memory for the actual training. The following code illustrates this concept.

dataset = lgb.Dataset(data=data, label=y, group=groups, free_raw_data=True)
del data
dataset.construct()
lgb.train(params=params, train_set=dataset)

TF-IDF

Term Frequency - Inverse Document Frequency is a well known method for representing a document as a bag of words. For a given corpus $\mathbb{C}$, we compute the IDF value for each word $w$ by taking $id{f}_w:=\frac{1}{logd{f}_w}$, with $d{f}_w$ denoting the number of documents in $\mathbb{C}$ containing the word $w$. The document $d$ is represented by a vector of length corresponding to the number of unique words in $\mathbb{C}$. Each element of the vector will be a tf-idf value for the word, i.e. $tfid{f}_w^d:=t{f}_w^d\cdot id{f}_w$, where $t{f}_w^d$ represents the term frequency of the word $w$ in document $d$. Sometimes, we may l1 or l2 normalize the tf-idf vector so that the dot product between document vectors represents the cosine similarity between them.

Bayesian Smoothing

We may want to apply some bayesian smoothing to the $t{f}_w^d$ terms to avoid spurious matches. For example, suppose that a rare word $w_r$ appears only in documents $d_1$ and $d_2$ in the entire corpus just by random chance. The $id{f}_{w_r}$ will be a large value, and hence documents $d_1$ and $d_2$ will have a high cosine similarity just because of this rare word.

For the specific setting I am considering, we can deal with this problem using bayesian smoothing. The setting is as follows:

• Each document $d$ represents a job, and each job is tagged to an occupation $o$
• An occupation can have one or more jobs tagged to it
• We wish to represent each occupation as a TF-IDF vector of words

To apply bayesian smoothing to this scenario, notice that we only need to smooth the term frequencies $t{f}_w^d$. Since the IDF values $id{f}_w$ are estimated across the whole corpus, we can assume that those are relatively reliable. And since term frequencies are counts, we can use a poisson random variable to represent them. See reference for a primer on the gamma-poisson bayesian inference model.

Specifically, we assume that ${\theta }_w^o$ is the poisson parameter that dictates the term frequency of $w$ in any document belonging to $o$, i.e. $t{f}_w^d\sim Poisson({\theta }_w^o)$. We treat the observed term frequency $t{f}_w^d$ for each document $d$ belonging to $o$ as a data sample to update our beliefs about ${\theta }_w^o$. We start with an uninformative gamma prior for ${\theta }_w^o$, and obtain the MAP estimate for ${\hat{\theta }}_w^o$ as below, with $d{f}_{d\in o}$ denoting the number of documents that belong to occupation $o$.

$${\hat{\theta }}_w^o=\frac{a+{\sum }_{d\in o}t{f}_w^d-1}{b+d{f}_{d\in o}}$$

We can thus use this formula to obtain posterior estimates for each ${\hat{\theta }}_w^o$. One possible choice of the prior parameters $a$ and $b$ is to set $a$ to be the mean term frequency for word $w$ per document in the entire corpus, and to set $b:=1$. This prior corresponds to ${\hat{\theta }}_w^o$ following a gamma distribution with mean $a$ and variance $a$, which seems to be a reasonable choice that can be overrided by a reasonable amount of data.

The posterior variance, which may also be helpful in quantifying the confidence of this estimate, is: $$Var({\hat{\theta }}_w^o)=\frac{a+{\sum }_{d\in o}t{f}_w^d}{(b+d{f}_{d\in o}{)}^2}$$

Finally, after obtaining the posterior estimates for each ${\hat{\theta }}_w^o$, we can just use them as our term frequencies and multiply them by the IDF values as per normal. We can also apply l1 or l2 normalization thereafter to the tf-idf vectors. This method should produce tf-idf vectors that are more robust to the original problem of spurious matches.

For illustration, for a very rare word $w_r$, $a$ will be a low value close to 0 (say 0.01). Suppose we were to observe $n$ number of new documents, each containing one occurrence of word $w_r$. Then the posterior estimate of ${\hat{\theta }}_w$ will update as follows:

As desired, the estimate for ${\hat{\theta }}_w$ starts off at a very small value and gradually approaches the true value $1$. This will help mitigate the effect of spurious matches. If we desire for the update to match the data more quickly, we can simply scale $a$ and $b$ down by some factor, e.g. now $a:=\frac{a}{5}=0.002$ and $b:=\frac{b}{5}=0.2$. Then we have:

As a final detail, note that the update formula can result in negative estimates if $a<1$ and ${\sum }_{d}t{f}_w^d=0$. The small negative value is probably not a big problem for our purposes, but we could also resolve it by setting the negative values to zero if desired.

Cross Encoders

Cross encoder is a type of model architecture used for re-ranking a relatively small set of candidates (typically 1,000 or less) with great precision. In the Question-Answering or machine reading literature, typically the task involves finding the top matching documents to a given query. A typical task is the MS MARCO dataset, which seeks to find the top documents that are relevant to a given bing query.

Basic Setup

Typically, the base model is some kind of pre-trained BERT model, and a classification head is added on top to output a probability. Each (query, document) pair is concatenated with [SEP] token in-between to form a sentence. The sentence is fed into the classification model to output a probability. The model is trained using binary cross-entropy loss against 0,1 labels (irrelevant or relevant).

This is the setup used by <Nogeuira 2019>, possibly the first paper to propose the cross encoder. Some specifics for their setup:

• The query is truncated to max 64 tokens, while the passage is truncated such that the concatenated sentence is max 512 tokens. They use the [CLS] embedding as input to a classifier head.
• The loss for a single query is formulated as below. $s_i,s_j$ refers to the score from the classifier model, $J_{pos}$ refers to the documents that are relevant, and $J_{neg}$ refers to documents in the top 1,000 retrieved by BM25 that are not relevant. Note that this results in a very imbalanced dataset.

$$L=-\sum _{j\in J_{pos}}log(s_j)-\sum _{j\in J_{neg}}log(1-s_j)$$

• The model is fine-tuned with a batch size of 128 sentence pairs for 100k batches.

As opposed to bi-encoders (or dual encoders), which take a dot product between the query embedding and the document embedding, we cannot pre-compute embeddings in the cross encoder setting, because the cross encoder requires a forward pass on the concatenated (query, document) pair. Due to the bi-directional attention on the full concatenated sentence, we need the full sentence before we can compute the score, which requires the query that we only see at inference time. Hence, the cross encoder is limited to reranking a small set of candidates as it requires a full forward pass on each query, candidate_document pair separately.

Contrastive Loss

The vanilla binary cross entropy loss proposed above may be thought of as a loss, in which each document is either relevant or irrelevant in absolute terms. However, treating relevance as a concept often better reflects reality. For example, given the first page of search results for a Google query, most of the documents should be relevant to some extent, but some are more relevant than the rest (and get clicked on). Simply treating all clicks as relevant and all non-clicks as irrelevant naively ignores the context (i.e. the neighbouring search results) in which the clicks were generated. It assumes that across query sessions, the average level of relevance of the results is comparable. Treating relevance as a concept within the same query session weakens this assumption and hence often works better.

Thus <Gao 2021> proposes the Local Contrastive Estimation loss. For a given query q, a positive document $d_q^{+}$ is selected, and a few negative documents $d_q^{-}$ are sampled using a retriever (e.g. BM25). The contrastive loss then seeks to maximize the softmax probability of the positive document against the negative documents.

$$L_{LCE}=\frac{1}{|Q|}\sum _{q\in Q,G_q}-log\frac{exp(s_{q,d_q^{+}})}{{\sum }_{d_q^{-}\in G_q}exp(s_{q,d_q^{-}})}$$

It is confirmed in multiple experiments in Gao 2021 and Pradeep 2022 that LCE consistently out-performs point-wise cross entropy loss. Furthermore, the performance consistently improves as the number of negative documents per query (i.e. $|G_q|$) increases. In Gao 2021, up to 7 negatives (i.e. batch size of 8) were used. Pradeep 2022 shows that increasing the batch size up to 32 continues to yield gains consistently (albeit diminishingly).

Other details

Pradeep 2022's experiments show that using a stronger retrieval model (a ColBERT-based model) during inference generates slight gains in final performance (as opposed to BM25). Although Gao 2021 argues that it is also important to use the same retrieval model during model training (so that the cross encoder sees the same distribution of negatives during training and inference), Pradeep 2022 argues that the alignment is not as important as the stronger retrieval performance during inference.

References

• Nogueira 2019 - Passage Re-ranking with BERT
• Gao 2021 - Rethink Training of BERT Rerankers in Multi-Stage Retrieval Pipeline
• Pradeep 2022 - A Bag of Tricks for Improving Cross Encoder Effectiveness

SentenceTransformers

Collaborative Filtering

This post is trying to grok my colleague Yuxuan's sharing on various losses for collaborative filtering. All credit goes to him.

Collaborative filtering is typically done with implicit feedback in the RecSys setting. In this setting, interactions are often very sparse. Most of the time, only positive signals are recorded, but a non-interaction could either mean (i) user dislikes the item or (ii) the user was not exposed to the item. Hence, we cannot use algorithms like SVD which assume no interactions as irrelevance.

A generic and fairly common architecture for the collaborative filtering model is to embed each user and item into separate fixed size vectors, and use the cosine similarity between the vectors to represent a score. This score is fed into a cross entropy loss against the labelled relevance of user to item to train the embeddings.

Setup

Let $f(u)$ and $f(i)$ denote the $k$ dimensional embedding vector for user $u$ and item $i$. Let the similarity function be $s(u,i)$ which is typically $f(u{)}^{T}f(i)$, and distance function $d(u,i)$ which is typically $||f(u)-f(i)|{|}_2^2$. Then some common loss functions may be denoted as below.

Pointwise Losses are typically low-performing. For a given (u, i) pair, pointwise losses assume the presence of a 0, 1 label for relevance, and tries to predict it. The typical pointwise loss is Binary Cross Entropy, which may be expressed as:

$${\mathcal{L}}_{BCE}=\sum _{(u,i)\in \mathcal{D}}\log \sigma (s(u,i))-\sum _{(u,j)\notin \mathcal{D}}\log (1-\sigma (s(u,j)))$$

Pairwise Losses assume the presence of training triplets (u, i, j) which correspond to user, positive item and negative item. A typical pairwise loss is Bayesian Personalized Ranking, as follows:

$$-\sum _{(u,i,j)\in \tau }\log \sigma \left[s(u,i)-s(u,j)\right]$$

AB Testing

References

• General:

  • Experiments at Airbnb
  • Evan Miller's AB Testing Tools
  • PostHog AB Testing Examples
  • How we experiment at Monzo
  • Ron Kohavi 2022 - AB Testing Intuition Busters
  • Francesco Casalegno - AB Testing A Complete Guide

• On Bayesian AB Testing:

  • Convoy - The Power of Bayesian AB Testing
  • Chris Stucchio's White Paper
  • VeepeeTech Intro to Bayesian AB Testing

Examples

These summaries are based on reading PostHog's article on AB testing and studying the ones of interest further.

AB Testing at Airbnb

AB Testing is crucial because the outside world often has a larger effect on metrics than product changes. Factors like seasonality, economy can cause metrics to fluctuate greatly, hence a controlled experiment is necessary to control for external factors and isolate the effects of the product change.

Airbnb has a complex booking flow of search -> contact -> accept -> book. While they track the AB impact of each stage, the main metric of interest is the search to book metric.

One pitfall is stopping the AB Test too early. Airbnb noticed that their AB tests tend to follow a pattern of hitting significance early on but returning to neutral when the test has run its full course. This is a phenomenon known as peeking, which is to repeatedly examine an ongoing AB test. It makes it much more likely to find a significant effect when there isn't, since we are doing a statistical test each time we peek. For Airbnb, they hypothesize that this is also a phenomenon caused by the long lead time it takes from search -> book, such that early converters have a disproportionately large influence at the beginning of the experiment.

The natural solution to this problem is to conduct power analysis and determine the desired sample size prior to the experiment. However, Airbnb runs multiple AB tests at the same time, hence they required an automatic way to track ongoing AB tests and report when significance has been reached. Their solution back then was to create a dynamic p-value graph. The idea is that on day 1 of the experiment, we would require a very low p-value to declare success. As time goes on and more samples are collected, we can gradually increase the p-value until it hits 5%. The shape of this graph is unique to their platform and they did extensive simulations to create it, so this solution is not very generalizable.

Another pitfall was assuming that the system is working. After running an AB test for shifting from more words to more pictures, the initial result was neutral. However, they investigated and found that most browsers had a significant positive effect except for Internet Explorer. It turned out that the change had some breaking effect on older IE browsers. After fixing that, the overall result became positive. Hence some investigation is warranted when the AB test results are unintuitive. However, one needs to be cautious of multiple-testing when investigating breakdowns, since we are conducting multiple statistical tests.

Airbnb has a strong AB testing culture - only 8% of AB tests are successful (see Ron Kohavi's LinkedIn Post).

AB Testing at Monzo

Monzo has a big AB testing culture - they ran 21 experiments in 6 months. Monzo has a bottom-up AB testing culture where anyone can write a proposal on Notion. Some of the best ideas come from the customer operations staff working on the frontlines. A proposal comprises the following sections:

• What problem are you trying to solve?
• Why should we solve it?
• How should we solve it (optional)?
• What is the ideal way to solve this problem (optional)?

Many proposals end up becoming AB experiments. Monzo prefers launching pellets rather than cannonballs. This means that each experiment comprises small changes, is quick to build, and helps the team learn quickly.

AB Testing at Convoy

Convoy argues that bayesian AB testing is more efficient than frequentist AB testing and allows them to push out product changes faster while still controlling risk.

The argument against frequentist AB testing is as follows. Under traditional AB testing, we define a null hypothesis using the control group (call it A), and declare a treatment (call it B) as successful if the treatment value has a significant p-value, i.e. it falls outside of the range of reasonable values under the null. Based on power analysis and an expected effect size, we predetermine the necessary sample size to achieve sufficient power, and once this sample size is reached, we have a binary success or failure result based on the p-value.

Convoy argues that this approach is safe but inefficient. This is because prior to the sample size being reached, we do not have a principled way of saying anything about the effectiveness of the treatment, even if it is performing better. Furthermore, frequentist AB testing gives us a binary result, but it does not quantify the size of the difference. Specifically, an insignificant test where E(A)=10%, E(B)=11% is quite different from E(A)=15%, E(B)=10%. For the former case, one can argue for launching B even if the p-value did not hit significance, whereas for the latter we should definitely not launch.

Bayesian analysis comes in to make the above intuition concrete. Suppose we are interested in the clickthrough rate (CTR) of variant A vs B. Bayesian analysis provides a distribution of the average CTR for each variant A, B at any point of the AB test, based on the results that it has seen thus far. These posterior distributions reflect both the mean of the data (how far apart $E(A)$ is from $E(B)$) and the variance of the data (how spread out the distributions are), allowing us to quantify how much we stand to gain if we were to pick either variant A or B at this point in time.

Concretely, they define a loss function as follows. Let $\alpha$ and $\beta$ be the unobserved true CTR for variants A and B respectively, and let the variable $x$ denote which variant we decide to choose. Then our loss for choosing each variant can be expressed as:

$$\begin{array}{r}\mathcal{L}(\alpha ,\beta ,x)=\{\begin{array}{cl}max(\beta -\alpha ,0)\text{if }& x=A\\ max(\alpha -\beta ,0)\text{if }& x=B\end{array}\end{array}$$

In other words, the loss above expresses how much we stand to lose by picking the unfortunately wrong variant based on incomplete evidence at this point in time. Of course, we do not know the true values of $\alpha$ and $\beta$, so we need to estimate the loss using our posterior distributions which we computed from data. We then compute the expected loss based on the posterior distributions $\hat{\alpha }\sim A$, $\hat{\beta }\sim B$ as such: $$\mathbf{E}(\hat{\mathcal{L}})={\int }_{\hat{\alpha }\sim A}{\int }_{\hat{\beta }\sim B}\mathcal{L}(\hat{\alpha },\hat{\beta },x)\cdot f(\hat{\alpha },\hat{\beta })d\alpha d\beta$$

Here, $f(\hat{\alpha },\hat{\beta })$ is the joint posterior distribution, which I believe we can obtain by multiplying the two independent posterior distributions $\hat{\alpha }$, $\hat{\beta }$ together. We can also perform random draws from the posterior distributions to estimate this statistic. Finally, we make a decision by choosing the variant that dips below a certain loss threshold, which is usually a very small value.

The appeal of the bayesian approach is two-fold:

1. It allows us to make faster decisions. Suppose an experiment is wildly successful, and it is clear within a day that variant B is better. Bayesian analysis will be able to reveal this result, whereas frequentist analysis will tell us to wait longer (since we estimated the effect size to be smaller).
2. It allows us to control risk. Since we are making decisions based on minimizing risk (supposing we had picked the poorer variant), we may be sure that even if we are wrong, it will not severely degrade our product. So supposing that there is no significant engineering cost between variant A and B, we can more rapidly roll out new iterations with the assurance that on average our product will be improving.

Power Analysis

Reference: Probing into Minimum Sample Size by Mintao Wei

How to determine the minimum sample size required to achieve a certain significance level and power desired?

The following table helps us understand how type I and type II errors come into play:

Type I Error refers to rejecting the null hypothesis when it is actually true, e.g. when we think that an AA test has significant difference. In short, it means we were too eager to deploy a poor variant. This should happen with probability $\alpha$, which is the significance level which we set (typically 5%). We have a better handle on type I error because the baseline conversion rate is typically known prior to an experiment.

Type II Error refers to failing to reject the null hypothesis when the alternate is actually true, i.e. we failed to get a significant effect on an improvement that is known to be better. In short, we were too conservative and failed to deploy a winning variant. In order to reason about type II error, we need to make a guess on what is the distribution of test variant B. Typically, this is done by assuming a minimum effect $\delta$ we wish to detect, and setting ${\mu }_B={\mu }_A+\delta$, and re-using the standard deviation from A. With these assumptions in place, we use $power=1-\beta$ to determine the type II error that should only occur with probability $\beta$ (typically 20%). Note that since $\delta$ is the minimum effect we wish to detect, if the actual effect turned out to be larger, the type II error can only be smaller than our desired amount, which is ok.

Now we can derive the formula for the minimum sample size required to achieve the desired levels of type I and type II error respectively.

Let us define the baseline conversion rate as $p$, and the minimum relative detectable effect rate as $d$. Consequently, the minimum detectable delta is $\delta =d\times p$. Let the desired power level be $1-\beta$, and the desired significance level as $\alpha$. Assume the scenario where we are running an AA or AB test with two variants of sample size $N$ each.

Firstly, we write down the distribution of the sample mean difference supposing we knew the true population means and standard deviations. Let $\mathbb{E}(X_A)={\mu }_A,Var(X_A)={\sigma }_A^2$ and $\mathbb{E}(X_B)={\mu }_B,Var(X_B)={\sigma }_B^2$. Note that $X_A,X_B$ may have arbitrary distributions, e.g. they could measure proportions, revenue etc.

Under the central limit theorem, the sample means will be distributed like so with $N_A,N_B$ samples: ${\bar{X}}_A\sim N\left({\mu }_A,\frac{{\sigma }_A^2}{N_A}\right)$, ${\bar{X}}_B\sim N\left({\mu }_B,\frac{{\sigma }_B^2}{N_B}\right)$. Importantly, the difference of the sample means will have the distribution below. Note that we add the variances together because $Var(B-A)=Var(B)+Var(A)$ for any two independent random variables $A,B$.

$$\begin{array}{r}{\bar{X}}_D={\bar{X}}_B-{\bar{X}}_A\sim N\left({\mu }_B-{\mu }_A,\frac{{\sigma }_A^2}{N_A}+\frac{{\sigma }_B^2}{N_B}\right)\end{array}$$

Now we can start working from the desired $\alpha ,\beta$ levels to the minimum sample size. We need to ensure that both objectives below are achieved with our sample size $N_A,N_B$:

1. Assuming null hypothesis to be true, ensure that type I error $\le \alpha$.
2. Assuming alternate hypothesis to be true, ensure that type II error $\le 1-\beta$.

Let us define some notation first.

• Let $z(\varphi )$ denote the critical value under the standard normal distribution such that $P(Z\le z(\varphi ))=\varphi$. This is basically the scipy.stats.norm.ppf function, e.g. $z(0.975)=1.96$.
• We also want to denote the critical value under the distribution ${\bar{X}}_D$ of the sample mean difference under the null or alternate hypothesis (these are non-standard normal distributions). Let these be $z_{{\bar{X}}_D|H_0}(\varphi )$ and $z_{{\bar{X}}_D|H_1}(\varphi )$ respectively.

Illustration for Power Analysis Derivation

For objective 1, assuming the null hypothesis and using equation (1) above, we have ${\bar{X}}_D|H_0\sim N(0,\frac{2{\sigma }_A^2}{N_A})$. Since $\alpha$ is a two-tailed probability and we want the critical region on the right-side, let ${\alpha }^{\prime }=1-\alpha /2$. E.g. $\alpha =0.05$ implies ${\alpha }^{\prime }=0.975$. Then:

$$\begin{array}{r}{z}_{{\bar{X}}_D|H_0}({\alpha }^{\prime })=z({\alpha }^{\prime })\cdot \sqrt{\frac{2{\sigma }_A^2}{N_A}}\end{array}$$

Note that equation (2) above tells us the critical value such that we will reject the null hypothesis if the sample mean of $B$ is greater than this value. To satisfy objective 2, we must thus ensure that the probability of rejecting the null hypothesis is at least $power=1-\beta$. In other words, we want $\delta -z_{{\bar{X}}_D|H_1}(1-\beta )\ge z_{{\bar{X}}_D|H_0}({\alpha }^{\prime })$. Assuming the alternate hypothesis and again using equation (1), we have ${\bar{X}}_D|H_1\sim N(\delta ,\frac{{\sigma }_A^2}{N_A}+\frac{{\sigma }_B^2}{N_B})$. So then:

$$\begin{array}{rl}\delta -z_{{\bar{X}}_D|H_1}(1-\beta )& \ge z_{{\bar{X}}_D|H_0}({\alpha }^{\prime })\\ \delta -z(1-\beta )\times \sqrt{\frac{{\sigma }_A^2}{N_A}+\frac{{\sigma }_B^2}{N_B}}& \ge z({\alpha }^{\prime })\times \sqrt{\frac{2{\sigma }_A^2}{N_A}}\end{array}$$

For the purpose of getting a minimum $N$, we assume $N=N_A=N_B$. Then using this and squaring both sides with some rearranging gives us:

$$\begin{array}{rl}N& \ge \frac{{\left(z(1-\beta )\sqrt{{\sigma }_A^2+{\sigma }_B^2}+z({\alpha }^{\prime }){\sigma }_A\sqrt{2}\right)}^2}{{\delta }^2}\end{array}$$

Which gives us the required minimum sample size equation. If we assume ${\sigma }_A={\sigma }_B$, as is often assumed because we do not know the variance of the treatment, then it simplifies to the following form (as seen in Ron Kohavi's paper).

$$\begin{array}{rl}N& \ge \frac{2{\sigma }_A^2\cdot {\left(z(1-\beta )+z({\alpha }^{\prime })\right)}^2}{{\delta }^2}\end{array}$$

Bernoulli Events

The equation (3) above for the minimum sample size requires us to know the standard deviation under the null and alternate hypotheses. Usually, the standard deviation under the null is computed from historical data, and it is assumed that ${\sigma }_A={\sigma }_B$. However, if the event we are interested in may be represented as a bernoulli random variable (e.g. an impression is shown and user either clicks or does not click with some probability), the equation may be simplified.

Specifically, the variance of a bernoulli random variable with probability $p$ is $p\cdot (1-p)$. Thus, if $X_A\sim Bernoulli(p_A)$, then $Var(X_A)=p_A\cdot (1-p_A)$, and likewise for $X_B$.

So we can use ${\sigma }_A=\sqrt{p_A\cdot (1-p_A)}$ and ${\sigma }_B=\sqrt{(p_A+\delta )\cdot (1-p_A-\delta )}$ and substitute these into equation (3). We will then be able to have a minimum sample size formula by just specifying $\alpha$, $\beta$, baseline conversion $p_A$ and minimum relative difference $d$. This is the formula used by Evan Miller's sample size calculator.

Large Language Models

LLMs are generally used in an auto-regressive way, where the user supplies a prompt and the LLM returns a generated response. This framework makes it amenable to a wide range of tasks.

HuggingFace has a great blog post explaining how we can run LLMs on humble hardware. Typically, LLMs have billions of parameters. The following rules of thumb helps us know how much memory we need to load the LLM into memory. For a model with X billion parameters:

• Loading in float32 requires 4X GB of VRAM
• Loading in bfloat16 requires 2X GB of VRAM
• Loading in int8 requires X GB of VRAM

Hence we see that we can load a ~7 billion parameters model with around 14GB of VRAM if loaded in bfloat16, which makes it feasible to run on GPUs like Tesla T4 with 16GB of VRAM. This can be done when loading the model with from_pretrained(..., torch_dtype=torch.bfloat16). Most models are trained in bfloat16 anyway, so it makes sense to load them at that precision.

Current popular open source LLM of that size includes mosaicml/mpt-7b, which can be easily downloaded and used using huggingface.

Quantization

It turns out that we can lower the precision of models even further than 16 bits if we use a quantization method (see e.g. Dettmers 2022, this paper is the basis for the package bitsandbytes used for quantization). The general idea is akin to encoding - we encode each number from a higher precision into a "codeword" in the lower precision (i.e. quantization). Numbers that are close to one another in the higher precision may get mapped to the same "codeword". When we want to use the encoded value, we look up the value in the higher precision that it maps to (i.e. de-quantization).

When applying this to quantizing a neural network, the steps involved are:

1. Quantize all model weights to target precision (e.g. int8)
2. Pass the input vector at bfloat16
3. At each layer, dequantize the weights and perform matmul in bfloat16
4. Quantize the weights again for storage

Hence while quantization lowers the memory footprint of the model, it may increase inference time. To use quantization, we need to do the following (also make sure bitsandbytes is pip installed). We can also pass load_in_4bit=True for 4bit quantization. More info on quantization usage is available at HuggingFace. An important note is that a GPU is required for quantization, at least in the bitsandbytes package.

model = AutoModelForCausalLM.from_pretrained(..., load_in_8bit=True)

Flash Attention

The self-attention mechanism (Dao 2022) is at the heart of the transformer performance but is also a major bottleneck in terms of memory and computational cost. One of the most successful optimizations for the attention mechanism is the Flash Attention paper.

Suppose we have an input sequence of embeddings $X=(x_1,...,x_N)$ where $x_i\in {\mathbb{R}}^k$, such that $X\in {\mathbb{R}}^{k\times N}$. The transformer stores parameters $W_q,W_k,W_v\in {\mathbb{R}}^{k\times d}$, such that $Q=X^T\cdot W_q,K=X^T\cdot W_k,V=X^T\cdot W_v$ such that $Q,K,V\in {\mathbb{R}}^{N\times d}$. The self-attention matrix $S=Q{K}^T\in {\mathbb{R}}^{N\times N}$ is then computed to represent the pairwise interaction between tokens at position $i$ ($i^{th}$ row) and position $j$ ($j^{th}$ column). The row-wise softmax is taken $P=softmax(S)$ to convert these into probabilities and finally the output is $O=P\cdot V$.

Typically, $N$ is much larger than the hidden dimensions $k,d$, as $N$ can be 2,048 or larger. Hence the $Q{K}^T\in {\mathbb{R}}^{N\times N}$ matrix is the bottleneck for memory and computation. The flash attention proposes to do this computation in a block-wise manner to reduce the memory usage. Furthermore, the algorithm also speeds up the computation compared to naive attention because the block-wise implementation minimizes the number of read-write operations between the faster SRAM and slower HBM.

More details can be found in the notebook at Dao 2022 - FlashAttention. We can utilize flash attention like so:

%pip install optimum

model.to_bettertransformer()

Note that this is only supported for models that have implemented flash attention, e.g. gpt-neox, bloom etc.

Flash Attention is now support natively within Pytorch as torch.nn.functional.scaled_dot_product_attention (see blog). The usage is like below. We need transformers>=4.36 and torch>=2.1.1 to use it.

from optimum.bettertransformer import BetterTransformer
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("gpt2-large", torch_dtype=torch.float16)
model = BetterTransformer.transform(model, keep_original_model=False)

Position Representation

Recent innovations in position encoding has led to accuracy improvements for long input sequences. The initial attention papers used absolute position embeddings. Given an input sequence of embeddings $X=(x_1,...,x_N)\in {\mathbb{R}}^{d\times N}$, absolute position embeddings $p_1,...,p_N\in {\mathbb{R}}^d$ are generated by the model. These position embeddings are added to the input sequence $\hat{X}=(x_i+p_1,...,x_N+p_N)$, thereby allowing to model to use these position cues.

It turns out that fixed positional embeddings are not ideal because they require the model to learn a fixed, unique representation of each position $1,...,N$. This does not represent language well, because the word in position i in one sentence does not necessarily serve the same purpose as a word in the same position in another sentence. Rather, it is the relative distance between words that we want to encode in the model. Furthermore, training absolute position embeddings makes it difficult for our model to generalize to texts with longer sequences than what it was trained with.

Recent papers advocate for relative positional embeddings, with the following differences:

• Relative position encoding rather than absolute position encoding
• The encoding of relative position is done most naturally within the $Q{K}^T$ self-attention matrix, since that is where the relative degree of interaction between tokens at different positions is encoded
• The encoding should be such that tokens further apart have a lower value in the self-attention matrix and tokens closer together have a higher value

Rotational Position Embeddings (RoPE) (Su 2021) proposes rotating the query and key vectors by an angle proportional to the relative distance between the two positions. Specifically ${\hat{q}}_i^T{\hat{k}}_j=q_i^T{R}_{\theta ,i-j}{k}^j$, where $R_{\theta ,i-j}$ is a rotational matrix that performs the rotation.

Attention with Linear Biases (ALiBi) (Press 2022) proposes an even simpler method. It simply subtracts $|i-j|/m$ from row $i$ and column $j$ of the self-attention matrix, where $m$ is a fixed scalar (specific to each attention head). Intuitively, it penalizes the attention proportional to the distance between the tokens. The study shows that this method outperforms RoPE as we extrapolate to longer sequences, and is conceptually simpler.

Key-Value Cache

Most LLMs work in an auto-regressive manner, i.e. we provide an input sequence, generate the next token with the LLM, then append this token to the input sequence for the next iteration. Most LLMs are also trained with the causal language modelling objective and mask the upper triangle of the self-attention matrix, so that each query token $q_i$ can only interact with key token $k_j$ and value token $v_j$ if $j\ge i$. This setup encourages us to cache results from previous time steps, since a lot of computation is repeated.

The following is based on how I imagine this to work, after reading Cameron R. Wolfe's LinkedIn post. During training, we compute the projections $Q,K,V\in {\mathbb{R}}^{N\times d}$, where $N$ is the maximum sequence length and $d$ is the hidden dimension. The final output $O=softmax(Q{K}^T)V\in {\mathbb{R}}^{N\times d}$ actually provides a $d$-dimension representation of the model's prediction at each of the $N$ positions.

For next token generation, we can add a projection head, say $W_p\in {\mathbb{R}}^{d\times p}$, where $p$ represents the size of the vocabulary, such that $A=O{W}_p\in {\mathbb{R}}^{N\times p}$ can represent the activations at each of the $N$ positions for the next token. Specifically, $A_{[0,:]}$ represents the predictions of position $1$ given input tokens $[0]$, $A_{[1,:]}$ represents the predictions of position $2$ given input tokens $[0,1]$, and so on. These activations will then be fed into some cross-entropy loss such that activations at the correct token for each position gets rewarded. This allows us to do efficient training, since we simultaneously provide losses for the prediction at each of the $N$ positions to the model for backpropagation.

However, when we are doing inference generation, we only need to predict for the final position of the input sequence (suppose it is position $c$), i.e. we are only interested in $A_{[c,:]}$ and $O_{[c,:]}$. Hence for starters, we only need $q_c:=Q_{[c,:]}$ instead of the entire $Q$ matrix, since only that row comes into play. However, we still need the entire $K$ and $V$ matrices, since we want $q_c$ to interact with all tokens in the input sequence. This is where the KV cache comes in - we cache the existing $K$ and $V$ matrices, so that we only need to project the final token of the input sequence $x_c^T{W}_k\in {\mathbb{R}}^{1\times d}$ and $x_c^T{W}_v\in {\mathbb{R}}^{1\times d}$ at each step and append it to the existing cached $K$ and $V$. We can then compute $O_c=softmax(q_c{K}^T)V\in {\mathbb{R}}^{1\times d}$.

As one can imagine, this saves a lot of computation, but also increases memory costs. Kwon 2023 - PagedAttention shows that serving a 13B model on NVIDIA A100 with 40GB of memory:

• $65\%$ of memory is model parameters
• $>30\%$ is the KV cache
• A small amount of memory is used ephemerally for activation

The usage of KV cache is like so:

model = AutoModelForCausalLM.from_pretrained(...)
model.generate(..., use_cache=True)

How to fine-tune an LLM

• trl RL example
  • Fine tune a 20B GPT model on text generation on IMDB dataset (loaded in 8 bit)
  • Since step 1 used PEFT, we need to merge the adapter weights with the base model
  • Finally, use RLHF to generate positive movie reviews. They used a BERT IMDB sentiment classifer to generate rewards
• DataCamp example
  • Using SFTTrainer from the trl library to do supervised fine-tuning
• PEFT - based on LORA - PEFT is built by hugginface to support LORA.

Fine-tuning

LLMs are typically trained with next-token prediction task on large amounts of text in an unsupervised manner. Some of the behaviour in these texts are not desirable to imitate. For example, while Github is full of code repositories with common programming mistakes, we do not want the LLM to replicate such behaviour. Hence, a process of alignment is necessary to encourage the model to produce desired responses.

There are typically two stages to this fine-tuning: Supervised Fine-Tuning and Reinforcement Learning from Human Feedback.

Supervised Fine-Tuning (SFT). In this stage, pairs of (prompt, desired response) are provided to the LLM. The desired responses are often called "demonstrations" by human annotators. Some form of cross-entropy loss is then used to update the LLM to encourage it to generate the desired response. This is a straightforward approach that is similar to the next-token prediction task (except we are predicting the desired response given the prompt). In the InstructGPT paper, the authors show that an SFT-aligned 1.3B model (using ~13k training data) generates human-preferred outputs compared to a 175B GPT model, showing the importance of SFT.

The problem with SFT is that it trains the model to provide a very specific response to a particular prompt, which makes it hard for the model to generalize. It also fails to express the natural ambiguity of responses for an open-ended prompt (e.g. write a poem about AI). Hence, unless we can collect quality responses for a very wide variety of prompts, SFT is limited in its ability to generalize.

Reinforcement Learning from Human Feedback (RLHF). This is where RLHF comes in: given triplets of (prompt, preferred response, not preferred response), we train the model to generate the preferred response in a more generalizable way.

More recently, a method called Direct Preference Optimization is used to solve this problem without needing to do reinforcement learning. More on that here.

Useful Models

There are too many models on Huggingface, so I try to keep track of useful ones here. Generally the open_llm_leaderboard is a good place to start.

Testing SOLAR-10.7B

Here are some results from testing out SOLAR-10.7B-Instruct-v1.0 on Google Colab, with a T4 GPU. First we install packages and load the model. Note that we are quantizing it to 4bits to save memory, and specifying the compute type of torch.float16 as that is the tensor type that this model was trained in. Using torch.float32 (which is the default) will incur unnecessary conversion cost and slow down inference.

%pip install transformers~=4.36 torch~=2.1.1 bitsandbytes accelerate optimum

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import torch

quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16
)

MODEL_NAME = "Upstage/SOLAR-10.7B-Instruct-v1.0"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    device_map="auto",
    quantization_config=quantization_config,
)

We define a generate function.

def generate(prompt: str, max_length: int = 2048):  
    conversation = [{'role': 'user', 'content': prompt}] 
    prompt = tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(**inputs, use_cache=True, max_length=max_length)
    output_text = tokenizer.decode(outputs[0])
    return output_text

And now we can test it out on a few prompts.

Prompt: Can you explain web3 to a 5 year old? Response: Web3 is like a special world on the internet that lets people share and use things in a new and more open way. Imagine a playground where everyone can add their own swings, slides, and sandboxes without needing permission from a big boss. This makes it easier for people to work together and create new fun things online. (Inference time: 8s)

Prompt: Can a firefighter transit into the finance industry? Response: While it may not be a typical career transition, it is possible for a firefighter to transition into the finance industry. To make a successful move, the firefighter would need to acquire the necessary skills and qualifications relevant to the finance sector. This could involve pursuing additional education or training in finance, accounting, or business administration. Networking, gaining relevant work experience, and showcasing transferable skills from the firefighting profession can also help in making the transition.\n\nSome transferable skills from firefighting to finance could include strong problem-solving abilities, teamwork, decision-making under pressure, and the ability to manage risk. However, it's essential to note that the finance industry has its unique requirements, and the firefighter would need to demonstrate a strong commitment to learning and adapting to the new environment. (Inference time: 19s)

Seems to be pretty decent performance at a somewhat reasonable inference time out of the box.

Speeding up GPU inference.

Gemma Models

References:

• Huggingface Reference
• Fine-tuning Gemma with PEFT
• gemma-7b Model Card

Google released open source small models under the Gemma label. It comes in two sizes: 2B and 7B parameters, each with base (pretrained) and instruction-tuned versions. All have a context length of 8k tokens. Gemma integration is supported with transformers==4.38. Note that user consent is needed to access Gemma models, so one needs to create a HuggingFace account and create an access token.

Sample Code

# Use model_name="google/gemma-7b-it"
def load(model_name: str):
    tokenizer = AutoTokenizer.from_pretrained(model_name, token=TOKEN)
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        quantization_config=BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_use_double_quant=True,  # nested quantization
            bnb_4bit_quant_type="nf4",  # normalized quantization
        ),
        token=TOKEN,
    )
    return model, tokenizer

Inference Speed

Tested out the inference speed of the model using the following prompt: Can you recommend good alternative occupations for a occupation_name? Please be concise.

An example is as follows, for an occupation counsellor:

Sure, here are some alternative occupations for a counsellor:

* **Social worker:** Similar to counselling, social workers help individuals and families facing various challenges.
* **Community organizer:** Organizes and leads community initiatives to address social issues.
* **Teacher:** Can be a good fit for those who enjoy working with children and adolescents.
* **Human resource manager:** HR managers handle employee relations and conflicts.
* **Public relations specialist:** Helps organizations communicate effectively with the public.

The inference speed is about 8-10 seconds with 4-bit quantization.

Gemma-2b

Testing out google/gemma-2b-it shows that it has poorer performance but faster inference speed at 2 billion parameters. Here is an example of the results for the same prompt.

* **Life coach**
* **Coach in education**
* **Career coach**
* **Executive coach**
* **Leadership coach**
* **Management consultant**
* **Business development consultant**
* **Financial advisor**
* **Investment banker**

The inference speed is as follows:

• Single: 3s
• Batch of 4: 14s
• Batch of 8: 16s
• Batch of 16: 20s
• Batch of 32: OOM on T4 GPU

Phi-2

Mobius Labs fine-tuned the phi-2 model from Microsoft which seems promising, and released it under mobiuslabsgmbh/aanaphi2-v0.1. The output seems better than gemma-2b-it.

1. Social worker
2. Mental health therapist
3. School counselor
4. Employee assistance program (EAP) specialist
5. Rehabilitation counselor
6. Family therapist
7. Substance abuse counselor
8. Career counselor
9. Trauma-focused therapist
10. Child and adolescent therapist

These occupations involve working with individuals, families, and communities to promote mental health and well-being, and may provide similar skills and experiences to those of a counsellor.

• bitsandbytes example notebook from HuggingFace

Flash attention, torch.compile, quantization.

Encoder vs Decoder

There are broadly two categories of LLMs: Encoder-Decoder architecture (typified by BERT) and Decoder only architecture (typified by GPT-2 series). There are some innate differences between the two that affect the type of application each is well suited for.

Encoder-Decoder Architecture

The Encoder-Decoder architecture was proposed in the original Attention is All You Need paper by Vaswani et al. As the name suggests, there is a separate encoder and decoder module. The design choice of two distinct modules is easier to understand when we consider that the original paper was designed for the machine translation task, such that the inputs may be in English and the outputs may be in French.

On the encoder-side, the encoder module sees the full input sequence and encodes it into an numeric encoding representation matrix of size (seq_len, embed_dim). There is no input masking required because we assume we see the full English text before we begin translation. This encoding representation is passed to the decoder module as context.

On the decoder-side, we cannot allow the decoder module to see the full output sequence during training, since that would cause data leakage and the model to collapse. Note that the decoder simultaneously (in one forward pass) predicts a token for each position in the output sequence during training. Hence to avoid data leakage, causal masking is applied such that for each position being predicted, the decoder module is only allowed to "see" tokens in previous positions.

Note that this can be implemented simply by masking the attention matrix in each self-attention block in a triangular manner. The feedforward layers only map information from position i to position i, so there is no need to adjust that.

BERT Encoder Decoder Architecture (From Attention is All You Need)

The decoder module itself creates a representation of the output sequence (with causal masking) of size (seq_len, embed_dim). Now for the tricky part of merging the encoder and decoder representations. Naively, we can simply concatenate the two to form a representation of (seq_len x 2, embed_dim), and add further self-attention blocks (with appropriate causal masking). However, this would increase the number of parameters of the model.

The paper instead chose to take the Q matrix from the decoder and the K, V matrices from the encoder. Since the attention mechanism is $softmax(Q\cdot K^T)V$, this allows every token on the decoder representation to attend to every token on the encoder representation. The result is a weighted combination of the V vectors taken from the encoder representation which is eventually used to generate the tokens for each decoder position.

Here we can already observe some inductive biases baked into the encoder-decoder architecture. These comments are inspired by / taken from Hyung Won Chung's lecture:

1. There is a clear distinction between the tokens on the encoder-side and the decoder-side, such that tokens on either side cannot directly attend to each other until the final block when the mixing happens. This may be suitable for the translation task but not for generic chat completion, where it is quite arbitary which point in the conversation we determine as the encoding and which point we determine as the decoding.
2. The V matrix is taken entirely from the encoder-side, which may again limit the expressiveness of the model (what if there are important tokens on the decoder-side that would be useful to taken the V values from?). One may argue that in translation, the tokens on the encoder-side capture the full meaning of the sentence to be translated, so it is comprehensive. But the same is not true of generic chat completions.
3. The separation of an encoder and decoder module suggests that there is significant difference between the tokens on the encoder-side and decoder-side. Again, this inductive bias

Miscellaneous Notes

A collection of miscellaneous, useful notes.

f-strings

To surround a text with a symbol (say =) to a fixed length:

>>> text = "Title"
>>> print(f"{text:=^20}")

=======Title========

Vim

Command to interactively change each 'foo' to 'bar'. :%s triggers the substitute global command, followed by the search and replace phrases respectively. Finally g means replace all occurrences and c means with confirmation. Note that :s will only do the same for one line.

:%s/foo/bar/gc

Find Files

To find files anywhere on the system with the filename python using bash, use:

find . -name python

We can add * before and/or after the filename to allow other characters before or after our keyword:

find . -name *python*

To search not just in the filename but also in the full path (e.g. we only want to search in Desktop), we can do:

find . -wholename "*Desktop*python*"

Note that if we want to locate executable binaries, another useful command is whereis:

whereis cat
---
cat: /usr/bin/cat /usr/share/man/man1/cat.1.gz

Numpy Indexing

Suppose we have a 2D array X and would like to take a slice of certain rows and columns. We might try, for e.g., to take the first two rows of X and the 3rd/4th column of X, i.e. we expect to get a 2 by 2 matrix.

import numpy as np
X = np.random.randn(5, 5)
X[[0, 1], [2, 3]]

Unfortunately, this will return an array of items array(X[0, 2], X[1, 3]), which is not what we want. Instead, a slightly inefficient but clear way is to slice each axis separately. It is not optimal because the first slice creates a temporary array view before the second slice is applied.

X[[0, 1], :][:, [2, 3]]

Finally, the recommended way seems to be to use np.ix_. Doing np.ix_([0, 1], [2, 3]) creates a tuple of two elements.

idx = np.ix_([0, 1], [2, 3])
idx
>> (array([[0],
           [1]]),
>>  array([[2, 3]]))

Indexing the original array with this output X[idx] will then give us what we want.

asyncio

reference

asyncio is a single-threaded framework that does not use multi-threading or multi-processing to speed up tasks. Instead, a coordinator (or event loop) passes control from a time-consuming blocking function (e.g. time.sleep or an I/O operation) to other functions to run. This passing of control occurs with the await keyword. When the blocking function is completed, it notifies the coordinator and control returns to where the blocking function left off.

asyncio does not speed up CPU-bound tasks, due to its single-threaded design. It only works when the function being awaited is an I/O operation that is supported by asyncio. This includes stuff like:

• HTTP (supported through aiohttp)
• DB calls (e.g. aioredis)

asyncio is preferred for such tasks (e.g. making multiple I/O calls and doing something when they all return) over multi-processing because of its single-thread design, making it easier to debug.

relatedness

In the context of information retrieval, Trey Grainger in AI-Powered Search suggests a relatedness measure to connect arbitrary entities together. Suppose we have a collection of jobs and each job is tagged with a set of skills. Suppose we wish to retrieve relevant skills to an arbitrary free text query $q$.

The relatedness idea is to define a foreground of documents, e.g. based on a retrieval of documents using query $q$ which are related to the query, and to compare the attributes of the foreground against the background, i.e. all documents.

Mathematically, we can think of the foreground documents as a sample, and the background documents as the population. The strength of the relationship between each skill $t$ to the query $q$ may then be defined as the z-statistic of the one-sample z-test of proportions of the occurrence of skill $t$ in the foreground sample compared against the background population. A significantly greater occurrence in the sample compared to the population suggests a strong relationship between $t$ and $q$, and vice versa. Specifically:

$$z=\frac{\hat{p}-p}{\sqrt{\frac{p(1-p)}{n}}}$$

Where:

• $\hat{p}=\frac{df(q,t)}{df(q)}$ is the sample proportion.
• $df(q,t)$ is the number of documents in the foreground corresponding to query $q$ and contains skill $t$.
• $df(q)$ is the total number of documents in the foreground corresponding to query $q$. It is also the number of samples $n$.
• $p=P(t)$ is the probability of skill t appearing across all documents.

By performing a retrieval and ranking skills based on the z-statistic, we can compute a relationship between any arbitrary query and attribute of the documents (on the fly, if necessary). This functionality is implemented in solr.

Package Versioning

link

SemVer is the versioning standard for Python packages. For a package version of e.g. 1.2.3:

• The first number 1 is the major version. We update it when we make major API changes. A major of 0 indicates that the package is under initial development and anything may change.
• The second number 2 is the minor version. We update it when we add functionality in a backward compatible manner.
• The third number 3 is the patch version. We update it when we patch bugs in a backward compatible manner.

Poetry is a useful tool to manage package dependencies in a Python library. In the pyproject.toml file, we specify package dependencies in the following manner:

[tool.poetry.dependencies]
python = ">=3.8,<3.11"
requests = "^2.26.0"
pytz = "~2022.1"

The caret requirement (e.g. requests = "^2.26.0") means that SemVer compatible changes are allowed, i.e. an update is allowed if the version number does not modify the left-most non-zero digit. E.g.:

• ^1.2.3 means that 1.3.0 is allowed but not 2.0.0
• ^0.2.4 means that 0.2.5 is allowed but not 0.3.0

The tilde requirement is somewhat stricter. It specifies a minimal version but allows room to upgrade depending on how many digits are supplied. If major.minor.patch or major.minor is supplied, only patch-level changes are allowed. If major is supplied, minor and patch level changes are allowed.

• ~1.2.3 means that 1.2.7 is allowed but not 1.3.0
• ~1.2 means that 1.2.7 is allowed but not 1.3.0
• ~1 means that 1.3.0 is allowed but not 2.0.0

The poetry versioning specification depends on developer updating the library versions in a disciplined way. It allows us to provide some flexibility in package versioning while avoiding updating a dependent package to a version that breaks our code. Having more flexibility in specifying package dependencies reduces the risk of dependencies version conflicting.

Parallelism in Python

References: realpython, superfastpython

The Python Global Interpreter Lock (GIL) is a lock that allows only one thread to run at a time. The reason for this design feature in Python is because the Python interpreter is not thread-safe, meaning that if the GIL did not exist, the Python interpreter could introduce race conditions and other fatal errors due to multiple threads modifying memory at the same time. The GIL also allowed other non-thread-safe C programs to be easily integrated into the Python ecosystem, which contributed to Python's success as a language.

The GIL also improves performance for single-threaded programs, as it only requires a single lock to be managed (not sure how this works). There have been attempts to remove the GIL from Python but none have been successful because they would degrade single-threaded performance.

For many C extensions (e.g. numpy), multi-threading is still possible because these extensions are allowed to manually release the GIL (as long as they restore things back to normal when the functions return). This allows us to still use multi-threading for CPU-intensive functions with the GIL around. Similarly for Rust, we can release the GIL to achieve parallelism.

Alternatively, we can use multiprocessing to create multiple processes (instead of threads). Each process contains its own Python interpreter (and GIL) and hence can run truly in parallel. The downside is that the overhead of creating and managing processes is much more than that for threads, meaning that the benefits of multiprocessing are much dampened compared to multi-threading.

Memory Profiling

It is often useful to profile the memory usage of our script. In python, we can use memory_profiler to check the memory usage of our program line by line.

from memory_profiler import profile
import sys


@profile(stream=sys.stdout)
def f():
    a = [1] * (10**6)
    b = [2] * (2 * 10**7)
    del b


if __name__ == "__main__":
    f()

This will print the following useful information to stdout. Note that even before we did anything, there is background memory usage of 17MiB.

Filename: memory_profiling.py

Line #    Mem usage    Increment  Occurrences   Line Contents
=============================================================
     5     17.1 MiB     17.1 MiB           1   @profile(stream=sys.stdout)
     6                                         def f():
     7     24.5 MiB      7.5 MiB           1       a = [1] * (10**6)
     8    177.2 MiB    152.6 MiB           1       b = [2] * (2 * 10**7)
     9     24.8 MiB   -152.4 MiB           1       del b

We might also want to track memory usage of a function over time. We can use memory_usage instead for that.

import time
from memory_profiler import memory_usage
def g(a, n: int = 100):
    time.sleep(1)
    b = [a] * n
    time.sleep(1)
    del b
    time.sleep(1)

if __name__ == "__main__":
    usage = memory_usage((g, (1,), {"n": int(1e7)}), interval=0.5)
    print(usage)

This will give us an array like so, showing the spike in memory in the middle of g.

[17.375, 17.5234375, 17.5234375, 19.34765625, 93.59765625, 93.59765625, 17.53125, 17.53125, 17.53125]

CUDA

PyTorch may sometimes throw errors if the installed torch version does not match the installed CUDA version. To address this, we need to first check the CUDA version using the nvcc command:

/usr/local/cuda/bin/nvcc --version
---
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2022 NVIDIA Corporation
Built on Wed_Jun__8_16:49:14_PDT_2022
Cuda compilation tools, release 11.7, V11.7.99
Build cuda_11.7.r11.7/compiler.31442593_0

Then install the correct version of torch:

pip install -U torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

List Comprehension

When doing nested list comprehensions in python, the later loop becomes nested as the inner loop. So this:

l = [i * j for i in range(4) for j in range(5)]

Is equivalent to:

l = []
for i in range(4):
    for j in range(5):
        l.append(i*j)

This also explains why flattening a list of lists is of the following form:

list_of_lists = [[0, 1, 2], [3, 4, 5], [6, 7, 8]]
flattened = [item for sublist in list_of_lists for item in sublist]

When we translate it into a nested for loop, we can see that the inner loop should indeed be placed at the end:

flattened = []
for sublist in list_of_lists:
    for item in sublist:
        flattened.append(item)

Bradley-Terry Model

Based on wikipedia. The Bradley-Terry model (1952) is a probability model that allows us to infer scores for individual objects based on a dataset of pairwise comparisons between them. Specifically, it estimates the probability that $i▹j$ (i.e. $i$ is preferred to $j$) as:

$$P(i▹j)=\frac{p_i}{p_i+p_j}$$

Where $p_i$ is a positive real-valued score assigned to object $i$ (not necessarily a probability). Typically, $p_i$ is parametrized as an exponential score $p_i=e^{{\beta }_i}$, and the goal is to learn the parameters ${\beta }_i$ from pairwise comparisons. This results in:

$$P(i▹j)=\frac{e^{{\beta }_i}}{e^{{\beta }_i}+e^{{\beta }_j}}$$

Parameter Estimation

Parameter estimation is typically done using maximum likelihood. Starting with a set of pairwise comparisons between individual objects, let $w_{ij}$ be the number of times object $i$ beats object $j$. Then the likelihood of a given set of parameters $\mathbf{p}:=[p_1,...,p_n]$ ($n$ denotes number of objects) is as follows:

$$\begin{array}{rl}L(\mathbf{p})& =ln\prod _{ij}P(i▹j{)}^{w_{ij}}\\ & =\sum _{i=1}^n\sum _{j=1}^{n}ln{\left(\frac{p_i}{p_i+p_j}\right)}^{w_{ij}}\\ & =\sum _{i=1}^n\sum _{j=1}^n{w}_{ij}\left[ln{p}_i-ln(p_i+p_j)\right]\end{array}$$

This likelihood function can then be minimized by differentiating wrt $p_i$ and solved by setting to zero.

BT Model as a Sigmoid Function

We can also express the likelihood as a function of the difference in scores ${\beta }_i-{\beta }_j$. Recall that the sigmoid function is $\sigma (x)=1/(1+e^{-x})$. Then: $$\begin{array}{rl}L(\mathbf{p})& =\sum _{i=1}^n\sum _{j=1}^n{w}_{ij}\cdot ln\left[\frac{e^{{\beta }_i}}{e^{{\beta }_i}+e^{{\beta }_j}}\right]\\ & =\sum _{i=1}^n\sum _{j=1}^n{w}_{ij}\cdot ln\sigma ({\beta }_i-{\beta }_j)\end{array}$$

The derivation for the second line above is found at /Identities/sigmoid. This re-parametrization shows that the BT-model is really modelling the preference as a difference in scores and then running that through the sigmoid function to convert it into a probability. This means that we are basically running a pairwise logistic regression.

Setting up WSL

We may need to uninstall and re-install WSL (Windows Subsystem for Linux) from time to time. Here is the step-by-step.

1. Tear down and delete all files. wsl --unregister Ubuntu-22.04
2. Re-install wsl. wsl --install Ubuntu-22.04
3. Set up keychain to auto-find ssh-agent (from Windows) and add keys
   • Copy .ssh folder from Windows to wsl ~/.ssh
   • Add to ~/.bashrc the following: eval $(keychain --eval --agents ssh id_rsa)
4. Install Docker
   • Add users to docker to allow vscode access
   • https://docs.docker.com/engine/install/linux-postinstall/
5. Essentials (get C linker)
   • sudo apt install build-essential
6. Install rust
   • https://www.rust-lang.org/tools/install
7. Install mdbook
   • cargo install mdbook mdbook-katex

To Read

• https://netflixtechblog.com/improve-your-next-experiment-by-learning-better-proxy-metrics-from-past-experiments-64c786c2a3ac
• https://netflixtechblog.com/sequential-a-b-testing-keeps-the-world-streaming-netflix-part-1-continuous-data-cba6c7ed49df
• https://arxiv.org/abs/2407.02464
• https://arxiv.org/abs/2305.12102
• https://arxiv.org/abs/1709.03933
• https://arxiv.org/abs/2101.08769
• https://arxiv.org/abs/1205.2618 # BPR
• https://arxiv.org/abs/2312.08520 # InfoNCE+
• https://arxiv.org/abs/2308.06091 # MAWU
• https://arxiv.org/abs/2109.12613 # CCL
• https://arxiv.org/abs/2201.02327 # SSM
• https://arxiv.org/abs/2206.12811 # DirectAU
• KG-BERT

Packages

KeyBERT, KeyLLM

References:

• KeyBERT Article
• KeyLLM Article

KeyBERT and KeyLLM are packages to perform unsupervised keyword extraction from text. KeyBERT relies on BERT-based models, and the main idea is to extract n-gram phrases which have high semantic similarity to the overall document embedding. Some additional features are:

• Allow user to specify phrase length for extraction
• Add diversification via MMR to get diverse phrases

KeyLLM taps on LLMs to enhance the keyword extraction. Basically, it creates a prompt to ask an LLM to extract keywords from a document. It integrates with KeyBERT such that we can use KeyBERT to cluster documents, and only run KeyLLM on one document per cluster to save costs. It can also use KeyBERT to suggest candidates and use the LLM to verify.

Skills

This documents LinkedIn's approach to constructing a knowledge graph around skills.

November 2022 - Building LinkedIn's Skills Graph to Power a Skills-First World

An important model for LinkedIn is the skills graph - it maps 39k skills over 26 languages and over 347k aliases. The components of their structured skills graph:

• 39k nodes, each of which is a skill
• Each skill has multiple aliases
• Skills are connected via edges which specify a hierarchical relationship. (e.g. Marketing is a parent of Demand Generation which is a parent of Email Marketing)

Skills Usage. Skills are automatically extracted from job postings and member histories. Job postings are allowed to tag up to 10 skills. Users can see something like 4/10 skills match your profile to help them determine job fit.

The skills graph is constructed both using ML and manual review by taxonomists.

March 2023 - Building and maintaining the skills taxonomy that powers LinkedIn's Skills Graph

This article is more about how the Skills Graph is constructed.

Each skill node includes the following foundational details (e.g. Machine Learning):

• Description of the skill: the study of computer algorithms...
• Aliases: ML, ...
• Skills type: Hard or Soft
• Skill ID

Each skill is represented by a node in the graph and nodes are linked via edges called knowledge lineages. Skills can relate for various reasons. Edges are directed and represent a parent-child relationship (e.g. Software Development -> Back-end Software Development -> Python)

Each node can have multiple parents and/or children.

• Multiple parents example: Both Back-end Software Development and Mobile Software Development are parents of Java
• Multiple children example: Supply Chain Management has children Supply Chain Engineering, Logistics Management and Digital Supply Chain

The hierarchical relationship allows us to enrich skills understanding, since if a person knows a particular skill, we may infer that he knows something about all the parent nodes and perhaps some of the "sibling" nodes.

LinkedIn has certain quality guardrails on the structured skills, one of which is discouraging ambiguity. For example, an ambiguous skill like Networking may be mapped to skills like Computer Networking or Professional Networking. This type of ambigious relationship is not allowed, and LinkedIn either removes such edges or disallows such a node. The meaning of a phrase is determined by analyzing how the skill is used predominantly in LinkedIn. In cases where a phrase can have divergent meanings, the skill is disambiguated by expanding the phrase. For example, Cadence is disambiguated to Cadence Software and Boundary to Boundary Line.

Architecture. The components to their ecosystem are as follows:

1. Human Curated KG. Presumably, this is a purely human-curated KG where all nodes and edges are taken as true, and constantly curated by human taxonomists.
2. AI-generated KG. This is generated using ML models which use the human curated KG as training and validation data. The model behind this is KGBERT, which will be briefly covered below.
3. Serving. Both the AI-generated KG and human-generated KG are made available to all LinkedIn services via a REST API for online serving and also on HDFS to power offline inference.

KGBERT is a model for automatic edge prediction, which can generate the AI-generated KG from the human curated one. The basic idea is that the human curated KG is used to generate training and validation data in the following form:

[CLS] Tensorflow [SEP] Machine Learning [SEP] -> label: child - parent

Two skills are concatenated to form the context. The skill is represented at random by its title, description or title + description. The context is fed into a BERT model and a linear + softmax layer is attached to the [CLS] token to generate a softmax probability over 3 options:

• Parent -> Child
• Child -> Parent
• No relation

The positive labels are taken from edges in the human curated graph. The negative labels are generated by some heuristics:

• Skill pairs from different industries
• Niece / Nephew pairs e.g. Tensorflow vs Cognitive Computing
• Sibling pairs with the same parent e.g. Tensorflow vs Pytorch
• Loosely related pairs that are 3 or more steps apart

The nice thing about this architecture is the clean separation of the human curated graph from the AI-generated one. The human-in-the-loop setup allows humans to review the AI-generated graph and add new edges to the human curated graph, which in turn improves the AI-generated graph. This seems superior over mixing both AI-generated edges and human-curated edges in the same graph.

December 2023 - Extracting skills from content to fuel the LinkedIn Skills Graph

This post takes a deeper look into how skills are extracted from data by LinkedIn from various contexts, such as job listings or member profiles.

Note that skills can be mentioned either directly or indirectly (e.g. you are expected to know how to apply different techniques to extract information from data and communicate insights through meaningful visualizations). Hence, a simple span extraction approach will not be exhaustive in extracting skills. On the other hand, not constraining the problem to span extraction could lead to false positive errors, so our model has to be very accurate.

An overview of the skills extraction architecture

At a high level, the steps are:

1. The Skill Segmenter organizes the raw text into structured input, e.g. a resume is split into the qualification, career history etc. sections
2. The next step aims to generate a candidate list of skills from the context:
   • 2a. The Text Skill Tagger is a trie-based model that simply finds matches in the text that exactly match a node in the Skills Taxonomy. This puts the burden on the Skills Taxonomy to have all the aliases that cover every possible utterance of the skill. However, this model is very fast
   • 2b. The Semantic Tagger aims to overcome the coverage issues above and uses BERT semantic match to surface more candidates
   • 2c. The skills from 2a. and 2b. are expanded into more skills using the Skills Graph. It seems like they add immediate parents, children and siblings of each skill.
3. Each skill candidate is scored against the context to generate a confidence score. This section has a Shared Model and Domain-specific Model.
   • The Shared Model contains a context-encoder, which encodes the text surrounding the skill into an embedding, and a entity-encoder, which encodes surrounding skills, the job title, etc. All these embeddings are fed as context into the next stage. The Shared Model assumes that the relationship between the context surrounding each skill and the skill itself is constant across the different domains and thus benefit from a shared module. Their AB tests confirm this hypothesis and show that such multi-task training does increase lift in online metrics.
   • Each domain (job posting, member profile, course) has its own scoring model. The embeddings from the Shared Model and each skill candidate are fed into the Domain-specific Model, which is presumably a form of cross encoder that generate a confidence score for each skill candidate. Finally, some threshold is applied and a final list of skills is generated.

LinkedIn needs to generate skill extractions within 100ms for a high volume of edits. Hence, they used knowledge distillation to distil the BERT model down 80% in parameters.

As we can see, the LinkedIn architecture is fairly complex and context specific. A natural question for smaller players is how we can tap on LLMs today to simplify parts of this architecture to achieve comparable performance.

Identities

Sigmoid

Sigmoid Relationship to Bradley-Terry Model

Suppose we have scores $s_i,s_j$ for objects $i$ and $j$. Under the Bradley-Terry model, we can express the preference probability as follows:

$$\begin{array}{rl}P(i▹j)& =\frac{e^{s_i}}{e^{s_i}+e^{s_j}}\\ & =\frac{1}{(e^{s_i}+e^{s_j})/e^{s_i}}\\ & =\frac{1}{1+e^{s_j}/e^{s_i}}\\ & =\frac{1}{1+e^{s_j-s_i}}\\ & =\frac{1}{1+e^{-(s_i-s_j)}}\\ & =\sigma (s_i-s_j)\end{array}$$

Where $\sigma (x)=1/(1+e^{-x})$ is the sigmoid function.

Statistics

Adding / Subtracting RVs

Let $X,Y$ be two independent random variables. Then $Var(X\pm Y)=Var(X)+Var(Y)$.

First, show that $Var(X)=\mathbf{E}(x^2)-{\mu }_X^2$ where $x\sim X$. Proof: $$\begin{array}{rl}Var(X)& =\mathbf{E}\left[\right]\end{array}$$

Reference: https://apcentral.collegeboard.org/courses/ap-statistics/classroom-resources/why-variances-add-and-why-it-matters

Summaries of individual papers that I have taken a deeper look into.

Weinberger 2009 - Hashing for Multitask Learning

Weinberger 2009 - Feature Hashing for Large Scale Multitask Learning

This paper proposes a method to represent a large feature set in a smaller space by using hashing. It shows analytically that with a sufficiently large hash dimension $m$:

• The inner product between instances is preserved, i.e. doing a dot product between instances in the hashed dimension approximates the true dot product in the original dimension
• The same applies to learning a weight vector to generate predictions in the hashed space: the error of approximation goes to zero as $M$ increases

Setup

Consider having data points $x^{(1)},...,x^{(n)}\in {\mathbb{R}}^d$, where $d$ can be very large (e.g. millions). This setting is easily realized when we use, for example, word bi-grams and tri-grams as term-frequency features to perform some kind of text classification task. Such a large feature vector is unwieldy, and also inefficient since the feature vector is very sparse for a given text.

The hashing trick maps the high dimensional input vector to a smaller dimension feature space with the notation $\varphi :\mathcal{X}\to {\mathbb{R}}^m$, such that $m<<d$.

We start with the following definitions:

• Let $h$ be a hash function $h:\mathbb{N}\to \{1,...,m\}$
• Let $\mathcal{E}$ be a hash function $\mathcal{E}:\mathbb{N}\to \{\pm 1\}$

Note that while the definitions map from an input integer, we may apply them to texts as well, since any finite-length text may be assigned to a unique integer. This is typically done in practice by applying some hash algorithm to a given string, and then using the modulo function to restrict it to the desired range.

With this, and given two vectors $x,x^{\prime }\in {\mathbb{R}}^d$, we define the hash feature map:

$${\varphi }_j^{(h,\mathcal{E})}(x)=\sum _{i\in \mathbb{Z}:h(i)=j,1\le i\le d}\mathcal{E}(i)x_i$$

Where $j\in 1,...,m$ is an index in the hashed dimension space, and $i\in 1,...,d$ is an index in the input dimension space. We get a hash collision if more than one $i$ term is hashed into a given position $j$. For brevity, we may just write ${\varphi }_j^{(h,\mathcal{E})}:=\varphi$.

Analysis

With this setup, the paper aims to prove analytically that hashing in this manner preserves the characteristics of the original space. In other words, we can significantly reduce the dimension of our features but achieve the same predictive effect as the original space by doing the hashing trick. This also means that the detrimental effect of hash collisions is minimal with a sufficiently large $m$.

We won't trace through all the results, just the important and simple ones.

Lemma 2 The hash kernel is unbiased, i.e. ${\mathbb{E}}_{\varphi }\left[⟨x,x^{\prime }{⟩}_{\varphi }\right]=⟨x,x^{\prime }⟩$.

Proof. The proof simply starts by expanding the inner product in the hashed space as follows: $$⟨x,x^{\prime }{⟩}_{\varphi }=\sum _{i=1,...,d}\sum _{j=1,...,d}\mathcal{E}(i)\mathcal{E}(j)\cdot x_i{x}_j^{\prime }\cdot {\delta }_{h(i),h(j)}$$

Where ${\delta }_{h(i)=h(j)}$ is an indicator variable which takes $1$ if $h(i)=h(j)$ (i.e. they are hashed to the same position) and $0$ otherwise.

To see that this expansion is true, consider a position in the hashed space, e.g. $k$. The value at position $k$ looks something like the following. We just need to move the summands to the left and use the $\delta$ variable to denote the common hash positions where $x_i$ and $x_j^{\prime }$ interact (if i and j are hashed to different positions, they clearly do not interact in an inner product). $${\left[⟨x,x^{\prime }{⟩}_{\varphi }\right]}_k=\left[\sum _{i\in \mathbb{Z}:h(i)=k,1\le i\le d}\mathcal{E}(i)x_i\right]\left[\sum _{j\in \mathbb{Z}:h(j)=k,1\le j\le d}\mathcal{E}(j)x_j\right]$$

Now note that we can decompose the expectation over $\varphi$ into its independent constituents, i.e. $h$ and $\mathcal{E}$ respectively (since the two hashes are independent):

$${\mathbb{E}}_{\varphi }\left[⟨x,x^{\prime }{⟩}_{\varphi }\right]={\mathbb{E}}_h\left[{\mathbb{E}}_{\mathcal{E}}⟨x,x^{\prime }{⟩}_{\varphi }\right]$$

Now we just need to observe that the hashed values $\mathcal{E}(i),\mathcal{E}(j)$ are independent from all other terms in general, but also independent from each other whenever $i\ne j$ (provided our hash function is pairwise independent). Thus when $i\ne j$, the summand is:

$${\mathbb{E}}_{\mathcal{E}}\left[\mathcal{E}(i)\right]\cdot {\mathbb{E}}_{\mathcal{E}}\left[\mathcal{E}(j)\right]\cdot {\mathbb{E}}_{\mathcal{E}}\left[x_i{x}_j^{\prime }\cdot {\delta }_{h(i),h(j)}\right]$$

These are clearly $0$ because $\mathbb{E}\left[\mathcal{E}(i)\right]=0$. So the original summation reduces to:

$$\begin{array}{rl}{\mathbb{E}}_{\varphi }\left[⟨x,x^{\prime }{⟩}_{\varphi }\right]& ={\mathbb{E}}_{\varphi }\left[\sum _{i=1,...,d}\mathcal{E}(i{)}^2\cdot x_i{x}_i^{\prime }\right]\\ & =⟨x,x^{\prime }⟩\end{array}$$

Not only is the hashed inner product unbiased, it also has a variance that scales down in $O(\frac{1}{m})$. The proof does a similar but more tedious expansion as the above, and assumes that $x,x^{\prime }$ have l2-norm of $1$. This suggests that the hashed inner product will be concentrated within $O(\frac{1}{\sqrt{m}})$ of the true value.

These results are sufficient to justify use of the hashed inner product space in practice. That is, we can perform recommendations in the hashed space with sufficiently large $m$ (we can tune that using validation error) to make the large feature space tractable. The paper goes on to prove more detailed bounds on the error and norm which are of less practical significance.

Multi-task Learning

The authors argue that this method is especially useful in the multi-task learning setting. Consider an email spam classification task where the vocab space is $\mathcal{V}$ and the user space is $\mathcal{U}$. The parameter space is thus $\mathcal{V}\times \mathcal{U}$, i.e. we wish to learn a user-specific weight vector $w_u\in {\mathbb{R}}^{|\mathcal{V}|}$ for each user $u$, which allows us to personalize the spam filter for each user (different users have slightly differing definitions of what is spam).

The authors suggest the following approach:

• Use the hashing trick to hash each term $v$ into the hashed space. e.g. data is passed into a global hash function ${\varphi }_0$ and assigned to a position
• Each user gets his/her own hash function ${\varphi }_u$. This may be implemented by using the same hash function but appending the user_id like so: user1_data, which hashes the same term into a new position.
• We may thus represent each instance by ${\varphi }_0(x)+{\varphi }_u(x)\in {\mathbb{R}}^m$, capturing both a global element (some terms are universally spam-indicative) and a personalized element (some terms are specifically indicative for a user)
• Finally, we learn a weight parameter $w_h\in {\mathbb{R}}^m$ by training it in the hashed space

Empirically, for their task of $|\mathcal{V}|=\text{40 million}$, $|\mathcal{U}|=\text{400,000}$, they found that performance starts to saturate with $m\approx \text{4 million}$. This is a very small fraction of the total space $|\mathcal{V}|\times |\mathcal{U}|$, showing the effectiveness of their method. Nevertheless, we should note that 4 million is still a rather large space.

From RankNET $\to$ LambdaRank $\to$ LambdaMART

Paper Link.

This is an overview paper that explains the model behind LambdaMART, a technique used to learn Gradient Boosted Trees that optimize for NDCG on recommendation-type of datasets.

RankNET

RankNET is a generic framework for training a learn to rank model. Given a differentiable model $f:{\mathbb{R}}^d\mapsto \mathbb{R}$ that produces a score from an n-dimensional input vector, RankNET is able to train $f$ such that for a given query session with items $i=1,...n$ and corresponding features $x_i\in {\mathbb{R}}^d:i=1,...,n$, $f$ learns to predict the relationship $f(x_i)>f(x_j)$ for any two items $i,j$ in the query session when $i$ is a better recommendation than $j$. The differentiable model $f$ is typically a neural network or boosted trees.

RankNET uses revealed pair-wise preferences within each query session to train $f$. Specifically, suppose we have a data for one query session as follows:

We can transform this data into a pairwise dataset as follows, where $y^{jk}$ denotes the preference relationship between $ite{m}^j$ and $ite{m}^k$ which we inferred from the click data. Note that the pairwise comparisons are only made within the same query session (e.g. qid1), as it reflects a given user's preferences in the context of a query and the items impressed to him/her in that session.

The pairwise setting is now more amenable to modelling (compared to directly optimizing for a good ranking), since we can now treat the task as a classification problem. For each row of the pairwise dataset, we only need to model the probability that $ite{m}^j$ is preferred (or not) to $ite{m}^k$. This can be formalized using a cross entropy loss comparing the predicted preference of our model to the revealed preference in the dataset.

First, we model the predicted probability from the model. Given row $i$ of the pairwise dataset and $ite{m}^j$ and $ite{m}^k$ respectively, we model the predicted probability that $ite{m}^j$ is preferred to $ite{m}^k$ (using $▹$ to denote a preference relationship) by passing the score difference between the predicted scores ${\hat{y}}_i^j:=f(x_i^j)$ and ${\hat{y}}_i^k:=f(x_i^k)$ for items j and k respectively through a sigmoid function, like so:

$${\hat{P}}_i^{jk}:=\hat{P}(ite{m}_i^j▹ite{m}_i^k)=\frac{1}{1+exp(-a({\hat{y}}_i^j-{\hat{y}}_i^k))}$$

Now let us denote the revealed probability that $ite{m}^j$ is preferred to $ite{m}^k$ as $P_i^{jk}$ such that:

• $P_i^{jk}=1$ if we prefer item j to item k
• $P_i^{jk}=0.5$ if we have no preference between the two items
• $P_i^{jk}=0$ if we prefer item k to item j

The cross entropy loss of our model can then be expressed as:

$$L=\sum _i\left[-P_i^{jk}log{\hat{P}}_i^{jk}-(1-P_i^{jk})log(1-{\hat{P}}_i^{jk})\right]$$

For convenience, let us denote $y_i^{jk}:=2P_i^{jk}-1$ (and conversely, $P_i^{jk}=\frac{y_i^{jk}+1}{2}$), which translates into the following:

• $y_i^{jk}=1$ if we prefer item j to item k
• $y_i^{jk}=0$ if we have no preference between the two items
• $y_i^{jk}=-1$ if we prefer item k to item j

Let us also define the convenience variable $z:=-log{\hat{P}}_i^{jk}=log\left[1+exp(-a({\hat{y}}_i^j-{\hat{y}}_i^k))\right]$. The cross entropy loss then simplifies to: $$\begin{array}{rl}L& =\sum _i\left[-P_i^{jk}log{\hat{P}}_i^{jk}-(1-P_i^{jk})log(1-{\hat{P}}_i^{jk})\right]\\ & =\sum _i\left[\frac{1}{2}(1+y_i^{jk})z-\frac{1}{2}(1-y_i^{jk})\left[(-a)({\hat{y}}_i^j-{\hat{y}}_i^k)-z\right]\right]\\ & =\sum _i\left[z+\frac{a}{2}(1-y_i^{jk})({\hat{y}}_i^j-{\hat{y}}_i^k)\right]\end{array}$$

Note that in line 2 of the above, we use the useful identity that $log(1-sigmoid(x))=x+log(sigmoid(x))$ and $log{\hat{P}}_i^{jk}=-z$. In line 3, the first and last term of line 2 cancel out to simply return $z$.

Having written out the loss function, we now need to differentiate the loss with respect to the model scores and parameters to obtain the gradient descent formula used to train the RankNET model. Differentiating $L$ wrt ${\hat{y}}_i^j$ and ${\hat{y}}_i^k$ gives:

$$\begin{array}{rl}\frac{\partial L}{\partial {\hat{y}}_i^j}& =\frac{a}{2}(1-y_i^{jk})+\frac{-a\cdot exp(-a({\hat{y}}_i^j-{\hat{y}}_i^k))}{1+exp(-a({\hat{y}}_i^j-{\hat{y}}_i^k))}\\ & =\frac{a}{2}(1-y_i^{jk})-\frac{a}{1+exp(a({\hat{y}}_i^j-{\hat{y}}_i^k))}\\ & =-\frac{\partial L}{\partial {\hat{y}}_i^k}\end{array}$$

Note that the first line of the above uses the result $\frac{d}{dx}lnf(x)=\frac{f^{\prime }(x)}{f(x)}$. We obtain line 2 by multiplying the right term by $exp(a({\hat{y}}_i^j-{\hat{y}}_i^k))$ in both the numerator and denominator. We obtain line 3 by observing that $L$ is a function of $d:={\hat{y}}_i^j-{\hat{y}}_i^k$, such that $\frac{\partial L}{\partial {\hat{y}}_i^j}=\frac{\partial L}{\partial d}\cdot \frac{\partial d}{\partial {\hat{y}}_i^j}=\frac{\partial L}{\partial d}$ and likewise $\frac{\partial L}{\partial {\hat{y}}_i^k}=\frac{\partial L}{\partial d}\cdot \frac{\partial d}{\partial {\hat{y}}_i^k}=-\frac{\partial L}{\partial d}$. The symmetry of the derivative wrt $j$ and $k$ will be important for the next section on factorizing RankNET to speed it up.

Finally, we use the gradient to update individual parameters $w_l\in \mathbb{R}$ of the model $f$. In the below, $n$ denotes the number of data points in the pairwise dataset. This update procedure rounds up the discussion on RankNET and is sufficient for training a generic differentiable model $f$ from ranking data. $$\begin{array}{rl}{w}_l& \leftarrow w_l-\eta \frac{\partial L}{\partial w_l}\\ & =w_l-\frac{\eta }{n}\sum _i\left[\frac{\partial L}{\partial {\hat{y}}_i^j}\cdot \frac{\partial {\hat{y}}_i^j}{\partial w_l}+\frac{\partial L}{\partial {\hat{y}}_i^k}\cdot \frac{\partial {\hat{y}}_i^k}{\partial w_l}\right]\end{array}$$

Factorizing RankNET

Schroff 2015 - FaceNET

Schroff 2015 - FaceNet: A Unified Embedding for Face Recognition and Clustering

This paper proposes the Triplet Loss for learning a face recognition system.

The task is that given a large number of person identities with a number of images associated with each person, to learn a model representation for an image in euclidean space, such that images belonging to the same person are close by and images for different persons are far away. This paper was impactful because it improved the SOTA on face verification by a large margin.

At this point, representation learning often trained a CNN classification model to classify images to a known identity. A bottleneck layer of relatively low dimensionality in the middle of the network is chosen as the representation of an image. In contrast to this indirect method, this paper directly optimizes the representation using contrastive learning.

Setup

Let the embedding of image $x$ be represented by $f(x)\in {\mathbb{R}}^d$, and constrain $||f(x)|{|}_2=1$. Now for a given person i, we want the anchor image $x_i^a$ to be closer to other images of the same person $x_i^p$ (positive) than images of any other person $x_i^n$ (negative), with some margin $\alpha$. That is, we desire a model $f$ such that:

$$\begin{array}{rlr}& ||f(x_i^a)-f(x_i^p)|{|}_2^2+\alpha <||f(x_i^a)-f(x_i^n)|{|}_2^2& \forall (x_i^a,x_i^p,x_i^n)\in \mathrm{T}\end{array}$$

Where $\mathrm{T}$ denotes the set of all possible triplets in the dataset. The authors note that most triplets would easily satisfy the desired constraint, hence we need some way to picking good triplets to optimize the learning process.

The loss we desire to minimize is thus as below. The intuition is that we wish to have small distances for the anchor-positive pair $d_{ap}$ but large distances for the anchor-negative pair $d_{an}$.

• Easy triplets where $d_{an}-d_{ap}>\alpha$ will contribute 0 loss
• Semi-hard triplets where $0<d_{an}-d_{ap}\le \alpha$ contribute small losses
• Hard triplets where $d_{an}-d_{ap}<0$ contribute large losses

$$\mathcal{L}=\sum _{i=1}^N{\left[||f(x_i^a)-f(x_i^p)|{|}_2^2-||f(x_i^a)-f(x_i^n)|{|}_2^2+\alpha \right]}_{+}$$

Triplet Selection

As with most contrastive learning methods, the authors observe that the selection of triplets is important for learning well. In an ideal world, we would want to select the hardest negatives for each anchor, i.e.: $$x_i^n=argmi{n}_{x_i^n}||f(x_i^a)-f(x_i^n)|{|}_2$$

However, this is undesirable because:

• We may end up mostly selecting mislabelled instances or corrupted images
• It is computationally infeasible to search across the whole corpus
• Trying to learn hard negatives at the start of training may cause learning to collapse

Hence, the authors propose:

• Mini batch negatives. Instead of finding the hardest negatives across the entire corpus, they mine for the hardest negatives in the mini-batch. This improves computational efficiency and mitigates the mislabelled/corrupted data issue.
• Curriculum Learning. At the start of training, the authors select easier triplets, and gradually increase the difficulty as training goes along. Presumably, this is done by thresholding negatives based on $d_{an}-d_{ap}>t$ and starting from high $t$ to low $t$.
• Semi-hard negatives. The authors mention that it is helpful to limit negatives to $d_{an}-d_{ap}>0$, but it is unclear whether they always do this or only at the start of training. The intuition is to cap the difficulty of the triplets, as triplets that are too difficult may actually hinder learning.

Application to Semantic Search

The triplet loss may be applied directly to semantic search. However, it is important to note that the paper assumes that for each person, all others-labelled instances are negatives. This is a suitable assumption for face recognition as each image can only belong to one person, but it is not true for semantic search, where a document may be relevant for multiple queries. Hence, the mislabelling issue when mining hard negatives is amplified for semantic search. I imagine that the selection of accurate negatives for semantic search would require more verification and filtering.

Recommendations as Treatments

Paper Link.

Training and evaluation data for recommender systems is subject to selection bias, because the probability of observing a data point depends on (i) users interacting with items which they like and (ii) the recommender system pushing items which they think the user likes. This leads to data Missing Not at Random (MNAR) and leads to biased model training and evaluation.

Consider users interacting with movies. Denote users $u\in \{1,...,U\}$ and movies $i\in \{1,...,I\}$. Let $Y\in {\mathbb{R}}^{U\times I}$ denote the true rating / interaction matrix between user and item, and $\hat{Y}$ the predicted matrix. Let $O\in \{0,1{\}}^{U\times I}$ be the observability matrix of whether an interaction was observed, and $P\in {\mathbb{R}}^{U\times I}$ be the probability of observation, i.e. $P_{u,i}=P(O_{u,i}=1)$. Ideally, we want the propensity matrix P to be uniform across all entries, which gives us the Missing Completely at Random (MCAR) condition.

Evaluating only on observed entries is biased

Given a model $\hat{Y}$, we often want to compute evaluation metrics ${\delta }_{u,i}(\hat{Y},Y)$ on how well $\hat{Y}$ approximates $Y$. For example:

$${\delta }_{u,i}(\hat{Y},Y)=(Y_{u,i}-{\hat{Y}}_{u,i}{)}^2\text{MSE}$$

$${\delta }_{u,i}(\hat{Y},Y)=\frac{Y_{u,i}}{logrank({\hat{Y}}_{u,i})}\text{DCG}$$

The Risk function may then be denoted:

$$R(\hat{Y})=\frac{{\sum }_u{\sum }_i{\delta }_{u,i}(\hat{Y},Y)}{U\times I}$$

However, $R(\hat{Y})$ cannot be computed because most entries in $Y$ are missing. Typically, we estimate $\hat{R}$ using only observed entries. This estimator is naive because it simply assumes the propensity matrix $P$ is uniform. This assumption is often false in recommender systems because a small set of popular items tend to get impressed a lot. Hence this estimator will favour models that lean towards the popular items compared to models that recommend rarer / new items.

$${\hat{R}}_{\text{naive}}(\hat{Y})=\frac{1}{|\{(u,i):O_{u,i}=1\}|}\cdot \sum _{u,i:O_{u,i}=1}{\delta }_{u,i}(\hat{Y},Y)$$

Unbiased estimators

The main idea in this paper is to view this problem as analogous to estimating treatment effects in causal inference. Think of recommending an item as an intervention analogous to treating a patient with a specific drug. The goal is to estimate the outcome of a new recommendation (clicked or not) or new treatment (recovered or not), while most outcomes between u,i pairs are not known.

The key to resolving this bias is to understand the assignment mechanism that generated $O$, namely the propensity matrix $P$. We can then correct for the propensity. We need to assume that $P_{u,i}>0\forall u,i$, i.e. full support, because otherwise the IPS estimator below is undefined.

IPS Estimator

The main estimator is the Inverse Propensity Score (IPS) estimator. Assuming that the assignment mechanism $P$ is known:

$${\hat{R}}_{\text{ips}}(\hat{Y}|P)=\frac{1}{|\{(u,i):O_{u,i}=1\}|}\cdot \sum _{u,i:O_{u,i}=1}\frac{{\delta }_{u,i}(\hat{Y},Y)}{P_{u,i}}$$

We have simply normalized each score by its inverse propensity, so that popular items with a high chance of being shown get their score reduced proportionally (and likewise in the opposite direction for rare items).

We can show that the IPS estimator is unbiased as we take the expectation over the random observability matrix. That is, suppose we are allowed to sample an infinitely large number of observations based on $P_{u,i}$, the average of the IPS estimator over these datasets will be the true risk function. This simply happens because the expected value of $O_{u,i}$ is simply the propensity, and if we know the propensity we can just cancel it out.