mmd.py

'''
MMD functions for theano variables.
For each variant, returns (mmd2, objective_to_maximize).

If you just want to call them, do e.g.:

    Xth, Yth = T.matrices('X', 'Y')
    sigmath = T.scalar('sigma')
    fn = theano.function([Xth, Yth, sigmath],
                         rbf_mmd2_and_ratio(Xth, Yth, sigma=sigmath))

    mmd2, ratio = fn(X, Y, 1)
'''
from __future__ import division
import numpy as np
import theano.tensor as T
from theano.tensor import slinalg

_eps = 1e-8

################################################################################
### Quadratic-time MMD with Gaussian RBF kernel

def rbf_mmd2(X, Y, sigma=0, biased=True):
    gamma = 1 / (2 * sigma**2)

    XX = T.dot(X, X.T)
    XY = T.dot(X, Y.T)
    YY = T.dot(Y, Y.T)

    X_sqnorms = T.diagonal(XX)
    Y_sqnorms = T.diagonal(YY)

    K_XY = T.exp(-gamma * (
            -2 * XY + X_sqnorms[:, np.newaxis] + Y_sqnorms[np.newaxis, :]))
    K_XX = T.exp(-gamma * (
            -2 * XX + X_sqnorms[:, np.newaxis] + X_sqnorms[np.newaxis, :]))
    K_YY = T.exp(-gamma * (
            -2 * YY + Y_sqnorms[:, np.newaxis] + Y_sqnorms[np.newaxis, :]))

    if biased:
        mmd2 = K_XX.mean() + K_YY.mean() - 2 * K_XY.mean()
    else:
        m = K_XX.shape[0]
        n = K_YY.shape[0]

        mmd2 = ((K_XX.sum() - m) / (m * (m - 1))
              + (K_YY.sum() - n) / (n * (n - 1))
              - 2 * K_XY.mean())
    return mmd2, mmd2


def rbf_mmd2_and_ratio(X, Y, sigma=0, biased=True):
    gamma = 1 / (2 * sigma**2)

    XX = T.dot(X, X.T)
    XY = T.dot(X, Y.T)
    YY = T.dot(Y, Y.T)

    X_sqnorms = T.diagonal(XX)
    Y_sqnorms = T.diagonal(YY)

    K_XY = T.exp(-gamma * (
            -2 * XY + X_sqnorms[:, np.newaxis] + Y_sqnorms[np.newaxis, :]))
    K_XX = T.exp(-gamma * (
            -2 * XX + X_sqnorms[:, np.newaxis] + X_sqnorms[np.newaxis, :]))
    K_YY = T.exp(-gamma * (
            -2 * YY + Y_sqnorms[:, np.newaxis] + Y_sqnorms[np.newaxis, :]))

    return _mmd2_and_ratio(K_XX, K_XY, K_YY, unit_diagonal=True, biased=biased)


################################################################################
### Linear-time MMD with Gaussian RBF kernel

# Estimator and the idea of optimizing the ratio from:
#    Gretton, Sriperumbudur, Sejdinovic, Strathmann, and Pontil.
#    Optimal kernel choice for large-scale two-sample tests. NIPS 2012.

def rbf_mmd2_streaming(X, Y, sigma=0):
    # n = (T.smallest(X.shape[0], Y.shape[0]) // 2) * 2
    n = (X.shape[0] // 2) * 2
    gamma = 1 / (2 * sigma**2)
    rbf = lambda A, B: T.exp(-gamma * ((A - B) ** 2).sum(axis=1))
    mmd2 = (rbf(X[:n:2], X[1:n:2]) + rbf(Y[:n:2], Y[1:n:2])
          - rbf(X[:n:2], Y[1:n:2]) - rbf(X[1:n:2], Y[:n:2])).mean()
    return mmd2, mmd2


def rbf_mmd2_streaming_and_ratio(X, Y, sigma=0):
    # n = (T.smallest(X.shape[0], Y.shape[0]) // 2) * 2
    n = (X.shape[0] // 2) * 2
    gamma = 1 / (2 * sigma**2)
    rbf = lambda A, B: T.exp(-gamma * ((A - B) ** 2).sum(axis=1))
    h_bits = (rbf(X[:n:2], X[1:n:2]) + rbf(Y[:n:2], Y[1:n:2])
            - rbf(X[:n:2], Y[1:n:2]) - rbf(X[1:n:2], Y[:n:2]))

    mmd2 = h_bits.mean()

    # variance is 1/2 E_{v, v'} (h(v) - h(v'))^2
    # estimate with even, odd diffs
    m = (n // 2) * 2
    approx_var = 1/2 * ((h_bits[:m:2] - h_bits[1:m:2]) ** 2).mean()
    ratio = mmd2 / T.sqrt(T.largest(approx_var, _eps))
    return mmd2, ratio


################################################################################
### MMD with linear kernel

# Hotelling test statistic is from:
#    Jitkrittum, Szabo, Chwialkowski, and Gretton.
#    Interpretable Distribution Features with Maximum Testing Power.
#    NIPS 2015.

def linear_mmd2(X, Y, biased=True):
    if not biased:
        raise ValueError("Haven't implemented unbiased linear_mmd2 yet")
    X_bar = X.mean(axis=0)
    Y_bar = Y.mean(axis=0)
    Z_bar = X_bar - Y_bar
    mmd2 = Z_bar.dot(Z_bar)
    return mmd2, mmd2


def linear_mmd2_and_hotelling(X, Y, biased=True, reg=0):
    if not biased:
        raise ValueError("linear_mmd2_and_hotelling only works for biased est")

    n = X.shape[0]
    p = X.shape[1]
    Z = X - Y
    Z_bar = Z.mean(axis=0)
    mmd2 = Z_bar.dot(Z_bar)

    Z_cent = Z - Z_bar
    S = Z_cent.T.dot(Z_cent) / (n - 1)
    # z' inv(S) z = z' inv(L L') z = z' inv(L)' inv(L) z = ||inv(L) z||^2
    L = slinalg.cholesky(S + reg * T.eye(p))
    Linv_Z_bar = slinalg.solve_lower_triangular(L, Z_bar)
    lambda_ = n * Linv_Z_bar.dot(Linv_Z_bar)
    # happens on the CPU!
    return mmd2, lambda_


def linear_mmd2_and_ratio(X, Y, biased=True):
    # TODO: can definitely do this faster for a linear kernel...
    K_XX = T.dot(X, X.T)
    K_XY = T.dot(X, Y.T)
    K_YY = T.dot(Y, Y.T)
    return _mmd2_and_ratio(K_XX, K_XY, K_YY, unit_diagonal=False, biased=biased)


################################################################################
### Helper functions to compute variances based on kernel matrices

def _mmd2_and_ratio(K_XX, K_XY, K_YY, unit_diagonal=False, biased=False,
                    min_var_est=_eps):
    mmd2, var_est = _mmd2_and_variance(
        K_XX, K_XY, K_YY, unit_diagonal=unit_diagonal, biased=biased)
    ratio = mmd2 / T.sqrt(T.largest(var_est, min_var_est))
    return mmd2, ratio


def _mmd2_and_variance(K_XX, K_XY, K_YY, unit_diagonal=False, biased=False):
    m = K_XX.shape[0]  # Assumes X, Y are same shape

    ### Get the various sums of kernels that we'll use
    # Kts drop the diagonal, but we don't need to compute them explicitly
    if unit_diagonal:
        diag_X = diag_Y = 1
        sum_diag_X = sum_diag_Y = m
        sum_diag2_X = sum_diag2_Y = m
    else:
        diag_X = T.diagonal(K_XX)
        diag_Y = T.diagonal(K_YY)

        sum_diag_X = diag_X.sum()
        sum_diag_Y = diag_Y.sum()

        sum_diag2_X = diag_X.dot(diag_X)
        sum_diag2_Y = diag_Y.dot(diag_Y)

    Kt_XX_sums = K_XX.sum(axis=1) - diag_X
    Kt_YY_sums = K_YY.sum(axis=1) - diag_Y
    K_XY_sums_0 = K_XY.sum(axis=0)
    K_XY_sums_1 = K_XY.sum(axis=1)

    Kt_XX_sum = Kt_XX_sums.sum()
    Kt_YY_sum = Kt_YY_sums.sum()
    K_XY_sum = K_XY_sums_0.sum()

    # TODO: turn these into dot products?
    # should figure out if that's faster or not on GPU / with theano...
    Kt_XX_2_sum = (K_XX ** 2).sum() - sum_diag2_X
    Kt_YY_2_sum = (K_YY ** 2).sum() - sum_diag2_Y
    K_XY_2_sum  = (K_XY ** 2).sum()

    if biased:
        mmd2 = ((Kt_XX_sum + sum_diag_X) / (m * m)
              + (Kt_YY_sum + sum_diag_Y) / (m * m)
              - 2 * K_XY_sum / (m * m))
    else:
        mmd2 = (Kt_XX_sum / (m * (m-1))
              + Kt_YY_sum / (m * (m-1))
              - 2 * K_XY_sum / (m * m))

    var_est = (
          2 / (m**2 * (m-1)**2) * (
              2 * Kt_XX_sums.dot(Kt_XX_sums) - Kt_XX_2_sum
            + 2 * Kt_YY_sums.dot(Kt_YY_sums) - Kt_YY_2_sum)
        - (4*m-6) / (m**3 * (m-1)**3) * (Kt_XX_sum**2 + Kt_YY_sum**2)
        + 4*(m-2) / (m**3 * (m-1)**2) * (
              K_XY_sums_1.dot(K_XY_sums_1)
            + K_XY_sums_0.dot(K_XY_sums_0))
        - 4 * (m-3) / (m**3 * (m-1)**2) * K_XY_2_sum
        - (8*m - 12) / (m**5 * (m-1)) * K_XY_sum**2
        + 8 / (m**3 * (m-1)) * (
              1/m * (Kt_XX_sum + Kt_YY_sum) * K_XY_sum
            - Kt_XX_sums.dot(K_XY_sums_1)
            - Kt_YY_sums.dot(K_XY_sums_0))
    )

    return mmd2, var_est