def compute_dci(ground_truth_data, representation_function, random_state,
artifact_dir=None,
num_train=gin.REQUIRED,
num_test=gin.REQUIRED,
batch_size=16):
"""Computes the DCI scores according to Sec 2.
Args:
ground_truth_data: GroundTruthData to be sampled from.
representation_function: Function that takes observations as input and
outputs a dim_representation sized representation for each observation.
random_state: Numpy random state used for randomness.
artifact_dir: Optional path to directory where artifacts can be saved.
num_train: Number of points used for training.
num_test: Number of points used for testing.
batch_size: Batch size for sampling.
Returns:
Dictionary with average disentanglement score, completeness and
informativeness (train and test).
"""
del artifact_dir
logging.info("Generating training set.")
# mus_train are of shape [num_codes, num_train], while ys_train are of shape
# [num_factors, num_train].
mus_train, ys_train = utils.generate_batch_factor_code(
ground_truth_data, representation_function, num_train,
random_state, batch_size)
assert mus_train.shape[1] == num_train
assert ys_train.shape[1] == num_train
mus_test, ys_test = utils.generate_batch_factor_code(
ground_truth_data, representation_function, num_test,
random_state, batch_size)
scores = _compute_dci(mus_train, ys_train, mus_test, ys_test)
return scores
def compute_sap(ground_truth_data,
representation_function,
random_state,
artifact_dir=None,
num_train=gin.REQUIRED,
num_test=gin.REQUIRED,
batch_size=16,
continuous_factors=gin.REQUIRED):
"""Computes the SAP score.
Args:
ground_truth_data: GroundTruthData to be sampled from.
representation_function: Function that takes observations as input and
outputs a dim_representation sized representation for each observation.
random_state: Numpy random state used for randomness.
artifact_dir: Optional path to directory where artifacts can be saved.
num_train: Number of points used for training.
num_test: Number of points used for testing discrete variables.
batch_size: Batch size for sampling.
continuous_factors: Factors are continuous variable (True) or not (False).
Returns:
Dictionary with SAP score.
"""
del artifact_dir
logging.info("Generating training set.")
mus, ys = utils.generate_batch_factor_code(ground_truth_data,
representation_function,
num_train, random_state,
batch_size)
mus_test, ys_test = utils.generate_batch_factor_code(
ground_truth_data, representation_function, num_test, random_state,
batch_size)
logging.info("Computing score matrix.")
return _compute_sap(mus, ys, mus_test, ys_test, continuous_factors)
def compute_modularity_explicitness(ground_truth_data,
representation_function,
random_state,
artifact_dir=None,
num_train=gin.REQUIRED,
num_test=gin.REQUIRED,
batch_size=16):
"""Computes the modularity metric according to Sec 3.
Args:
ground_truth_data: GroundTruthData to be sampled from.
representation_function: Function that takes observations as input and
outputs a dim_representation sized representation for each observation.
random_state: Numpy random state used for randomness.
artifact_dir: Optional path to directory where artifacts can be saved.
num_train: Number of points used for training.
num_test: Number of points used for testing.
batch_size: Batch size for sampling.
Returns:
Dictionary with average modularity score and average explicitness
(train and test).
"""
del artifact_dir
scores = {}
mus_train, ys_train = utils.generate_batch_factor_code(
ground_truth_data, representation_function, num_train, random_state,
batch_size)
mus_test, ys_test = utils.generate_batch_factor_code(
ground_truth_data, representation_function, num_test, random_state,
batch_size)
discretized_mus = utils.make_discretizer(mus_train)
mutual_information = utils.discrete_mutual_info(discretized_mus, ys_train)
# Mutual information should have shape [num_codes, num_factors].
assert mutual_information.shape[0] == mus_train.shape[0]
assert mutual_information.shape[1] == ys_train.shape[0]
scores["modularity_score"] = modularity(mutual_information)
explicitness_score_train = np.zeros([ys_train.shape[0], 1])
explicitness_score_test = np.zeros([ys_test.shape[0], 1])
mus_train_norm, mean_mus, stddev_mus = utils.normalize_data(mus_train)
mus_test_norm, _, _ = utils.normalize_data(mus_test, mean_mus, stddev_mus)
for i in range(ys_train.shape[0]):
explicitness_score_train[i], explicitness_score_test[i] = \
explicitness_per_factor(mus_train_norm, ys_train[i, :],
mus_test_norm, ys_test[i, :])
scores["explicitness_score_train"] = np.mean(explicitness_score_train)
scores["explicitness_score_test"] = np.mean(explicitness_score_test)
return scores
def compute_downstream_task(ground_truth_data,
representation_function,
random_state,
artifact_dir=None,
num_train=gin.REQUIRED,
num_test=gin.REQUIRED,
batch_size=16):
"""Computes loss of downstream task.
Args:
ground_truth_data: GroundTruthData to be sampled from.
representation_function: Function that takes observations as input and
outputs a dim_representation sized representation for each observation.
random_state: Numpy random state used for randomness.
artifact_dir: Optional path to directory where artifacts can be saved.
num_train: Number of points used for training.
num_test: Number of points used for testing.
batch_size: Batch size for sampling.
Returns:
Dictionary with scores.
"""
del artifact_dir
scores = {}
for train_size in num_train:
mus_train, ys_train = utils.generate_batch_factor_code(
ground_truth_data, representation_function, train_size,
random_state, batch_size)
mus_test, ys_test = utils.generate_batch_factor_code(
ground_truth_data, representation_function, num_test, random_state,
batch_size)
predictor_model = utils.make_predictor_fn()
train_err, test_err = _compute_loss(np.transpose(mus_train), ys_train,
np.transpose(mus_test), ys_test,
predictor_model)
size_string = str(train_size)
scores[size_string + ":mean_train_accuracy"] = np.mean(train_err)
scores[size_string + ":mean_test_accuracy"] = np.mean(test_err)
scores[size_string + ":min_train_accuracy"] = np.min(train_err)
scores[size_string + ":min_test_accuracy"] = np.min(test_err)
for i in range(len(train_err)):
scores[size_string +
":train_accuracy_factor_{}".format(i)] = train_err[i]
scores[size_string +
":test_accuracy_factor_{}".format(i)] = test_err[i]
return scores
def compute_mig(ground_truth_data,
representation_function,
random_state,
artifact_dir=None,
num_train=gin.REQUIRED,
batch_size=16):
"""Computes the mutual information gap.
Args:
ground_truth_data: GroundTruthData to be sampled from.
representation_function: Function that takes observations as input and
outputs a dim_representation sized representation for each observation.
random_state: Numpy random state used for randomness.
artifact_dir: Optional path to directory where artifacts can be saved.
num_train: Number of points used for training.
batch_size: Batch size for sampling.
Returns:
Dict with average mutual information gap.
"""
del artifact_dir
logging.info("Generating training set.")
mus_train, ys_train = utils.generate_batch_factor_code(
ground_truth_data, representation_function, num_train,
random_state, batch_size)
assert mus_train.shape[1] == num_train
return _compute_mig(mus_train, ys_train)
def compute_mig(ground_truth_data,
representation_function,
random_state,
num_train=gin.REQUIRED,
num_bins=50,
batch_size=16):
"""Computes the mutual information gap.
Args:
ground_truth_data: GroundTruthData to be sampled from.
representation_function: Function that takes observations as input and
outputs a dim_representation sized representation for each observation.
random_state: Numpy random state used for randomness.
num_train: Number of points used for training.
batch_size: Batch size for sampling.
Returns:
Dict with average mutual information gap.
"""
#try:
# print("MIG rs: ", random_state)
#except Exception as e:
# print()
# print("Random seed exception :", e)
# print()
print("MIG batch_:", batch_size)
print("MIG n_train: ",num_train)
#import IPython
#IPython.embed()
#logging.info("Generating training set.")
try:
mus_train, ys_train = utils.generate_batch_factor_code(
ground_truth_data, representation_function, num_train,
random_state, batch_size)
except:
mus_train, ys_train = utils.generate_batch_factor_code(
ground_truth_data, representation_function, num_train,
np.random.get_state(), batch_size)
print("MUS train ", mus_train.shape, ' : should equal : ', "NUM TRAIN ", num_train)
#assert mus_train.shape[1] == num_train
return _compute_mig(mus_train, ys_train, num_bins = num_bins)
def unsupervised_metrics(ground_truth_data,
representation_function,
random_state,
artifact_dir=None,
num_train=gin.REQUIRED,
batch_size=16):
"""Computes unsupervised scores based on covariance and mutual information.
Args:
ground_truth_data: GroundTruthData to be sampled from.
representation_function: Function that takes observations as input and
outputs a dim_representation sized representation for each observation.
random_state: Numpy random state used for randomness.
artifact_dir: Optional path to directory where artifacts can be saved.
num_train: Number of points used for training.
batch_size: Batch size for sampling.
Returns:
Dictionary with scores.
"""
del artifact_dir
scores = {}
logging.info("Generating training set.")
mus_train, _ = utils.generate_batch_factor_code(ground_truth_data,
representation_function,
num_train, random_state,
batch_size)
num_codes = mus_train.shape[0]
cov_mus = np.cov(mus_train)
assert num_codes == cov_mus.shape[0]
# Gaussian total correlation.
scores["gaussian_total_correlation"] = gaussian_total_correlation(cov_mus)
# Gaussian Wasserstein correlation.
scores[
"gaussian_wasserstein_correlation"] = gaussian_wasserstein_correlation(
cov_mus)
scores["gaussian_wasserstein_correlation_norm"] = (
scores["gaussian_wasserstein_correlation"] / np.sum(np.diag(cov_mus)))
# Compute average mutual information between different factors.
mus_discrete = utils.make_discretizer(mus_train)
mutual_info_matrix = utils.discrete_mutual_info(mus_discrete, mus_discrete)
np.fill_diagonal(mutual_info_matrix, 0)
mutual_info_score = np.sum(mutual_info_matrix) / (num_codes**2 - num_codes)
scores["mutual_info_score"] = mutual_info_score
return scores
def compute_irs(ground_truth_data,
representation_function,
random_state,
artifact_dir=None,
diff_quantile=0.99,
num_train=gin.REQUIRED,
batch_size=gin.REQUIRED):
"""Computes the Interventional Robustness Score.
Args:
ground_truth_data: GroundTruthData to be sampled from.
representation_function: Function that takes observations as input and
outputs a dim_representation sized representation for each observation.
random_state: Numpy random state used for randomness.
artifact_dir: Optional path to directory where artifacts can be saved.
diff_quantile: Float value between 0 and 1 to decide what quantile of diffs
to select (use 1.0 for the version in the paper).
num_train: Number of points used for training.
batch_size: Batch size for sampling.
Returns:
Dict with IRS and number of active dimensions.
"""
del artifact_dir
logging.info("Generating training set.")
mus, ys = utils.generate_batch_factor_code(ground_truth_data,
representation_function,
num_train, random_state,
batch_size)
assert mus.shape[1] == num_train
ys_discrete = utils.make_discretizer(ys)
active_mus = _drop_constant_dims(mus)
if not active_mus.any():
irs_score = 0.0
else:
irs_score = scalable_disentanglement_score(ys_discrete.T, active_mus.T,
diff_quantile)["avg_score"]
score_dict = {}
score_dict["IRS"] = irs_score
score_dict["num_active_dims"] = np.sum(active_mus)
return score_dict
def compute_fairness(ground_truth_data,
representation_function,
random_state,
artifact_dir=None,
num_train=gin.REQUIRED,
num_test_points_per_class=gin.REQUIRED,
batch_size=16):
"""Computes unfairness scores.
We first compute either the mean or maximum total variation for a given
sensitive and target variable. Then, we either average or take the maximum
with respect to target and sensitive variable. For convenience, we compute and
save all combinations. The score used in Section 4 of the paper is here called
mean_fairness:mean_pred:mean_sens.
Args:
ground_truth_data: GroundTruthData to be sampled from.
representation_function: Function that takes observations as input and
outputs a dim_representation sized representation for each observation.
random_state: Numpy random state used for randomness.
artifact_dir: Optional path to directory where artifacts can be saved.
num_train: Number of points used for training.
num_test_points_per_class: Number of points used for testing.
batch_size: Batch size for sampling.
Returns:
Dictionary with scores.
"""
del artifact_dir
factor_counts = ground_truth_data.factors_num_values
num_factors = len(factor_counts)
scores = {}
# Training a predictive model.
mus_train, ys_train = utils.generate_batch_factor_code(
ground_truth_data, representation_function, num_train, random_state,
batch_size)
predictor_model_fn = utils.make_predictor_fn()
# For each factor train a single predictive model.
mean_fairness = np.zeros((num_factors, num_factors), dtype=np.float64)
max_fairness = np.zeros((num_factors, num_factors), dtype=np.float64)
for i in range(num_factors):
model = predictor_model_fn()
model.fit(np.transpose(mus_train), ys_train[i, :])
for j in range(num_factors):
if i == j:
continue
# Sample a random set of factors once.
original_factors = ground_truth_data.sample_factors(
num_test_points_per_class, random_state)
counts = np.zeros((factor_counts[i], factor_counts[j]),
dtype=np.int64)
for c in range(factor_counts[j]):
# Intervene on the sensitive attribute.
intervened_factors = np.copy(original_factors)
intervened_factors[:, j] = c
# Obtain the batched observations.
observations = ground_truth_data.sample_observations_from_factors(
intervened_factors, random_state)
representations = utils.obtain_representation(
observations, representation_function, batch_size)
# Get the predictions.
predictions = model.predict(np.transpose(representations))
# Update the counts.
counts[:, c] = np.bincount(predictions,
minlength=factor_counts[i])
mean_fairness[i, j], max_fairness[i,
j] = inter_group_fairness(counts)
# Report the scores.
scores.update(compute_scores_dict(mean_fairness, "mean_fairness"))
scores.update(compute_scores_dict(max_fairness, "max_fairness"))
return scores
def compute_reduced_downstream_task(ground_truth_data,
representation_function,
random_state,
artifact_dir=None,
num_factors_to_remove=gin.REQUIRED,
num_train=gin.REQUIRED,
num_test=gin.REQUIRED,
batch_size=16):
"""Computes loss of a reduced downstream task.
Measure the information leakage in each latent component after removing the
k ("factors_to_remove") most informative features for the prediction task.
Args:
ground_truth_data: GroundTruthData to be sampled from.
representation_function: Function that takes observations as input and
outputs a dim_representation sized representation for each observation.
random_state: Numpy random state used for randomness.
artifact_dir: Optional path to directory where artifacts can be saved.
num_factors_to_remove: Number of factors to remove from the latent
representation.
num_train: Number of points used for training.
num_test: Number of points used for testing.
batch_size: Batch size for sampling.
Returns:
Dictionary with scores.
"""
del artifact_dir
scores = {}
# Loop on different sizes of the training 'batch', as specified with gin.
for train_size in num_train:
size_string = str(train_size)
mus_train, ys_train = utils.generate_batch_factor_code(
ground_truth_data, representation_function, train_size, random_state,
batch_size)
mus_test, ys_test = utils.generate_batch_factor_code(
ground_truth_data, representation_function, num_test, random_state,
batch_size)
# Create variables for aggregated scores.
reduced_factor_train_scores = []
other_factors_train_scores = []
reduced_factor_test_scores = []
other_factors_test_scores = []
# Compute the reduced representation and test it for each factor of
# variation.
for factor_of_interest in range(ground_truth_data.num_factors):
# Copy the training data and eliminate the k most informative factors.
reduced_mus_train = mus_train.copy()
reduced_mus_test = mus_test.copy()
for _ in range(num_factors_to_remove):
reduced_mus_train, reduced_mus_test =\
compute_reduced_representation(reduced_mus_train, ys_train,
reduced_mus_test, ys_test,
factor_of_interest)
predictor_model = utils.make_predictor_fn()
train_acc, test_acc = compute_predictive_accuracy(
np.transpose(reduced_mus_train), ys_train,
np.transpose(reduced_mus_test), ys_test, predictor_model)
# Save scores for reduced factor.
scores[size_string +
":reduced_factor_{}:mean_train_accuracy_reduced_factor".format(
factor_of_interest)] = train_acc[factor_of_interest]
scores[size_string +
":reduced_factor_{}:mean_test_accuracy_reduced_factor".format(
factor_of_interest)] = test_acc[factor_of_interest]
reduced_factor_train_scores.append(train_acc[factor_of_interest])
reduced_factor_test_scores.append(test_acc[factor_of_interest])
# Save the scores (accuracies) in the score dictionary.
local_other_factors_train_scores = []
local_other_factors_test_scores = []
for i in range(len(train_acc)):
scores[size_string +
":reduced_factor_{}:mean_train_accuracy_factor_{}".format(
factor_of_interest, i)] = train_acc[i]
scores[size_string +
":reduced_factor_{}:mean_test_accuracy_factor_{}".format(
factor_of_interest, i)] = test_acc[i]
if i != factor_of_interest:
local_other_factors_train_scores.append(train_acc[i])
local_other_factors_test_scores.append(test_acc[i])
# Save mean score for non-reduced factors.
scores[size_string +
":reduced_factor_{}:mean_train_accuracy_non_reduced_factor".format(
factor_of_interest)] = np.mean(
local_other_factors_train_scores)
scores[size_string +
":reduced_factor_{}:mean_test_accuracy_non_reduced_factor".format(
factor_of_interest)] = np.mean(local_other_factors_test_scores)
other_factors_train_scores.append(
np.mean(local_other_factors_train_scores))
other_factors_test_scores.append(np.mean(local_other_factors_test_scores))
# Compute the aggregate scores.
scores[size_string + ":mean_train_accuracy_reduced_factor"] = np.mean(
reduced_factor_train_scores)
scores[size_string + ":mean_test_accuracy_reduced_factor"] = np.mean(
reduced_factor_test_scores)
scores[size_string + ":mean_train_accuracy_other_factors"] = np.mean(
other_factors_train_scores)
scores[size_string + ":mean_test_accuracy_other_factors"] = np.mean(
other_factors_test_scores)
return scores