Treatment Assignment Strategy in A/B Test | by Chong Han Khai | Towards Data Science

In my previous post, I gave a step by step process of planning an A/B Test but did not cover the implementation part. This week, I focus on one of the components of an A/B Test implementation, treatment assignment.
One of the requirements of a valid randomized controlled experiment is random treatment assignment. Despite its importance, this is unfortunately also the part that most people seem to not care too much about. Perhaps because the concept of random seems so easy, I just randomly assign my experiment subjects to the control and treatment groups and that’s it, right? In a business environment, the reality is actually far from this because there are some requirements that make randomness hard to achieve.
In this blog post, I will attempt to highlight some of these requirements, common strategies people use to fulfill them, problems that arise with these strategies and lastly a simple and elegant solution that fulfills all the requirements.
Treatment given to an experiment subject has to be the consistent every time.
I will give a few examples here:
If Netflix wants to measure the effect of its new recommendation engine, then each user (experiment subject) should receive recommendations from the same recommendation engine as previous times.If I want to test if a new mobile app design has a faster loading time, then each device (experiment subject) needs to load the same mobile app design every time.
If Netflix wants to measure the effect of its new recommendation engine, then each user (experiment subject) should receive recommendations from the same recommendation engine as previous times.
If I want to test if a new mobile app design has a faster loading time, then each device (experiment subject) needs to load the same mobile app design every time.
If we randomly assign treatment every time when an experiment subject is exposed, then we will not be able to guarantee this. (e.g. user A might log in and see different treatments every time.)
For ease of explanation and the fact that most experiments are done using users as subjects, I am going to treat users like experiment subjects henceforth.
It is rather easy to fulfill the first requirement, we just log down which user is given which treatment when they are first exposed to the experiment. In subsequent login or subsequent exposure, we just look up for this user’s past treatment and assign the same treatment.
This may be feasible for a small company, but as you scale, you will find that this method is not very scalable at all. We can imagine that it is almost impossible to achieve this when we have tens of millions of users and hundreds of experiments.
In order to solve both of the issues above, one of the common strategies that companies use is to utilize the last character or digit of the user ID. In this case, treatment will be consistent and fully scalable because the treatment only depends on a fixed user ID.
This, however, poses another problem. The experiment, by itself, is indeed truly random. But as we do more and more experiments, we will find that the treatment assignments across different experiments are now correlated with each other.
Let’s use an example here.
+-------------------+-----------------+-------------------+-------+|                   | Exp 1 : Control | Exp 1 : Treatment | Total |+-------------------+-----------------+-------------------+-------+| Exp 2 : Control   |            2107 |              2929 |  5036 || Exp 2 : Treatment |            2916 |              2048 |  4964 || Total             |            5023 |              4977 | 10000 |+-------------------+-----------------+-------------------+-------+
Imagine if we have 10,000 users and we did two experiments. In the first experiment (Exp 1), we tested if a reduction in subscription fee increases the conversion rate of paid members and found a significant result. In the second experiment (Exp 2), we tested if an additional feature increases the number of purchases our users make and found that the effect is not statistically significant.
Here is why this might be the case:
Paid users make more purchases because of discounts.More users in the treatment group in Exp 1 are assigned control group in Exp 2, which implies that Exp 2 control group has more paid users.Users in Exp 2 control group naturally will make more purchases than the treatment group if both groups are given the placebo.Hence, even if the additional feature in Exp 2 did increase the number of purchases, we will not be able to measure the effect accurately because the treatment assignment in Exp 2 is correlated to treatment assignment in Exp 1.
Paid users make more purchases because of discounts.
More users in the treatment group in Exp 1 are assigned control group in Exp 2, which implies that Exp 2 control group has more paid users.
Users in Exp 2 control group naturally will make more purchases than the treatment group if both groups are given the placebo.
Hence, even if the additional feature in Exp 2 did increase the number of purchases, we will not be able to measure the effect accurately because the treatment assignment in Exp 2 is correlated to treatment assignment in Exp 1.
To be truly random, there should not be correlation between treatment assignment in different experiments.
To examine this, I performed the following simulation.
Generate 10,000 user IDs where each id is a hex string of length 10 (Example : d3ef2942d7).
Generate 10,000 user IDs where each id is a hex string of length 10 (Example : d3ef2942d7).
import pandas as pdimport randomexp_df = pd.DataFrame(['%010x' % random.randrange(16**10) for x in range(10000)], columns=['user_id'])
2. Create 20 experiments. For each of the experiments, randomly select 8 out of the 16 possible values (0 - 9, a - f) as control. If the last character of a user ID is in the control, assign them the control group, otherwise, assign them to treatment group.
import numpy as np# number of experimentsj = 20# get all 16 possible valuesrandom_list = exp_df['user_id'].apply(lambda x: x[-1]).unique()control_list = [set(np.random.choice(random_list, 8, replace=False)) for x in range(j)]treatment_list = [set(random_list) - x for x in control_list]    for k in range(j):    exp_df[f'exp_{k}'] = exp_df['user_id'].apply(lambda x: x[-1]).isin(control_list[k]).map({True:'control', False: 'treatment'})
3. For each pair of experiments, generate a contingency table like the one shown in the example above and perform a chi-square test.
from scipy import stats# initialize list to store chi-square results all experiment pairschi2_res = []for cols in combinations(exp_df.columns[1:], 2):    target_cols = list(cols)# generate contingency table    aggregate_df = exp_df[target_cols]\                   .groupby(target_cols)\                   .size()\                   .to_frame()\                   .reset_index()\                   .pivot(index=target_cols[1],                            columns=target_cols[0])# store chi-square test result    chi2_res.append(stats.chi2_contingency(aggregate_df)[1])# number of pairs that fail the chi-square test at alpha = 0.01print((np.array(chi2_res) < 0.01).sum())
Intuitively, if the treatment assignments are truly independent, we should expect to see the following (assuming 50% of the samples are assigned to control and 50% to treatment).
+-------------------+-----------------+-------------------+-------+|                   | Exp 1 : Control | Exp 1 : Treatment | Total |+-------------------+-----------------+-------------------+-------+| Exp 2 : Control   |            2500 |              2500 |  5000 || Exp 2 : Treatment |            2500 |              2500 |  5000 || Total             |            5000 |              5000 | 10000 |+-------------------+-----------------+-------------------+-------+
A chi-square test tells us how far the actual count is, from the independent case. The further away it is, the higher the chi-square test statistic and hence there is stronger evidence that there is some association between treatment assignments of the two experiments.
I ran the above simulation 100 times and see how many pairs of experiments failed the chi-square test within each simulation and got a mean of 62% (118 out of 190 pairs)! In other words, if we use this approach, on average, we will observe 62% of the experiment pairs failing the chi-square test. I also tried with 5 or 10 experiments (instead of the original 20) and got similar results (~60%).
The solution to fulfilling all the requirements above is actually very simple. All we have to do is to use a hash function. A hash function is any function that can be used to map data of arbitrary size to fixed-size values. This property ensures that the first two requirements (consistency and scalability) is fulfilled.
In order to fulfill the final requirement, instead of hashing only the user ID, we just have to hash the concatenation of user ID and experiment ID.
Below is an example of how to do this using python.
import mmh3def treatment_assignment(user_id, exp_id, control_bucket):        # calculates the number of buckets    num_buckets = len(control_bucket) * 2        # this generates a 32 bit integer    hash_int = mmh3.hash(user_id  + exp_id)        # get the mod of the 32 bit integer    mod = hash_int % num_buckets        if mod in control_bucket:        return 'control'    else:        return 'treatment'    # create 50 random integer as control groupcontrol_bucket = np.random.choice(np.arange(0,100,1), 50, replace=False)treatment_assignment('d3ef2942d7', 'exp_1', control_bucket)
The above function will return either treatment or control depending on your control bucket and experiment ID. We can also expand this to more-than-two-variant experiments.
I ran the above simulation again using this approach and below is the result.
We can see that 73 out of the 100 simulations result in 2 or fewer pairs of experiments (out of 190 pairs) failing the chi-square test, proving that it is a robust way of assigning treatments without violating any of the requirements.
In this blog post, I gave a detailed overview of the requirements of treatment assignment in experiments. I listed down some common practices I have seen and stated why they violate some of the requirements. Violating any of the requirements will mean the experiment results are no longer valid, which is why we should start putting more attention to treatment assignment.
Lastly, I gave a simple solution that fulfills all the requirements that are generalizable to more-than-two-variant experiments.
[1] A good hash is hard to find (OfferUp)https://blog.offerup.com/a-good-hash-is-hard-to-find-60e8a201e8ce