Методы балансировки в А/Б тестировании19.12.2023 10:15

Привет, Хабр! Как часто вы думаете о балансе? Балансе вселенной, личной жизни и работы, балансе БЖУ в своем рационе или балансе в банке. Мы в команде ad-hoc X5 Tech не только думаем о балансе, но и сталкиваемся с ним в работе. Сегодня поговорим о балансировке при анализе причинности. Это важный инструмент статистики, который помогает нам выяснить, как одни величины влияют на другие. Балансировка здесь — это способ убрать ошибки, которые могут возникнуть из-за разных распределений переменных в разных группах. Расскажем о различных методах балансировки, об их работе, преимуществах и недостатках каждого. Также затронем проблемы и ограничения, связанные с балансировкой. Запасайтесь чаем, мы начинаем!

Введение

В идеальном мире мы можем контролировать все факторы, которые влияют на объекты нашего исследования. Оценить влияние интересующего нас фактора в таких условиях несложно. Разделим объекты на две группы: контрольную и тестовую. Для первой группы оставим все факторы неизменными, а для второй изменим лишь интересующий нас фактор. В этом случае вся разница между группами будет заключаться в исследуемом факторе и, проведя эксперимент, мы установим его влияние на наши объекты.

В жизни не все факторы подвластны нашему контролю, из-за чего возникают более интересные задачи.

Классическими являются задачи проведения рандомизированных контролируемых исследований (randomized controlled trials). При такой постановке объекты исследования могут отличаться друг от друга по известным и неизвестным факторам, которые могут влиять на результаты исследования. Для контроля статистического влияния неподвластных нам факторов, объекты эксперимента случайно распределяются между контрольной и тестовой группой.

Подобные исследования стали популярны в начале прошлого века, когда Рональд Фишер решил проверить эффективность удобрений. Он столкнулся с проблемой: на рост растений влияют разные факторы, такие как солнце, ветер, свойства почвы и влажность. Устранить их влияние невозможно. Тогда Фишер предложил решение: разделить поле на участки и случайным образом определять, получит ли участок удобрение (то есть попадет в тестовую группу) или нет (останется в контрольной группе). Таким образом, эффективность удобрения можно оценить по разнице между урожаями тестовых и контрольных участков.

В реальности даже распределение по группам может оказаться вне нашего влияния, но оно может быть похожим на случайное. В таком случае мы имеем дело с естественными экспериментами (Natural experiments).

Но что делать, если распределение по группам не находится под нашим влиянием и даже не напоминает случайное? Именно в этом случае методы балансировки приходят к нам на помощь.

Допущения

Стоит определиться, в каких ситуациях корректно применять методы балансировки. Сформулируем допущения на условия эксперимента, при которых можно применять методы, описанные далее в нашей статье.

Stable unit treatment value assumption (SUTVA) — эффект на каждый объект оказывается независимо, то есть у нас нет «сетевых эффектов». Дополнительно предполагается, что уровень воздействия для каждого объекта схож — например, программа переобучения имеет сравнительно одинаковое качество для всех участников в каждом из регионов работы программы.

Ignorability — условная независимость потенциальных результатов (conditional independence of the potential outcomes). Условная независимость позволяет нам измерить влияние на результат исследуемого фактора, а не какой-либо другой переменной, скрывающейся вокруг. Это непростая, но важная концепция, поэтому далее в примере мы проиллюстрируем, что происходит, если ее не придерживаться.

$Y(0), Y(1)\perp\!\!\!\perp Z | X$

где Y(0) — значение целевой метрики при отсутствии воздействия исследуемого фактора, Y(1) — значение целевой метрики при воздействии исследуемого фактора, — статус присвоения к контрольной или тестовой группе, — все остальные известные признаки.

Overlap — означает, что у каждого объекта должен быть хоть какой-то шанс попасть в тестовую группу и попасть в контрольную группу. То есть у нас точно нет наблюдений, которые автоматически попадают в одну из групп. Если данное предположение не выполняется, то мы не можем обобщать результаты. К примеру, если исследование проводилось с людьми разных полов, но в тестовую группу попали только мужчины, то некорректно обобщать результаты эксперимента на людей обоих полов.

0 < Pr(Z = 1 | X) < 1

Обещанный пример о важности Ignorability

Пусть мы тестируем эффективность препарата. Рассмотрим следующую ситуацию:

У нас есть один признак, влияющий на выздоровления — уровень поддержки семьей и друзьями. Те, кого поддерживают, выздоравливают с вероятностью 0.6. Кого не поддерживают — с вероятностью 0.4;
Наше экспериментальное лечение не влияет на вероятность выздоровления;
Врач знает про пользу поддержки. И считает, что лечение чаще нужно назначать поддерживаемым пациентам (врач с вероятностью 0.9 отправит пациента с поддержкой на новое лечение и с вероятностью 0.5 отправит пациента без поддержки).

Что будет в случае случайного и неслучайного распределения по группам?

Распишем математическое ожидание выздоровления в случае, когда лекарство выдается случайно, то есть врач не влияет на назначение лекарства:

$\mathbb{E} (Y | Z = 0) = \sum_{x = 0}^{1} \mathbb{E} (Y | Z = 0, X = x) Pr (X = x | Z = 0) = 0.4 \cdot 0.5 + 0.6 \cdot 0.5 = 0.5$

Так как распределение случайное, то формула для Z=1 аналогична. То есть разница между тестовой и контрольной группами — 0 (лечение не влияет на вероятность выздороветь).

Теперь посчитаем математические ожидания выздоровления пациента из тестовой и контрольной групп, когда решение о назначении лекарства принимает врач:

$\mathbb{E} (Y | Z = 1) = \frac{\sum_{x = 0}^{1} \mathbb{E} (Y | Z = 1, X = x) Pr (Z = 1 | X = x) Pr (X = x)}{\sum_{x = 0}^{1} Pr (Z = 1 | X = x) Pr (X = x)} =$ $= \frac{0.4 \cdot 0.5 \cdot 0.5 + 0.6 \cdot 0.9 \cdot 0.5}{0.5 \cdot 0.5 + 0.9 \cdot 0.5} = 0.53$ $\mathbb{E} (Y | Z = 0) = \frac{\sum_{x = 0}^{1} \mathbb{E} (Y | Z = 0, X = x) Pr (Z = 0 | X = x) Pr (X = x)}{\sum_{x = 0}^{1} Pr (Z = 0 | X = x) Pr (X = x)} =$ $= \frac{0.4 \cdot 0.5 \cdot 0.5 + 0.6 \cdot 0.1 \cdot 0.5}{0.5 \cdot 0.5 + 0.1 \cdot 0.5} = 0.43$

Очевидная разница между тестовой и контрольной группами — 0.1 (казалось бы, лечение влияет на вероятность выздороветь).

Что такое балансировка

Как продемонстрировано в примере выше, при наличии смещений в группах можно получать неверные результаты. При балансировке мы хотим приблизить оценку эффекта к реальному значению, и баланс по ковариатам помогает нам достичь этой цели. При балансировке мы хотим получить некоторую оценку эффекта, группируя объекты или изменяя веса так, чтобы получить баланс по ковариатам, нашим признакам , между группами, что является воплощением простого принципа сравнения подобного с подобным — если группы схожи во всех отношениях, за исключением исследуемого фактора, то можно говорить, что различия результатов между группами объясняется именно им.

Propensity score

Propensity score впервые описывается в статье Rosenbaum P.R., Rubin D.B. «Assessing sensitivity to an unobserved binary covariate in an observational study with binary outcome» в 1983 году. Смысл propensity score — вероятность попасть в тестовую группу при условии наблюдаемых ковариат.

e_i = Pr(Z_i = 1 | X_i)

Но чем нам может помочь использование propensity score? А тем, что вместо контроля за всеми ковариатами (коих может быть очень много), будем контролировать только одно число — нашу условную вероятность попасть в тестовую группу.

Полезность propensity score можно показать с помощью графа причинности (causal graph) — ориентированного графа, отражающего причинно-следственные связи между переменными. Он включает в себя набор переменных (узлов), где каждый узел связан стрелкой с другими узлами, на которые он оказывает влияние. Если мы контролируем propensity score = P (x), то прямое использование не даст нам дополнительной информации о том, каким будет (назначенная группа). Можно воспринимать это следующим образом. Если у нас есть два объекта с одинаковым propensity score, и разными группами (один попал в тест, другой — в контроль), то можно считать, что попадание в тест или контроль было определено фактором случайности, как в случае классического рандомизированного эксперимента.

Методы стратификации

Одним из вариантов балансировки является стратификация. При стратификации выборка делится на некоторый набор подгрупп с непересекающимися значениями ковариат. Подгруппы могут разделяться по некоторым признакам: полу, возрастной группе, доходам и так далее. При случайном разделении выборки может случиться, что группы будут отличаться по доле объектов из важной ковариаты (например, перекос в гендере), особенно на малых выборках.

Для снижения влияния дисбаланса классов на оценку, объекты могут быть включены в выборку с учетом стратификации, с соблюдением пропорции объектов из каждой страты в процессе сэмплирования. В случаях, когда проведение стратифицированного сэмплирования невозможно, оказывается, что все еще можно воспользоваться стратифицированным средним, взвешивая полученное среднее по весу каждой страты. Этот метод известен как пост-стратификация. Подробнее про стратификацию и про повышение чувствительности тестов при ее применении можно почитать в статье от X5 Tech.

Плюсы:

Включение требуемых подгрупп с учетом их распространенности особенно важно при ограниченном объеме выборки;
Уменьшение дисперсии и повышение чувствительности теста;
Снижение затрат на сбор данных и увеличение управляемости, например, проведение опроса в одном географическом регионе с последующим учетом диспропорции при дизайне эксперимента на всей территории.

Минусы:

При неправильной стратификации, например, при ошибочной оценке весов подгрупп, возможно смещение оценок;
С увеличением гранулярности становится сложнее получить корректную выборку. В крайнем случае, при выборе слишком многих признаков каждый элемент выборки может оказаться в уникальной подгруппе;
В случае нечетких страт, таких как этническая или религиозная принадлежность, могут возникнуть сложности при проведении эксперимента.

Методы мэтчинга

При условии, что объект характеризуется множеством ковариат, более предпочтительным методом балансировки данных может оказаться мэтчинг. Мэтчинг представляет собой инструмент, позволяющий сопоставить объектам из тестовой группы объекты из контрольной группы. Такой подход позволяет получить сбалансированные группы, сократить влияние факторов, которые могут внести потенциальный дисбаланс в выборку, и получить менее смещенную оценку эффекта.

В рамках проведения A/B-тестов это означает, что для каждого объекта из тестовой группы, подвергнутой воздействию, подбирается максимально похожий объект из потенциальной контрольной группы с целью исключить влияние таких факторов, как, например, география.

Общий процесс выглядит так:

Определение меры «близости» между объектами;
Реализация метода мэтчинга, учитывающего установленную меру «близости»;
Оценка качества полученного набора «близких» объектов (возможно, повторение шагов 1 и 2 до достижения «хорошего» набора);
Оценка разницы между полученными группами.

Существует значительное число различных методов мэтчинга, и в данной статье мы подробно рассмотрим только некоторые из них. Самым простым является Exact Matching, основная идея которого заключается в том, что для каждого объекта из тестовой выборки ищется идентичный объект в контрольной выборке. Если объект нашелся, то расстояние до него равно 0, иначе расстояние равно $\infty$ .

Развитием данного метода является Coarsened Exact Matching (CEM), в нем непрерывные признаки превращаются в категориальные, и вместо поиска идентичного объекта осуществляется поиск похожих объектов в «загрубленном пространстве».

В основе Mahalanobis Distance Matching (MDM) лежит расстояние Махаланобиса. Если говорить упрощенно, то в нем учитывается близость к центру масс с учетом разброса:

$D^2 = (x-m)^T * C^{-1} * (x-m)$

где D^2 — квадрат расстояния Махаланобиса, — вектор признаков, — вектор средних значений независимых переменных, $C^{-1}$ — обратная ковариационная матрица независимых переменных.

Последний рассматриваемый нами метод — это Propensity Score Matching (PSM). Суть метода в том, что строится модель определения propensity score на основе признаков и осуществляется поиск похожих объектов в данном пространстве. Подробнее про propensity score можно почитать в этом докладе. Про другие методы matching можно почитать по этой ссылке.

При использовании методов мэтчинга необходимо аккуратно проверять сопоставимость между группами. Если некоторые значения исключаются, следует внимательно контролировать размер оставшейся выборки, поскольку более маленькая выборка может исказить результаты. Кроме того, важно оценить распределение признаков в обеих группах, используя статистические тесты, такие как t-тест, тест Колмогорова-Смирнова, тест Хи-квадрат и другие. Хотя это может включать множественное сравнение, это позволит выявить различия между группами по отдельным признакам. Информация о сопоставимости также может быть извлечена из сравнения эмпирических функций распределения, плотностей распределения по признакам, standardized mean difference, variance ratio, propensity scores и из графиков эмпирических QQ-plots.

На графиках ниже представлены результаты проверки сопоставимости для набора данных из статьи LaLonde, R. (1986). На первом графике видно, что применение методов мэтчинга позволило сильно сбалансировать выборку.

На втором графике QQ-plots для параметров age и re75 выборки до и после мэтчинга не претерпели существенных изменений. Однако, для параметра married изначально наблюдались значительные выбросы, большая часть из которых была успешно исключена. Это позволило получить более сбалансированную выборку для данного параметра.

На третьем графике показано распределение propensity score. Из него видно, что для всех объектов из тестовой группы удалось сопоставить объекты из контрольной группы. После сопоставленных элементов распределение стало лучше, однако все равно можно выделить некоторый дисбаланс. Также заметно, что было исключено много объектов из потенциальной контрольной группы, поскольку вероятность попадания в тестовую группу у них практически нулевая. Это связано с тем, что они значительно отличаются от объектов тестовой группы, и рассматривать их в данном случае не имеет смысла.

Больше графиков, оценивающих сопоставимость между контрольной и тестовой группами, доступны в документации библиотеки MatchIt.

Плюсы:

Контроль за «похожестью» объектов, оценка баланса;
Эффективная комбинация с другими методами, такими как регрессионные методы;
Повышение чувствительности исследований;
Простота и гибкость метода с множеством вариаций;
Для Propensity Score Matching — поиск похожих объектов всего по одному параметру, а не по множеству.

Минусы:

Для успешного мэтчинга необходимо учитывать все важные ковариаты; пропуск важных факторов может привести к смещению;
Метод Propensity Score Matching требователен к объему данных;
С увеличением числа ковариат становится сложнее найти подходящие сопоставления;
Для Propensity Score Matching похожесть propensity score не обязательно гарантирует баланс ковариат.

Больше про критику PSM можно посмотреть в видео от Gary King.

Методы взвешивания

Третий метод балансировки — взвешивание. В то время как стратификация разделяет выборку на страты, а мэтчинг сопоставляет объекты, методы взвешивания предоставляют возможность получить псевдопопуляцию с лучшим балансом с помощью изменения весов наблюдений. Данный подход играет ключевую роль в уменьшении влияния факторов, которые могут повлиять на смещение оценки эффекта.

Рассмотрим следующий метод взвешивания — Inverse Probability of Treatment Weighting (IPTW). Суть данного метода в том, что в начале необходимо оценить propensity score, к примеру с помощью логистической регрессии. Далее, исходя из полученных propensity scores, происходит перевзвешивание наблюдений и производится оценка эффекта.

$\mathbb{IPTW} = \frac{1}{n} \sum_{i=1}^{n} \biggl (\frac{D_i*Y_i}{P(X_i)} - \frac{(1-D_i)*Y_i}{1-P(X_i)} \biggr )$

где P(X_i) — полученный propensity score, D_i — индикатор принадлежности -го объекта к тестовой группе, Y_i — значение целевой переменной (target) для -го объекта.

Суть перевзвешивания заключается в том, что если объект относится к тестовой группе, но его значение P(X_i) (вероятность принадлежности к тестовой группе) достаточно мало, то это указывает на то, что такие объекты в большинстве случаев попадают в контрольную группу, и, следовательно, им необходимо присвоить больший вес. Это особенно интересно в контексте анализа эффекта воздействия, так как объекты с низкой вероятностью попадания в тестовую группу представляют собой «граничные случаи». Напротив, если объект находится в тестовой группе и имеет довольно большое значение P(X_i) , то его вес следует уменьшить, так как таких объектов в тестовой группе и так много.

Более подробную информацию о применении методов взвешивания в контексте проведения A/B-тестов для оффлайн ритейла вы можете найти в статье от X5 Tech.

Плюсы:

Использование полной выборки без отбрасывания ее части, например, при применении методов, таких как clipping в методах мэтчинга;
Возможность комбинирования с другими методами, такими как doubly robust методы, для получения более корректных оценок;
Применимость для множественных тестовых групп.

Минусы:

Проблема возможного возникновения экстремальных весов;
Менее интуитивный и понятный алгоритм.

Примеры

Давайте более подробно рассмотрим, как упомянутые выше методы, представленные в статье, применяются на практике. Для оценки эффекта, воспользуемся библиотекой causallib для языка программирования Python. Будем проводить оценку эффекта на синтетическом наборе данных, который включает 2000 наблюдений (1000 в тестовой и 1000 в контрольной группах) и содержит 4 ковариаты. Целевая метрика представляет из себя сумму всех ковариат, кроме первой, и нормального шума. Первая ковариата представляет из себя категориальную переменную, принимающую значения 0 и 1; в случае если принимается значение 1, целевая метрика увеличивается на 25. В тестовую группу также добавим эффект, равный 10. Для обеспечения точности результатов каждую оценку будем повторять 50 раз, и затем усреднять.

Код

import numpy as np
import pandas as pd
from causallib.estimation import IPW, MarginalOutcomeEstimator, Matching
from causallib.preprocessing.transformers import PropensityTransformer
from sklearn.linear_model import LogisticRegression
from tqdm import tqdm

np.random.seed(80)


def generate_data(real_ate: float, sample_size: int) -> pd.DataFrame:
    """генерация синтетических данных"""
    df = pd.DataFrame(
        data={
            'treat': np.concatenate([np.zeros(sample_size), np.ones(sample_size)]),
            'cov1': np.random.randint(0, 2, 2*sample_size),
            'cov2': np.random.normal(20, 20, 2*sample_size),
            'cov3': np.random.normal(30, 20, 2*sample_size),
            'cov4': np.random.normal(40, 20, 2*sample_size),
        }
    )

    df['target'] = (
        25*df['cov1']
        + df['cov2']
        + df['cov3']
        + df['cov4']
        + np.random.normal(0, 10, 2*sample_size)
        + real_ate * df['treat']
    )
    return df


def estimate_diff_of_means(
    df: pd.DataFrame,
    target_col_name: str,
    treatment_indicator_col_name: str,
    covariate_col_names: str,
) -> float:
    """прямая разница между группами"""
    covariates, treatment_ind, target = (
        df[covariate_col_names],
        df[treatment_indicator_col_name],
        df[target_col_name],
    )
    moe = MarginalOutcomeEstimator(None).fit(covariates, treatment_ind, target)
    outcomes = moe.estimate_population_outcome(covariates, treatment_ind, target)
    effect = moe.estimate_effect(outcomes[1], outcomes[0])["diff"]

    return effect


def estimate_stratification_ate(
    df: pd.DataFrame,
    target_col_name: str,
    treatment_indicator_col_name: str,
    covariate_col_names: str,
    strat_col: str,
) -> float:
    """пост-стратификация по полю strat_col"""
    weights = df.groupby(strat_col)[target_col_name].count() / len(df)

    treatment_stratified_avg = (
        df[df[treatment_indicator_col_name] == 1]
        .groupby(strat_col)[target_col_name]
        .mean()
        * weights
    ).sum()
    control_stratified_avg = (
        df[df[treatment_indicator_col_name] == 0]
        .groupby(strat_col)[target_col_name]
        .mean()
        * weights
    ).sum()
    effect = treatment_stratified_avg - control_stratified_avg

    return effect


def estimate_matching_ate(
    df: pd.DataFrame,
    target_col_name: str,
    treatment_indicator_col_name: str,
    covariate_col_names: str,
    metric: str = "mahalanobis",
    use_propensity_score: bool = False,
) -> float:
    """применение методов Matching"""
    covariates, treatment_ind, target = (
        df[covariate_col_names],
        df[treatment_indicator_col_name],
        df[target_col_name],
    )
    propensity_transform = None
    if use_propensity_score:
        propensity_transform = PropensityTransformer(
            learner=LogisticRegression(solver="liblinear", class_weight="balanced"),
            include_covariates=False,
        )

    matcher = Matching(propensity_transform=propensity_transform, metric=metric)
    matcher.fit(covariates, treatment_ind, target)

    outcomes = matcher.estimate_individual_outcome(covariates, treatment_ind)
    effect = matcher.estimate_effect(outcomes[1], outcomes[0])["diff"]

    return effect


def estimate_iptw_ate(
    df: pd.DataFrame,
    target_col_name: str,
    treatment_indicator_col_name: str,
    covariate_col_names: str,
) -> float:
    """применение методов взвешивания"""
    covariates, treatment_ind, target = (
        df[covariate_col_names],
        df[treatment_indicator_col_name],
        df[target_col_name],
    )
    ipw = IPW(learner=LogisticRegression(solver="liblinear", class_weight="balanced"))

    ipw.fit(covariates, treatment_ind)
    outcomes = ipw.estimate_population_outcome(covariates, treatment_ind, target)
    effect = ipw.estimate_effect(outcomes[1], outcomes[0])["diff"]

    return effect


COLS = [
    "treat",
    "cov1",
    "cov2",
    "cov3",
    "cov4",
    "target",
]

TARGET_COL_NAME = "target"
TREATMENT_INDICATOR_COL_NAME = "treat"
COVARIATE_COL_NAMES = COLS[1:-1]

methods = [
    ("Прямая разница между группами", estimate_diff_of_means),
    ("Mahalanobis Distance Matching", estimate_matching_ate),
    ("Пост-стратификация по полю cov1", estimate_stratification_ate, "cov1"),
    ("Euclidean Distance Matching", estimate_matching_ate, "euclidean"),
    ("Propensity Score Matching", estimate_matching_ate, "mahalanobis", True),
    ("IPTW", estimate_iptw_ate),
]

methods_res = {
    "Прямая разница между группами": [],
    "Mahalanobis Distance Matching": [],
    "Пост-стратификация по полю cov1": [],
    "Euclidean Distance Matching": [],
    "Propensity Score Matching": [],
    "IPTW": [],
}

iters = 50
ATE = 10
size = 1000
for i in tqdm(range(iters)):
    df = generate_data(ATE, size)

    for method, estimator, *args in methods:
        ate = estimator(
            df,
            TARGET_COL_NAME,
            TREATMENT_INDICATOR_COL_NAME,
            COVARIATE_COL_NAMES,
            *args,
        )
        methods_res[method].append(ate)

for method, vals in methods_res.items():
    print(f"{method: <35}{round(np.mean(vals), 2): >12}")

На основе результатов оценки можно сделать вывод о том, что в случае синтетического набора данных наилучшие результаты продемонстрировал метод IPTW. С его помощью удалось получить оценку, максимально близкую к истинному эффекту.

Подлинный ATE	10
Прямая разница между группами	9.55
Mahalanobis Distance Matching	9.89
Пост-стратификация по первой ковариате	9.61
Euclidean Distance Matching	9.93
Propensity Score Matching	9.93
IPTW	9.95

Однако, при использовании реального набора данных нельзя с однозначностью утверждать, какой метод предоставляет истинную оценку эффекта, так как методы различаются в подходах к формированию контрольной группы. Поэтому стоит опираться на предыдущий опыт и использовать то, что хорошо показывало себя на похожих задачах. Также целесообразно оценивать вероятности ошибок на исторических данных для исключения неподходящих методов. Более подробно про выбор метода для оценки можно прочитать в нашей статье в главе 1.5 «Propensity Score Weighting».

Заключение

Подведем итоги:

Методы балансировки основаны на следующей идее — если неким образом сбалансировать наборы данных для теста и контроля, то мы получим менее искаженную оценку эффекта;
Существуют различные способы балансировки: учет пропорций объектов в совокупности, подбор похожих объектов, изменение весов объектов при оценке результатов;
Для каждой отдельной задачи могут лучше сработать разные методы;
При использовании методов балансировки лучше дополнительно перепроверить, что не случилось «перекосов», влияющих на оценку.

Надеемся, статья была полезна, и будем рады дельным комментариям и замечаниям!

Над статьёй работали: Артём Ерохин, Денис Лавров, Андрей Виданов, Иван Щербак, Инсаф Рашитов.