Learning Representations for Independence Testing

## Abstract

Abstract The Hilbert-Schmidt Independence Criterion (HSIC) is a powerful tool for nonparametric detection of dependence between random variables. It crucially depends, however, on the selection of reasonable kernels; commonly-used choices like the Gaussian kernel,
or the kernel that yields the distance covariance, are sufficient only for amply sized samples from data distributions with relatively simple forms of dependence. We propose a scheme for selecting the kernels used in an HSIC-based independence test, based on maximizing an estimate of the asymptotic test power. We prove that maximizing this estimate indeed approximately maximizes the true power of the test, and demonstrate that our learned kernels can identify forms of structured dependence between random variables
in various experiments.

## 1 Introduction

Independence testing, the question of using paired samples to determine whether a random variable $X$ and another $Y$ are associated with one another
or if they are statistically independent,
is one of the most common tasks across scientific and data-based fields.
Traditional methods
make strong parametric assumptions,
for instance assuming that $X$ and $Y$ are jointly normal so that dependence is characterized by covariance,
and/or operate only in limited settings,
for instance the tabular setting of the celebrated $\chi^{2}$ test or Fisher’s exact test.
More recent widely-applicable tests have been developed for discrete distributions
based on the $\ell_{1}$ distance []
or multinomial log-likelihood [].
Using these methods on continuous data, however,
requires a “bucketing” process,
which does not work well for high-dimensional data.

As an alternative for more general kinds of data distributions,
[] introduced the Hilbert-Schmidt Independence Criterion (HSIC),
which measures the total cross-covariance
between feature functions of $X$ and $Y$,
supporting even *infinite-dimensional* features
in a reproducing kernel Hilbert space (RKHS).
This statistic can also be efficiently estimated from samples.
With appropriate choices of kernel [, see], HSIC tests can eventually detect *any* form of dependence between random variables, given enough samples.
In a separate line of work, [] and [] proposed *distance covariance* tests,
which measure the covariance of pairwise distances between $X$ values with distances between $Y$ values;
it was later shown that these classes of tests are in fact equivalent [].

With standard kernels, such as the Gaussian kernel or the distance kernel that recovers distance covariance, these tests will *eventually* be able to detect any fixed dependence with enough samples;
indeed, most applications of HSIC or the distance covariance use these “default” choices.
Their finite-sample behavior, however, can differ dramatically based on the choice of kernel.
Dependence which is easy to identify with the “right lens” might be extremely subtle with another choice of how to “look at” the data,
and require an enormous number of samples to achieve acceptable power.

[Figure 1](#S1.F1) shows the test power of HSIC with a simple kernel on a high-dimensional Gaussian mixture dataset.
Choosing the lengthscale of a Gaussian kernel with the common median heuristic (HSIC-M) works reasonably in dimension 4, but requires enormous amounts of data in dimension 20 or 40.
Optimizing the lengthscale for test power – HSIC-O with our scheme, or the previous approach HSIC-Agg [] –
does not substantially help,
because any Gaussian kernel is poorly suited to this task.
Our scheme, however, allows optimizing the parameters of kernels defined by a deep network (HSIC-D),
identifying a good kernel and achieving high test power with far fewer samples.

Figure: (a) Samples $X,Y\sim\mathbb{P}_{xy}$
Refer to caption: /html/2409.06890/assets/x1.png

Using a rich class of kernels parameterized by a deep neural network allows us to effectively identify dependence even between complex data types.
With appropriate deep kernel architectures,
we can maintain the consistency of the test,
while achieving far more powerful tests at a given dataset size.

Our criterion for which kernel to select is based on the test power,
whose asymptotic expression is dominated by the ratio of the HSIC to its standard deviation.
We show that we can estimate this quantity from finite samples in a reliable way,
such that optimizing the estimate approximately optimizes the asymptotic power of the test.
To the best of our knowledge, this is the first theoretical analysis of general kernel selection in independence testing.
We then construct a valid test via data splitting and permutation testing [, e.g.].
Our overall scheme builds on prior work in kernel two-sample testing [, , , ],
but the important independence testing setting has thus far been unaddressed to our knowledge.
We also demonstrate the effectiveness of our scheme in numerical experiments.

## 2 Kernel Independence Testing

###### Problem 1 (Independence testing) .

Let $Z=(X,Y)\sim\mathbb{P}_{xy}$ on a domain $\mathcal{X}\times\mathcal{Y}$, and $\mathbb{P}_{x}$ and $\mathbb{P}_{y}$ the corresponding marginal distributions on $\mathcal{X}$ and $\mathcal{Y}$.
We observe $m$ independent samples $\{(x_{1},y_{1}),\dots,(x_{m},y_{m})\}$ from $\mathbb{P}_{xy}$.
We conduct a null hypothesis significance test,
with the null hypothesis $\mathfrak{H}_{0}:\mathbb{P}_{xy}=\mathbb{P}_{x}\times\mathbb{P}_{y}$ (so $X\mathchoice{\mathrel{\hbox to0.0pt{$\displaystyle\perp$\hss}\mkern 2.0mu{\displaystyle\perp}}}{\mathrel{\hbox to0.0pt{$\textstyle\perp$\hss}\mkern 2.0mu{\textstyle\perp}}}{\mathrel{\hbox to0.0pt{$\scriptstyle\perp$\hss}\mkern 2.0mu{\scriptstyle\perp}}}{\mathrel{\hbox to0.0pt{$\scriptscriptstyle\perp$\hss}\mkern 2.0mu{\scriptscriptstyle\perp}}}Y$),
and the alternative $\mathfrak{H}_{1}:\mathbb{P}_{xy}\neq\mathbb{P}_{x}\times\mathbb{P}_{y}$ (i.e. $X\centernot{\mathchoice{\mathrel{\hbox to0.0pt{$\displaystyle\perp$\hss}\mkern 2.0mu{\displaystyle\perp}}}{\mathrel{\hbox to0.0pt{$\textstyle\perp$\hss}\mkern 2.0mu{\textstyle\perp}}}{\mathrel{\hbox to0.0pt{$\scriptstyle\perp$\hss}\mkern 2.0mu{\scriptstyle\perp}}}{\mathrel{\hbox to0.0pt{$\scriptscriptstyle\perp$\hss}\mkern 2.0mu{\scriptscriptstyle\perp}}}}Y$).

We wish to solve this problem without making strong (parametric) assumptions about the form of $\mathbb{P}_{xy}$, $\mathbb{P}_{x}$, or $\mathbb{P}_{y}$.
Most independence tests are based on estimating the “amount” of dependence between $X$ and $Y$,
or equivalently the discrepancy between
between $\mathbb{P}_{xy}$ and $\mathbb{P}_{x}\times\mathbb{P}_{y}$.
Given a nonnegative quantity which is zero if $X\mathchoice{\mathrel{\hbox to0.0pt{$\displaystyle\perp$\hss}\mkern 2.0mu{\displaystyle\perp}}}{\mathrel{\hbox to0.0pt{$\textstyle\perp$\hss}\mkern 2.0mu{\textstyle\perp}}}{\mathrel{\hbox to0.0pt{$\scriptstyle\perp$\hss}\mkern 2.0mu{\scriptstyle\perp}}}{\mathrel{\hbox to0.0pt{$\scriptscriptstyle\perp$\hss}\mkern 2.0mu{\scriptscriptstyle\perp}}}Y$,
we can reject $\mathfrak{H}_{0}$ if our estimate is large enough that we are confident the true value is positive.

One well-known such measurement of dependence is the (Shannon) mutual information.
This quantity, the Kullback-Liebler divergence of $\mathbb{P}_{xy}$ from $\mathbb{P}_{x}\times\mathbb{P}_{y}$,
is zero if and only if $X\mathchoice{\mathrel{\hbox to0.0pt{$\displaystyle\perp$\hss}\mkern 2.0mu{\displaystyle\perp}}}{\mathrel{\hbox to0.0pt{$\textstyle\perp$\hss}\mkern 2.0mu{\textstyle\perp}}}{\mathrel{\hbox to0.0pt{$\scriptstyle\perp$\hss}\mkern 2.0mu{\scriptstyle\perp}}}{\mathrel{\hbox to0.0pt{$\scriptscriptstyle\perp$\hss}\mkern 2.0mu{\scriptscriptstyle\perp}}}Y$.
Unfortunately, despite substantial work on estimating the mutual information
in ways based on binning, nearest-neighbors, kernel density estimation, and so on (see e.g. the broad list implemented by [])
or an array of approaches based on variational bounds estimated with deep networks [, e.g.],
the mutual information is fundamentally extremely difficult to estimate from samples [, , ].

The Hilbert-Schmidt Independence Criterion (HSIC, []),
is also zero if and only if $X$ and $Y$ are independent when an appropriate kernel is chosen.
Unlike the mutual information,
it is easy to estimate from samples.
We will discuss its relationship to mutual information in [Section 3](#S3).

###### Problem 1 (Independence testing) .

#### Hilbert-Schmidt Independence Criterion (HSIC).

To define HSIC, we will first briefly review positive-definite kernels.
A (real-valued) *kernel* is a function $k:\mathcal{X}\times\mathcal{X}\to\mathbb{R}$
which can be expressed as the inner product between *feature maps* $\phi:\mathcal{X}\to\mathcal{F}$,
$k(x,x^{\prime})=\langle\phi(x),\phi(x^{\prime})\rangle_{\mathcal{F}}$,
where $\mathcal{F}$ is any Hilbert space.
For every kernel function, there exists a unique space called the *reproducing kernel Hilbert space* (RKHS),
which consists of functions $f:\mathcal{X}\to\mathbb{R}$.
The key *reproducing property* of an RKHS $\mathcal{F}$ states that for any function $f\in\mathcal{F}$ and any point $x\in\mathcal{X}$,
we have $\langle f,\phi(x)\rangle_{\mathcal{F}}=f(x)$.

Suppose we have a kernel $k$ on $\mathcal{X}$ with RKHS $\mathcal{F}$ and feature map $\phi$,
as well as another kernel $l$ on $\mathcal{Y}$ with RKHS $\mathcal{G}$ and feature map $\psi$.
Let $\otimes$ denote the outer product.(^1^11Analogously to Euclidean outer products, $f\otimes g:\mathcal{G}\to\mathcal{F}$ is given by $[f\otimes g]g^{\prime}=f\langle g,g^{\prime}\rangle_{\mathcal{G}}$.)
The *cross-covariance operator* is

$$ $C_{xy}=\operatorname*{\mathbb{E}}_{xy}\Bigl{[}\bigl{(}\phi(x)-\operatorname*{\mathbb{E}}_{x}\phi(x)\bigr{)}\otimes\bigl{(}\psi(y)-\operatorname*{\mathbb{E}}_{y}\psi(y)\bigr{)}\Bigr{]};$ $$

for kernels with finite-dimensional feature maps,
this is exactly the standard (cross-)covariance matrix between the features of $X$ and those of $Y$.
Under mild integrability conditions on the kernel and the distributions,(^2^22We require $\operatorname*{\mathbb{E}}[\sqrt{k(x,x)}],\operatorname*{\mathbb{E}}[\sqrt{k(y,y)}]<\infty$; this is guaranteed for any distribution when $k$ is bounded.)
the reproducing property shows that
$\langle f,C_{xy}g\rangle=\operatorname{Cov}(f(X),g(Y))$ for all $f\in\mathcal{F},\,g\in\mathcal{G}$.
One definition of independence is whether there exist any correlated “test functions” $f$ and $g$.
Thus, for rich enough choices of kernel
– using universal $k$ and $l$ suffices, but is not necessary [] –
we have that $X\mathchoice{\mathrel{\hbox to0.0pt{$\displaystyle\perp$\hss}\mkern 2.0mu{\displaystyle\perp}}}{\mathrel{\hbox to0.0pt{$\textstyle\perp$\hss}\mkern 2.0mu{\textstyle\perp}}}{\mathrel{\hbox to0.0pt{$\scriptstyle\perp$\hss}\mkern 2.0mu{\scriptstyle\perp}}}{\mathrel{\hbox to0.0pt{$\scriptscriptstyle\perp$\hss}\mkern 2.0mu{\scriptscriptstyle\perp}}}Y$ if and only if the operator $C_{xy}=0$.
We can thus check whether the operator is zero,
and hence whether $X\mathchoice{\mathrel{\hbox to0.0pt{$\displaystyle\perp$\hss}\mkern 2.0mu{\displaystyle\perp}}}{\mathrel{\hbox to0.0pt{$\textstyle\perp$\hss}\mkern 2.0mu{\textstyle\perp}}}{\mathrel{\hbox to0.0pt{$\scriptstyle\perp$\hss}\mkern 2.0mu{\scriptstyle\perp}}}{\mathrel{\hbox to0.0pt{$\scriptscriptstyle\perp$\hss}\mkern 2.0mu{\scriptscriptstyle\perp}}}Y$,
by checking the squared Hilbert-Schmidt norm of $C_{xy}$,
$\operatorname{HSIC}(X,Y)=\lVert C_{xy}\rVert_{\rm HS}^{2}$.
With finite-dimensional features, this corresponds to the squared Frobenius norm of the feature cross-covariance matrix.

There are two similar estimators of HSIC in wide use.
One, “the biased estimator,”
is a $V$-statistic:

$$ $\widehat{\operatorname{HSIC}}_{\rm b}(X,Y)=\frac{1}{m^{2}}\operatorname{tr}(\mathbf{K}\mathbf{H}\mathbf{L}\mathbf{H})=\frac{1}{m^{2}}\sum_{i,j}^{m}k_{ij}l_{ij}+\frac{1}{m^{4}}\sum_{i,j,q,r}^{m}k_{ij}l_{qr}-\frac{2}{m^{3}}\sum_{i,j,q}k_{ij}l_{iq},$ (1) $$

where here
$\mathbf{K}$ is the $m\times m$ matrix with entries $k_{ij}=k(x_{i},x_{j})$,
$\mathbf{L}$ similarly has entries $l_{ij}=l(y_{i},y_{j})$,
and $\mathbf{H}$ is the “centering matrix” $\mathbf{I}_{m}-\frac{1}{m}\mathbf{1}_{m}\mathbf{1}_{m}^{\top}$.
This estimator has $\operatorname{\mathcal{O}}(1/m)$ bias,
but is consistent.

The other common estimator, “the unbiased estimator,”
is a $U$-statistic []:

$$ $\small\widehat{\operatorname{HSIC}}_{\rm u}(X,Y)=\frac{1}{m(m-3)}\left[\operatorname{tr}\left(\boldsymbol{\tilde{K}}\boldsymbol{\tilde{L}}\right)+\frac{\boldsymbol{1}^{\top}\boldsymbol{\tilde{K}}\boldsymbol{1}\boldsymbol{1}^{\top}\boldsymbol{\tilde{L}}\boldsymbol{1}}{(m-1)(m-2)}-\frac{2\boldsymbol{1}^{\top}\boldsymbol{\tilde{K}}\boldsymbol{\tilde{L}}\boldsymbol{1}}{m-2}\right],$ (2) $$

where $\boldsymbol{\tilde{K}}$ and $\boldsymbol{\tilde{L}}$ are
$m\times m$ matrices whose diagonal entries are zero but whose off-diagonal entries agree with those of $\mathbf{K}$ or $\mathbf{L}$.
It is unbiased,
$\operatorname*{\mathbb{E}}\widehat{\operatorname{HSIC}}_{\rm u}(X,Y)=\operatorname{HSIC}(X,Y)$,
and also consistent.

#### HSIC tests.

To construct a test, we use the asymptotic distribution of the unbiased HSIC estimator.

###### Proposition 2 ( [ 63 ] , Theorem 5) .

Under the alternative hypothesis $\mathfrak{H}_{1}$: $\mathbb{P}_{xy}\neq\mathbb{P}_{x}\times\mathbb{P}_{y}$, the unbiased estimator of HSIC is asymptotically normal:
with $m$ samples,

$$ $\displaystyle\sqrt{m}(\widehat{\operatorname{HSIC}}_{\rm u}-\operatorname{HSIC})\overset{\textbf{d}}{\to}\mathcal{N}(0,\sigma^{2}_{\mathfrak{H}_{1}}).$ (3) $$

The asymptotic variance of $\sqrt{m}\,\widehat{\operatorname{HSIC}}_{\rm u}$,
$\sigma^{2}_{\mathfrak{H}_{1}}$,
can be consistently estimated from $n$ samples(^3^33Typically $m=n$, but we might want to use a few ($n$) samples to roughly estimate the power of an $m$-sample test with $m\gg n$, as done in a different context by [, ].) as
$16\left(R-{{\rm HSIC}}^{2}\right)$,
where
$R=\frac{1}{n}\sum_{i=1}^{n}\left(\frac{(n-4)!}{(n-1)!}\sum_{(j,q,r)\in{\bf i}^{n}_{3}/\ \{i\}}h(i,j,q,r)\right)^{2}$
.
Here
${\bf i}^{\ell}_{n}\setminus\{i\}$ denotes the set of all $\ell$-tuples drawn without replacement from the set $\{1,\dots,n\}\setminus\{i\}$,
and $h(i,j,q,r)=\frac{1}{24}\sum_{(s,t,u,v)}^{(i,j,q,r)}k_{st}(l_{st}+l_{uv}-2l_{su})$
,
where the sum ranges over all $4!=24$ ways to assign the distinct indices $\{i,j,q,r\}$
to the four variables $(s,t,u,v)$ without replacement.

The behavior of the estimator is different under the null, $X\mathchoice{\mathrel{\hbox to0.0pt{$\displaystyle\perp$\hss}\mkern 2.0mu{\displaystyle\perp}}}{\mathrel{\hbox to0.0pt{$\textstyle\perp$\hss}\mkern 2.0mu{\textstyle\perp}}}{\mathrel{\hbox to0.0pt{$\scriptstyle\perp$\hss}\mkern 2.0mu{\scriptstyle\perp}}}{\mathrel{\hbox to0.0pt{$\scriptscriptstyle\perp$\hss}\mkern 2.0mu{\scriptscriptstyle\perp}}}Y$;
then $m\,\widehat{\operatorname{HSIC}}_{\rm u}$ (rather than $\sqrt{m}\widehat{\operatorname{HSIC}}_{\rm u}$) converges in distribution
to something with complex dependence on $\mathbb{P}_{x}$, $\mathbb{P}_{y}$, $k$, and $l$.

Thus, under $\mathfrak{H}_{0}$ when $\operatorname{HSIC}=0$,
the statistic $m\,\widehat{\operatorname{HSIC}}_{\rm u}$ has mean zero and standard deviation $\Theta(1)$.
Under $\mathfrak{H}_{1}$,
the mean of $m\operatorname{HSIC}=\Theta(\sqrt{m})\to\infty$,
with standard deviation $\Theta(\sqrt{m})$.
For large enough $m$,
these two cases are distinguishable.
To select a threshold at any particular $m$,
we use permutation testing
to estimate the $1-\alpha$ quantile of the null distribution:
we repeatedly randomly shuffle the data
to break any possible dependence between $X$ and $Y$,
and estimate the $1-\alpha$ quantile of the distribution of $m\widehat{\operatorname{HSIC}}_{\rm u}$ values computed on shuffled samples.
If we also include the original order of the data in the shuffled set,
this scheme is exactly valid for any number of permutations [, Theorem 2],
and yields a consistent test [].

###### Proposition 2 ( [ 63 ] , Theorem 5) .

#### Impact of kernel selection.

Two-sample testing (given samples from $\mathbb{P}$ and $\mathbb{Q}$, is $\mathbb{P}=\mathbb{Q}$?) is closely related to independence testing, where kernel-based methods similar to HSIC are in wide use.
Similarly to independence testing,
any “reasonable” choice of kernel gives a consistent test,
such that a casual user might think they can simply always use a Gaussian kernel with unit bandwidth and be done with it.
In practice, however, this scheme can perform extremely poorly;
if the data varies on a very different scale, it will take exorbitant quantities of data to achieve any reasonable test power.
Thus, users of Gaussian kernels in two-sample testing often use the *median heuristic*
to choose a kernel relevant to the data at hand,
choosing a bandwidth based on the median pairwise distance among data points [].
While this is a reasonable first guess for many data types,
there exist datasets where it can be dramatically better to instead select a bandwidth that optimizes a measure of test power [, ].
Beyond that, there exist many distributions where no Gaussian kernel performs well – for instance, many problems on natural image data or involving sparsity – but a learned deep kernel can yield very powerful tests [].

Although it has not been to our knowledge studied in the same way, the situation is analogous for independence testing.
The synthetic HDGM dataset (for “high-dimensional Gaussian mixture”; details in
[Section 6](#S6))
is a mixture of two Gaussians,
where two dimensions are dependent (but uncorrelated), while all other dimensions are independent of one another.
Splitting this $Z$ into $X$ and $Y$ such that one of the dependent dimensions is in each variable,
detecting this dependence becomes quite difficult when there are many “distractor” dimensions present.

Several independence tests are considered in [Figure 1](#S1.F1), including HSIC-M based on a median-heuristic Gaussian kernel, HSIC-O based on optimizing the lengthscale of a Gaussian with our scheme (introduced shortly), and HSIC-D which uses our scheme to learn a kernel parameterized by a deep network.
(Others are described in [Section 6](#S6)).
Kernel selection significantly influences the test power.
Optimizing the lengthscale of the Gaussian helps a bit,
but much more improvement comes from using a more complex deep kernel class
– an option made available by our optimization scheme.

## 3 Relationship to Mutual Information Bounds

#### Population values.

Mutual information, $\operatorname{I}(X;Y)$, is also frequently used to measure the dependence among two random variables.
As mentioned, estimation is fundamentally challenging.
Recent approaches focus on finding variational bounds on $\operatorname{I}(X;Y)$ which can themselves be easily estimated. As an example,
consider the following lower bound [, ],
named for its connection to noise-contrastive estimation []:

$$ $\operatorname{I_{NCE}^{\mathit{f,K}}}(X;Y)=\operatorname*{\mathbb{E}}_{\begin{subarray}{c}(x_{i},y_{i})\sim\mathbb{P}_{xy}^{K}\end{subarray}}\left[\operatorname{\hat{I}_{NCE}^{\mathit{f,K}}}\right]\leq\operatorname{I}(X;Y)\qquad\operatorname{\hat{I}_{NCE}^{\mathit{f,K}}}=\frac{1}{K}\sum_{i=1}^{K}\log\frac{e^{f(x_{i},y_{i})}}{\frac{1}{K}\sum_{j=1}^{K}e^{f(x_{i},y_{j})}}.$ $$

Typically we parameterize $f$ as a deep network,
and estimate $\operatorname{I_{NCE}^{\mathit{f,K}}}$ from samples with
$\operatorname{\hat{I}_{NCE}^{\mathit{f,K}}}$.
To choose $f$,
we maximize $\operatorname{\hat{I}_{NCE}^{\mathit{f,K}}}$ on minibatches,
attempting to find the tightest lower bound.

To understand the relationship to HSIC,
we first use the equivalence
to the maximum mean discrepancy (MMD),
a notion of distance between distributions,
in its integral probability metric form.

###### Proposition 3 ( [ 59 ] , Theorem 25) .

Let $k$ and $l$ be kernels on $\mathcal{X}$ and $\mathcal{Y}$,
and define a kernel on $\mathcal{X}\times\mathcal{Y}$ by
$h\bigl{(}(x,y),(x^{\prime},y^{\prime})\bigr{)}=k(x,x^{\prime})l(y,y^{\prime})$
with RKHS $\mathcal{H}$.
Then

$$ $\sqrt{\operatorname{HSIC}(X,Y)}=\operatorname{MMD}(\mathbb{P}_{xy},\mathbb{P}_{x}\times\mathbb{P}_{y})=\sup_{\begin{subarray}{c}f\in\mathcal{H}\\ \lVert f\rVert_{\mathcal{H}}\leq 1\end{subarray}}\operatorname*{\mathbb{E}}_{(X,Y)\sim\mathbb{P}_{xy}}[f(X,Y)]-\operatorname*{\mathbb{E}}_{\begin{subarray}{c}X\sim\mathbb{P}_{x}\\ Y^{\prime}\sim\mathbb{P}_{y}\end{subarray}}[f(X,Y^{\prime})].$ $$

Now suppose $\sup_{x\in\mathcal{X},y\in\mathcal{Y}}h((x,y),(x,y))\leq\nu^{2}$ for $\nu\geq 0$.
Then
$\lvert f(x,y)\rvert=\lvert\langle f,\phi_{h}(x,y)\rangle_{\mathcal{H}}\rvert\leq\nu\lVert f\rVert_{\mathcal{H}}$ by Cauchy-Schwarz,
and so $\lVert f\rVert_{\mathcal{H}}\geq\frac{1}{\nu}\sup_{x,y}\lvert f(x,y)\rvert=\frac{1}{\nu}\lVert f\rVert_{\infty}$. We thus have that

$$ $\sqrt{\operatorname{HSIC}(X,Y)}\leq\sup_{f:\lVert f\rVert_{\infty}\leq\nu}\operatorname*{\mathbb{E}}_{(X,Y)\sim\mathbb{P}_{xy}}[f(X,Y)]-\operatorname*{\mathbb{E}}_{\begin{subarray}{c}X\sim\mathbb{P}_{x}\\ Y^{\prime}\sim\mathbb{P}_{y}\end{subarray}}[f(X,Y^{\prime})]=2\nu\operatorname{TV}(\mathbb{P}_{xy},\mathbb{P}_{x}\times\mathbb{P}_{y}),$ (4) $$

where $\operatorname{TV}$ is the total variation distance between distributions [](^4^44[] define the TV as twice the more common definition, which we use here.)
Applying standard bounds relating the total variation to the KL divergence which defines mutual information, we obtain the following.

###### Proposition 4 .

In the setting of [Proposition 3](#Thmtheorem3),
suppose $\sup_{x\in\mathcal{X},y\in\mathcal{Y}}h((x,y),(x,y))\leq\nu^{2}$.
Then

$$ $\frac{1}{2\nu^{2}}\operatorname{HSIC}(X,Y)\leq\operatorname{I}(X;Y)\qquad\text{and}\qquad-\log\left(1-\frac{1}{4\nu^{2}}\operatorname{HSIC}(X,Y)\right)\leq\operatorname{I}(X;Y).$ $$

###### Proof.

The first bound follows from applying Pinsker’s inequality, which relates total variation to the KL divergence, to ([4](#S3.E4)).
The second instead applies the bound of [] [, also see].
∎

The second bound is tighter for large values of $\operatorname{I}(X;Y)$;
most importantly,
both bounds are monotonic in $\operatorname{HSIC}(X,Y)/\nu^{2}$.

We could thus consider a similar scheme to prior work:
choose kernels $k$, $l$ to maximize
a mutual information lower bound,
by maximizing $\operatorname{HSIC}(X,Y)/\nu^{2}$.
Indeed, maximizing HSIC has been used by many previous applications in other areas [, , , , e.g.].

###### Proposition 3 ( [ 59 ] , Theorem 25) .

###### Proposition 4 .

###### Proof.

#### Tests.

To construct an appropriate test
and avoid overfitting,
we choose a kernel on a training set,
then run a test on a separate test set.

We can also use the same scheme to construct an independence test based on $\operatorname{I_{NCE}^{\mathit{}}}$,
which to our knowledge has not yet been done in the literature.
Given a critic function $f$ chosen independently of a test dataset
– perhaps obtained by maximizing $\operatorname{\hat{I}_{NCE}^{\mathit{}}}$ on a training set, but not necessarily –
we can again construct a permutation test with exact level control based on the test statistic

$$ $\operatorname{\hat{I}_{NCE}^{\mathit{f,K}}}=\frac{1}{K}\sum_{i=1}^{K}f(x_{i},y_{i})-\frac{1}{K}\sum_{i=1}^{K}\log\left(\frac{1}{K}\sum_{j=1}^{K}e^{f(x_{i},y_{j})}\right).$ (5) $$

The second term of ([5](#S3.E5))
considers all pairs $x_{i},y_{j}$,
regardless of how they are aligned in the dataset.
Thus it shifts all permutations by the same amount, and is irrelevant to the actual test.
The same is true for many mutual information bounds, including UBA, TUBA, NWJ, and DV considered by [];
they choose $f$ differently,
but given an $f$, permutation tests are identical.

Now
consider a test based on
$\operatorname{MMD}(\mathbb{P}_{xy},\mathbb{P}_{x}\times\mathbb{P}_{y})$
using a kernel
of the form
$h((x,y),(x^{\prime},y^{\prime}))=g(x,y)g(x^{\prime},y^{\prime})$ for some real-valued function $g$.
(If $g(x,y)=g_{1}(x)g_{2}(y)$, this is an HSIC test.)
Because $\phi_{h}(x,y)=g(x,y)\in\mathbb{R}$ is a valid feature map,
every function in $\mathcal{H}$
is of the form $f(x,y)=\alpha g(x,y)$
with $\lVert\alpha g\rVert_{\mathcal{H}}=\lvert\alpha\rvert$.
By [Proposition 3](#Thmtheorem3),

$$ $\operatorname{MMD}(\mathbb{P}_{xy},\mathbb{P}_{x}\times\mathbb{P}_{y})^{2}=\left(\sup_{\lvert\alpha\rvert\leq 1}\alpha\left(\operatorname*{\mathbb{E}}g(X,Y)-\operatorname*{\mathbb{E}}g(X,Y^{\prime})\right)\right)^{2}=\bigl{(}\operatorname*{\mathbb{E}}g(X,Y)-\operatorname*{\mathbb{E}}g(X,Y^{\prime})\bigr{)}^{2}.$ $$

Using the plug-in estimator corresponding to the biased HSIC estimator ([1](#S2.E1))
would yield the test statistic

$$ $\widehat{\operatorname{HSIC}}_{\rm b}=\left(\frac{1}{m}\sum_{i=1}^{m}g(x_{i},y_{i})-\frac{1}{m^{2}}\sum_{i=1}^{m}\sum_{j=1}^{m}g(x_{i},y_{j})\right)^{2}.$ (6) $$

Comparing ([6](#S3.E6)) to ([5](#S3.E5)) with $K=m$,
we can see that the main term is identical between the two;
call this $T$.
The other term is again permutation-invariant;
it is the mean of $T$ over all possible permutations, $\bar{T}$.
Thus, a permutation test based on ([6](#S3.E6))
asks how far the value of $T$ for the true data is from $\bar{T}$,
while a permutation test based on ([5](#S3.E5))
is equivalent to asking how much the value of $T$ exceeds $\bar{T}$.
The only difference is that ([6](#S3.E6)) gives a two-sided test,
while ([5](#S3.E5)) gives a one-sided test.(^5^55[, Section 3.1 ] also found a relationship between HSIC and $\operatorname{I_{NCE}^{\mathit{}}}$ for categorical $Y$.)

Our usual test uses $\widehat{\operatorname{HSIC}}_{\rm u}$ of ([2](#S2.E2)) instead of $\widehat{\operatorname{HSIC}}_{\rm b}$, but the difference in estimators is typically small.
Thus, if we use deep kernels of the form $k(x,x^{\prime})=f(x)f(x^{\prime})$ and $l(y,y^{\prime})=g(y)g(y^{\prime})$,
the HSIC test(^6^66This version is closely related to a witness two-sample test [] used for independence.) is nearly equivalent to the NCE test with a separable critic function $(x,y)\mapsto f(x)g(y)$.
The NCE test chooses a critic by maximizing the NCE estimate;
while we could choose a (different) critic by maximizing the HSIC, we will now see there is a better approach.

## 4 Choosing an HSIC kernel

#### Approximate test power.

The standard quality measure of a null hypothesis test,
as long as it satisfies its design level of Type-I error,
is *test power*:
the probability that, when $X\centernot{\mathchoice{\mathrel{\hbox to0.0pt{$\displaystyle\perp$\hss}\mkern 2.0mu{\displaystyle\perp}}}{\mathrel{\hbox to0.0pt{$\textstyle\perp$\hss}\mkern 2.0mu{\textstyle\perp}}}{\mathrel{\hbox to0.0pt{$\scriptstyle\perp$\hss}\mkern 2.0mu{\scriptstyle\perp}}}{\mathrel{\hbox to0.0pt{$\scriptscriptstyle\perp$\hss}\mkern 2.0mu{\scriptscriptstyle\perp}}}}Y$, we correctly reject $\mathfrak{H}_{0}$.
[Proposition 2](#Thmtheorem2) implies,
using $\Phi$ for the standard normal CDF,
that

$$ $\displaystyle{\Pr}_{\mathfrak{H}_{1}}\left(m\,\widehat{\operatorname{HSIC}}_{\rm u}>r\right)$ $\displaystyle={\Pr}_{\mathfrak{H}_{1}}\left(\frac{\sqrt{m}}{\sigma_{\mathfrak{H}_{1}}}(\widehat{\operatorname{HSIC}}_{\rm u}-\operatorname{HSIC})>\frac{r}{\sqrt{m}\,\sigma_{\mathfrak{H}_{1}}}-\frac{\sqrt{m}\operatorname{HSIC}}{\sigma_{\mathfrak{H}_{1}}}\right)$ $\displaystyle\sim\Phi\left(\frac{\sqrt{m}\operatorname{HSIC}}{\sigma_{\mathfrak{H}_{1}}}-\frac{r}{\sqrt{m}\,\sigma_{\mathfrak{H}_{1}}}\right),$ (7) $$

where $a\sim b$ means $\lim_{m\to\infty}a/b=1$.
Letting $r$ be the threshold found by permutation testing,
for any given $\mathbb{P}_{xy}$, $k$, and $l$,
we know from [Proposition 2](#Thmtheorem2) that
$r$ converges to a constant
and $\operatorname{HSIC}$ and $\sigma_{\mathfrak{H}_{1}}$ are also constants.
As $m$ grows,
the test power is thus dominated by the first term.(^7^77This statement, while true in the limit $m\to\infty$, does not fully describe the setting for smaller $m$: for instance, a test which is only near the rejection threshold will have $r\approx m\,\operatorname{HSIC}$, in which case the two terms are of approximately equal size. As in [Footnote 3](#footnote3), though, we can estimate the power of an $m$-sample test with $n$ samples: here we use $n$ samples to approximately maximize the power of a test with $m\to\infty$ samples. [] estimate an analogue of ([7](#S4.E7)) rather than just ([8](#S4.E8)) in a related setting.)
Thus, maximizing the following *approximate test power*
maximizes the limiting power of the test as $m\to\infty$:

$$ $J(X,Y;k,l)=\frac{\operatorname{HSIC}(X,Y;k,l)}{\sigma_{\mathfrak{H}_{1}}(X,Y;k,l)}.$ (8) $$

In practice, we cannot evaluate $J$ in ([8](#S4.E8)),
since we do not know $\mathbb{P}_{xy}$.
Thus, we estimate $J$ using

$$ $\hat{J}_{\lambda}(X,Y;k,l)=\frac{\widehat{\operatorname{HSIC}}_{\rm u}(X,Y;k,l)}{\hat{\sigma}_{\mathfrak{H}_{1}}^{\lambda}(X,Y;k,l)},$ (9) $$

using a regularized variance estimator
$\hat{\sigma}_{\mathfrak{H}_{1}}^{\lambda}(X,Y;k)=\sqrt{{16}(R-\widehat{\operatorname{HSIC}}^{2}_{\rm u})+\lambda}$.
Here $R$ is exactly the estimate from [Proposition 2](#Thmtheorem2);
following [],
it can be computed more efficiently with
$R=\frac{((n-4)!)^{2}}{4n((n-1)!)^{2}}\lVert\boldsymbol{h}\rVert^{2}$,
where the vector $\boldsymbol{h}$ is

$$ $\displaystyle\boldsymbol{h}={}$ $\displaystyle(n-2)^{2}\left(\boldsymbol{\tilde{K}}\circ\boldsymbol{\tilde{L}}\right)\boldsymbol{1}-n(\boldsymbol{\tilde{K}}\boldsymbol{1})\circ(\boldsymbol{\tilde{L}}\boldsymbol{1})+(\boldsymbol{1}^{\top}\boldsymbol{\tilde{L}}\boldsymbol{1})\boldsymbol{\tilde{K}}\boldsymbol{1}+(\boldsymbol{1}^{\top}\boldsymbol{\tilde{K}}\boldsymbol{1})\boldsymbol{\tilde{L}}\boldsymbol{1}-(\boldsymbol{1}^{\top}\boldsymbol{\tilde{K}}\boldsymbol{\tilde{L}}\boldsymbol{1})\boldsymbol{1}$ $\displaystyle+(n-2)\left((\boldsymbol{1}^{\top}(\boldsymbol{\tilde{K}}\circ\boldsymbol{\tilde{L}})\boldsymbol{1})\boldsymbol{1}-\boldsymbol{\tilde{K}}\boldsymbol{\tilde{L}}\boldsymbol{1}-\boldsymbol{\tilde{L}}\boldsymbol{\tilde{K}}\boldsymbol{1}\right),$ $$

with $\circ$ denoting elementwise multiplication on matrices,
and $\boldsymbol{1}=(1,\dots,1)\in\mathbb{R}^{n}$.
Given sets of possible kernels,
we can approximately identify
the kernels $(k,l)$ yielding the asymptotically most powerful HSIC test
by maximizing ([9](#S4.E9)).
Typically, we run some variant of a gradient-based optimization algorithm
to maximize ([9](#S4.E9)) with respect to the parameters of $k$ and/or $l$.

#### Kernel architecture.

Since simple kernels have limited power to observe the dependence between $X$ and $Y$, we consider using deep kernels [], which have been successfully used in two-sample testing [, , ] and many other settings [, , , , e.g.].

Specifically, similarly to [], we use the following deep kernels for $X$ and $Y$:

$$ $\displaystyle k_{\omega}(x,x^{\prime})$ $\displaystyle=(1-\epsilon_{X})\,\kappa_{X}(f_{\omega}(x),f_{\omega}(x^{\prime}))+\epsilon_{X}\,q_{X}(x,x^{\prime})$ $\displaystyle l_{\gamma}(y,y^{\prime})$ $\displaystyle=(1-\epsilon_{Y})\,\kappa_{Y}(g_{\gamma}(y),g_{\gamma}(y^{\prime}))+\epsilon_{Y}\,q_{Y}(y,y^{\prime}).$ $$

Here $f_{\omega}$ and $g_{\gamma}$ are deep networks with parameters in $\omega$, $\gamma$,
which extract relevant features from $\mathcal{X}$ or $\mathcal{Y}$
to a feature space $\mathbb{R}^{D}$.
These features are then used inside a Gaussian kernel $\kappa$ on the space $\mathbb{R}^{D}$, to compute the baseline “similarity” between data points.
We then take a convex combination of that kernel with a Gaussian kernel $q$ on the input space;
the weight of this component is determined by a parameter $\epsilon\in(0,1)$.
Using $\epsilon>0$
provides a “backup” to the deep kernel,
perhaps giving some signal early in optimization when the deep kernel features are not yet useful,
and guaranteeing that the overall kernel is characteristic.
The lengthscale of $\kappa$ as well as the mixture parameter $\epsilon$ are included in the overall parameters, $\omega$ or $\gamma$,
and learned during the optimization process.

With $\epsilon=0$, scalar-output $f_{\omega}$ and $g_{\gamma}$, and $\kappa(a,b)=ab$,
we obtain the kernel used for the near-equivalence to mutual information-based tests in [Section 3](#S3).
Using non-scalar features and a more powerful $\kappa$ allows the kernel to take on more of the “work,”
and in most of the mentioned prior work has performed better than scalar features with linear $\kappa$.

#### Overall learning algorithm.

The overall procedure, written based on full-batch gradient ascent for simplicity, is shown in [Algorithm 1](#alg1).
In practice, we use AdamW [] instead, and draw minibatches in epochs;
experimental details are given in [Section C.1](#S3.SS1).

Let $E_{X}$ be the cost of computing an embedding $f_{\omega}(x)$, $E_{Y}$ be the cost of computing an embedding $f_{\gamma}(y)$,
and $L$ the cost of computing
$k_{\omega}(x,x^{\prime})$ and $l_{\gamma}(y,y^{\prime})$
given the relevant embeddings.
Each training iteration costs $\mathcal{O}\left(KE_{X}+KE_{Y}+K^{2}L\right)$, where $K$ is the minibatch size.
Typically for practical values of $K$,
$E_{X}+E_{Y}\gg KL$,
so the cost is “almost” linear in practice.(^8^88Equation ([9](#S4.E9)) could use block estimators [] or incomplete $U$-statistics [] to reduce $\operatorname{\mathcal{O}}(K^{2}L)$ to $\operatorname{\mathcal{O}}(K^{\beta}L)$ for any $\beta\leq 2$, at the cost of increased variance [, see].)

Figure: Algorithm 1 Testing with a learned deep-HSIC

## 5 Theoretical Analysis

Does optimizing the estimator $\hat{J}_{\lambda}$ of ([9](#S4.E9))
in fact approximately optimize the asymptotic test power $J$ in ([8](#S4.E8))?
[Theorem 5](#Thmtheorem5) shows that, if the training set size is reasonably large and we can optimize the estimates successfully, we can learn a kernel that generalizes nearly optimally and is hence powerful for large $m$, not just overfitting to the training set (and finding a bad test).

###### Theorem 5 .

Under [Assumptions (A)](#S3.EGx6), [(B)](#S3.EGx7) and [(C)](#S1.Ex19),
let $\Theta\subseteq\Omega\times\Gamma$ be a set of kernel parameters $\theta\in\Theta$ for which $\sigma_{\theta}^{2}\geq s^{2}$, and $n$ be the training set size; take $\lambda=\Theta\left(n^{-1/3}\right)$.
Then

$$ $\displaystyle\sup_{\theta:=(\omega,\gamma)\in\Theta}\left|\hat{J}_{\lambda}(X,Y;k_{\omega},l_{\gamma})-J_{\lambda}(X,Y;k_{\omega},l_{\gamma})\right|=\tilde{\operatorname{\mathcal{O}}}_{p}\left(\frac{1}{s^{2}n^{1/3}}\left[\frac{1}{s}+L_{k}+L_{l}+\sqrt{D_{\Omega}}+\sqrt{D_{\Gamma}}\right]\right).$ $$

If there are unique best kernels $(k_{\omega^{*}},l_{\gamma^{*}})$ maximizing $J_{\lambda}(X,Y;k_{\omega},l_{\gamma})$, then the maximizer of $\hat{J}_{\lambda}(X,Y;k_{\omega},l_{\gamma})$ converges in probability to $(k_{\omega^{*}},l_{\gamma^{*}})$
as $n\to\infty$.

[Section A](#S1a) states and proves a nonasymptotic version of this result;
the assumptions and proof techniques
are based on those of [].

###### Theorem 5 .

## 6 Experiments

Baselines.
We compare our method with various baselines, all based on permutation testing:

- •
HSIC-D: HSIC using deep kernels on each space $\mathcal{X}$ and $\mathcal{Y}$; simultaneously trained via [Section 4](#S4).
- •
HSIC-Dx: HSIC using a tied deep kernel, i.e. $k_{\omega}=l_{\gamma}$, and trained via [Section 4](#S4).
- •
HSIC-O: HSIC using Gaussian kernels, with each bandwidth parameter optimized via [Section 4](#S4).
- •
HSIC-M: HSIC using Gaussian kernels, with bandwidth selected via the median heuristic.
- •
HSIC-Agg []: aggregating Gaussian kernels, with their default settings.
- •
MMD-D: The method of [] applied to $\mathbb{P}_{xy}$ vs $\mathbb{P}_{x}\times\mathbb{P}_{y}$,(^9^99Specifically, we compare the given samples to a single shuffling of the given samples.) with a kernel on $\mathcal{X}\times\mathcal{Y}$.
- •
C2ST-S [] / C2ST-L []: Sign/logit-based classifier two-sample test for $\mathbb{P}_{xy}$ vs $\mathbb{P}_{x}\times\mathbb{P}_{y}$.
- •
InfoNCE [, ]: the statistic $\operatorname{\hat{I}_{NCE}^{\mathit{}}}$ as in [Section 3](#S3).
- •
NWJ
[, ]: another mutual information bound statistic $\hat{I}_{\text{NWJ}}$.

#### Datasets.

We consider three informative synthetic datasets, where the true answers are known.

Figure: (a) HDGM-10
Refer to caption: /html/2409.06890/assets/x6.png

$\bullet\leavevmode\nobreak\$High-dimensional Gaussian mixture.
The distribution HDGM-$d$ has $d$ total dimensions (divided between $X$ and $Y$), but has dependence only between two of them:

$$ $\displaystyle[X_{1},Y_{\left\lfloor{d/2}\right\rfloor},\dots,X_{\left\lceil{d/2}\right\rceil},Y_{1}]\sim\sum_{i=1}^{2}\frac{1}{2}\mathcal{N}\left(\textbf{0}_{d},\begin{bmatrix}1&0.5(-1)^{i}&\textbf{0}^{T}_{d-2}\\ 0.5(-1)^{i}&1&\textbf{0}^{T}_{d-2}\\ \textbf{0}_{d-2}&\textbf{0}_{d-2}&I_{d-2}\end{bmatrix}\right),$ $$

where the odd dimensions are taken to be from $\mathbb{P}_{x}$ and even dimensions to be from $\mathbb{P}_{y}$. Moreover, for $d\geq 4$ the dependent variables $X_{1}$ and $Y_{\lfloor d/2\rfloor}$ are at different dimensions. We perform independence tests at dimensions 4, 8, 10, 20, 30, 40, and 50.

$\bullet\leavevmode\nobreak\$Sinusoid []. We sample from sinusoidally dependent data with distribution $\mathbb{P}_{xy}\propto 1+\sin(\ell x)\sin(\ell y)$ on support $\mathcal{X}\times\mathcal{Y}=[-\pi,\pi]^{2}$. Higher frequencies $\ell$ produce subtler departures from the uniform distribution, resulting in a harder independence problem; we use $\ell=4$. A visualization of this density is given in [Figure 9](#S3.F9).

$\bullet\leavevmode\nobreak\$RatInABox [].
RatInABox simulates hippocampal cells of a rat in motion. In particular, we test for dependence between firing rates of grid cells and the rat’s head direction. Grid cells respond near points in a grid covering the environment surface, and should be subtly connected to head direction because of the geometry of the “box” ([Figure 9](#S3.F9)). We consider 8 grid cells, and simulate motion for 100 000 seconds, taking a measurement every 5 seconds as our dataset.

Power versus test size.
We first compare how well methods identify dependency with a large available training set,
by comparing the rate at which the learned tests achieve perfect power (1.0) as the test set size $m$ increases.
Results on the HDGM problems are given in [Figures 1](#S1.F1) and [5](#S2.F5) for $m$ in {100, 200, 500, 1000, 2000}. We use a training size of 10,000 for HDGM $\leq$ 30 and 100,000 for HDGM $>$ 30, and 2000 validation samples for all dimensions. Overall, HSIC-D outperforms baselines, and is able to reach perfect power using smaller $m$. Both InfoNCE and C2ST-L also achieve high power at lower dimensions, but are unable to observe dependence for HDGM-$d$ with $d\geq 40$. Similarly, optimized bandwidth HSIC-O and median heuristic HSIC-M perform moderately well at lower dimensions, but at dimension of 15 and above fail to learn anything meaningful and exhibit power levels approximately equal to the type-I error rate.

Sinusoid results are shown in [Figure 2](#S6.F2) (d). We train on 5,000 samples and use 1,000 for validation. Once again, HSIC-D outperforms other baselines and can achieve perfect power using the smallest number of samples. In general, all methods apart from the median heuristic are able to detect dependency given enough test samples.

RatInABox results are shown in [Figure 2](#S6.F2) (c). We use a training size of 4,000 samples with no validation. Again, HSIC-D exhibits higher empirical power than baselines.
Notably, HSIC-Agg underperforms HSIC-O in these settings, despite considering a similar class of kernels.
On these problems it seems that the aggregation process is worse than data splitting;
more relevantly, data splitting and optimization allows for far more complex (deep) kernels
than would be reasonable with a finite set of available kernels in HSIC-Agg.

Power versus dimension.
We demonstrate the effectiveness of our method at various dimensions $d/2$ by examining the empirical test power at HDGM-$d$ for $d\in\{4,8,10,20,30,40,50\}$ with fixed test sizes $m$. We use the same training splits as before. Results are shown in [Figure 3](#S6.F3). HSIC-D still exhibits the highest test power across all dimensions. When a small number of test samples are used (i.e., $m=100$) the performance of HSIC-D slightly degrades with increasing dimension, while at larger test sample sizes it consistently has near-perfect power. C2ST-L and InfoNCE demonstrate similar behavior but with worse power levels. HSIC-O and HSIC-M show failure at high dimensions. This is most noticeable at $m=2000$, where the power levels drastically curve downwards with increasing dimension.

Figure: (a) $m=100$
Refer to caption: /html/2409.06890/assets/x11.png

Power versus dataset size.
A drawback to kernel selection via optimization is that we must hold out a split of the data for training. In contrast, HSIC-M and HSIC-Agg are able to utilize the entirety of the dataset for their test. To examine this trade-off, we consider consistent data splitting at datasets of varying sizes. For the HDGM and Sinusoid problems we use a 7:2:1 train-val-test split, and for RatInABox we use a 3:2 train-test split. Conversely, HSIC-M and HSIC-Agg use the *entire* dataset for testing.
Results are shown in [Figure 4](#S6.F4). Notice that HSIC-M is able to take advantage of the additional test samples and achieves better power than some optimization-based methods at smaller datasets; however, as the size increases, optimization-based tests catch up and excel over heuristic methods. These results suggest that when the total number of samples is limited or when the dependence is straightforward, it may be advantageous to use heuristic-based tests. Conversely, if a dataset is large enough and the dependence is obscure, deep learning-based tests can be far more powerful.

Optimizing J versus HSIC. Rather than optimizing objective $J$, we can choose to maximize the test statistic HSIC; however, doing so makes no guarantees on the test power. [Figure 6](#S2.F6) compares these objectives and demonstrates that optimizing the approximate test power $J$ is preferred over optimizing the test statistic HSIC.

Figure: (a) HDGM-10
Refer to caption: /html/2409.06890/assets/x16.png

## 7 Conclusion and Future Work

Kernel independence testing aims to see if two paired random variables are statistically independent given a pre-defined kernel, which is an important research line in both statistics and machine learning.
As demonstrated in [Figures 1](#S1.F1) and [2](#S2.SS0.SSS0.Px3),
simple kernels as are overwhelmingly used for independence testing
can be resoundingly outperformed by learned kernels tailored to a particular problem.
We also, to our knowledge, are the first to propose testing based on variational mutual information estimation,
which does not work quite as well as HSIC-D in our experiments but is much better than simple fixed tests.
We further prove theoretically that the kernels learned HSIC-D result in powerful tests, and confirm this in experiments.

In the future, it is appealing to extend deep kernel learning schemes to conditional independence testing and apply this to causal discovery. Meanwhile, the data-splitting procedure used in this paper has to sacrifice some test power because we reduce the sample size for the final independence testing.
Thus, how to make this learning scheme more data-efficient,
as in HSIC-Agg [] or the scheme of [];
it is currently unclear, however, how to generalize this to complex data-dependent kernels.

## Acknowledgments

The authors would like to thank Roman Pogodin for productive discussions.
This work was enabled in part by support provided by the Natural Sciences and Engineering Research Council of Canada,
the Canada CIFAR AI Chairs program,
Calcul Québec,
the BC DRI Group,
and the Digital Research Alliance of Canada. This work was also supported by the Australian Research Council (ARC) with grant numbers DE240101089 and DP230101540.

[sorting=nyt]

## A Uniform Convergence Analysis

### A.1 Preliminaries

We start by defining the notation used in the proofs. Let kernels $k_{\omega}$ and $l_{\gamma}$ be parameterized by some $\omega$ and $\gamma$, and samples $(X_{i},Y_{i})\sim\mathbb{P}_{xy}$ drawn i.i.d. from the joint distribution. We denote the $n\times n$ gram matrices of $k_{\omega}$ and $l_{\gamma}$ by $K^{(\omega)}$ and $L^{(\gamma)}$ respectively. We will often omit the kernel parameters $\omega$ and $\gamma$ when it is clear from the context.

Let $\eta$ be the HSIC test statistic and $\hat{\eta}$ to be its U-statistic estimator given by

$$ $\hat{\eta}=\frac{1}{(n)_{4}}\sum_{(i,j,q,r)\in\mathbf{i}^{n}_{4}}H_{ijqr},$ $$

where $(n)_{k}=n!/(n-k)!$ is the Pochhammer symbol and $\mathbf{i}^{n}_{4}$ is all possible 4-tuples drawn without replacement from 1 to n. $H$ is the kernel gram matrix of the U-statistic defined by

$$ $H_{ijqr}=\frac{1}{4!}\sum^{(i,j,q,r)}_{(a,b,c,d)}K_{ab}\left(L_{ab}+L_{cd}-2L_{ac}\right),$ $$

where sum represents all $4!$ combinations of tuples $(a,b,c,d)$ that can be selected without replacement from $(i,j,q,r)$.

### A.2 Main Results

Our main uniform convergence results require the following assumptions.

- (A)
The set of kernel parameters $\Omega$ lies in a Banach space of dimension $D_{\Omega}$, and the set of kernel parameters $\Gamma$ lies in a Banach space of dimension $D_{\Gamma}$. Furthermore, each parameter space is bounded by $R_{\Omega}$ and $R_{\Gamma}$ respectively, i.e.,
$\displaystyle\Omega\subseteq\{\omega\,\,|\,\,\|\omega\|\leq R_{\Omega}\},$
$\displaystyle\Gamma\subseteq\{\gamma\,\,|\,\,\|\gamma\|\leq R_{\Gamma}\}.$
- (B)
The kernels $k_{\omega}$ and $l_{v}$ are uniformly bounded:
$\displaystyle\sup_{\omega\in\Omega}\sup_{x\in\mathcal{X}}k_{\omega}(x,x)\leq\nu_{k},$
$\displaystyle\sup_{\gamma\in\Gamma}\sup_{x\in\mathcal{Y}}l_{\gamma}(y,y)\leq\nu_{l}.$
For the kernels we use in practice, $\nu_{k}=\nu_{l}=1$.
- (C)
Both kernels $k$ and $l$ are Lipschitz with respect to the parameter space: for all $x,x^{\prime}\in\mathcal{X}$ and $\omega,\omega^{\prime}\in\Omega$
$|k_{\omega}(x,x^{\prime})-k_{\omega^{\prime}}(x,x^{\prime})|\leq L_{k}\|\omega-\omega^{\prime}\|,$
and for all $y,y^{\prime}\in\mathcal{Y}$ and $\gamma,\gamma^{\prime}\in\Gamma$
$|l_{\gamma}(y,y^{\prime})-l_{\gamma^{\prime}}(y,y^{\prime})|\leq L_{l}\|\gamma-\gamma^{\prime}\|.$

###### Theorem 6 .

Under [Assumptions (A)](#S3.EGx6), [(B)](#S3.EGx7) and [(C)](#S1.Ex19), let $\Theta\subseteq\Omega\times\Gamma$ be the set of kernel parameters $\theta\in\Theta$ for which $\sigma_{\theta}^{2}\geq s^{2}$, and take $\lambda=n^{-1/3}$. Assume $\nu_{k},\nu_{l}\geq 1$. Then, with probability at least $1-\delta$,

$$ $\displaystyle\sup_{\theta\in\Theta}\left|\frac{\hat{\eta}_{\theta}}{\hat{\sigma}_{\theta,\lambda}}-\frac{\eta_{\theta}}{\sigma_{\theta}}\right|\leq\,$ $\displaystyle\frac{2\nu_{k}\nu_{l}}{s^{2}n^{1/3}}\Bigg{[}\frac{1}{s}+\frac{9216\nu^{2}_{k}\nu^{2}_{l}}{\sqrt{n}}$ $\displaystyle+\left(12288\nu^{2}_{k}\nu^{2}_{l}+\frac{8s}{n^{1/6}}\right)\left(\frac{L_{k}}{\nu_{k}}+\frac{L_{l}}{\nu_{l}}+\sqrt{2\log\frac{4}{\delta}+2D_{\Omega}\log(4R_{\Omega}\sqrt{n})+2D_{\Gamma}\log(4R_{\Gamma}\sqrt{n})}\right)\Bigg{]},$ $$

and thus, treating $\nu_{k},\nu_{l}$ as constants,

$$ $\displaystyle\sup_{\theta\in\Theta}\left|\frac{\hat{\eta}_{\theta}}{\hat{\sigma}_{\theta,\lambda}}-\frac{\eta_{\theta}}{\sigma_{\theta}}\right|=\tilde{\operatorname{\mathcal{O}}}_{P}\left(\frac{1}{s^{2}n^{1/3}}\left[\frac{1}{s}+L_{k}+L_{l}+\sqrt{D_{\Omega}}+\sqrt{D_{\Gamma}}\right]\right).$ $$

###### Proof.

Let $\hat{\sigma}_{\theta,\lambda}^{2}:=\hat{\sigma}_{\theta}^{2}+\lambda$ be our regularized variance estimator from which we can assume is positive. We start by decomposing

$$ $\displaystyle\sup_{\theta\in\Theta}\left|\frac{\hat{\eta}_{\theta}}{\hat{\sigma}_{\theta,\lambda}}-\frac{\eta_{\theta}}{\sigma_{\theta}}\right|$ $\displaystyle\leq\sup_{\theta\in\Theta}\left|\frac{\hat{\eta}_{\theta}}{\hat{\sigma}_{\theta,\lambda}}-\frac{\hat{\eta}_{\theta}}{\sigma_{\theta,\lambda}}\right|+\sup_{\theta\in\Theta}\left|\frac{\hat{\eta}_{\theta}}{\sigma_{\theta,\lambda}}-\frac{\hat{\eta}_{\theta}}{\sigma_{\theta}}\right|+\sup_{\theta\in\Theta}\left|\frac{\hat{\eta}_{\theta}}{\sigma_{\theta}}-\frac{\eta_{\theta}}{\sigma_{\theta}}\right|$ $\displaystyle=\sup_{\theta}\frac{|\hat{\eta}_{\theta}|}{\hat{\sigma}_{\theta,\lambda}\cdot\sigma_{\theta,\lambda}}\frac{|\hat{\sigma}_{\theta,\lambda}^{2}-\sigma_{\theta,\lambda}^{2}|}{\hat{\sigma}_{\theta,\lambda}+\sigma_{\theta,\lambda}}+\sup_{\theta}\frac{|\hat{\eta}_{\theta}|}{\sigma_{\theta,\lambda}\cdot\sigma_{\theta}}\frac{|\sigma_{\theta,\lambda}^{2}-\sigma_{\theta}^{2}|}{\sigma_{\theta,\lambda}+\sigma_{\theta}}+\sup_{\theta}\frac{1}{\sigma_{\theta}}|\hat{\eta}_{\theta}-\eta_{\theta}|$ $\displaystyle\leq\frac{4\nu_{k}\nu_{l}}{s\sqrt{\lambda}(s+\sqrt{\lambda})}\sup_{\theta}|\hat{\sigma}_{\theta}^{2}-\sigma_{\theta}^{2}|+\frac{4\nu_{k}\nu_{l}\lambda}{s\sqrt{s^{2}+\lambda}(s+\sqrt{s^{2}+\lambda})}+\frac{1}{s}\sup_{\theta}|\hat{\eta}_{\theta}-\eta_{\theta}|$ $\displaystyle\leq\frac{4\nu_{k}\nu_{l}}{s^{2}\sqrt{\lambda}}\sup_{\theta}|\hat{\sigma}_{\theta}^{2}-\sigma_{\theta}^{2}|+\frac{1}{s}\sup_{\theta}|\hat{\eta}_{\theta}-\eta_{\theta}|+\frac{2\nu_{k}\nu_{l}\lambda}{s^{3}}.$ $$

[Proposition 7](#Thmtheorem7) and [Proposition 8](#Thmtheorem8) show the uniform convergence of $\hat{\eta}_{\theta}$ and $\hat{\sigma}_{\theta}$, from which we get that with probability at least $1-\delta$, the error is at most

$$ $\displaystyle\sup_{\theta\in\Theta}\left|\frac{\hat{\eta}_{\theta}}{\hat{\sigma}_{\theta,\lambda}}-\frac{\eta_{\theta}}{\sigma_{\theta}}\right|\leq\,$ $\displaystyle\frac{2\nu_{k}\nu_{l}\lambda}{s^{3}}+\frac{18432\nu^{3}_{k}\nu^{3}_{l}}{s^{2}n\sqrt{\lambda}}+\left[\frac{8192\nu^{3}_{k}\nu^{3}_{l}}{s^{2}\sqrt{\lambda n}}+\frac{8\nu_{k}\nu_{l}}{s\sqrt{n}}\right]\left(\frac{L_{k}}{\nu_{k}}+\frac{L_{l}}{\nu_{l}}\right)$ $\displaystyle+\left[\frac{24576\nu^{3}_{k}\nu^{3}_{l}}{s^{2}\sqrt{\lambda n}}+\frac{16\nu_{k}\nu_{l}}{s\sqrt{n}}\right]\sqrt{2\log\frac{4}{\delta}+2D_{\Omega}\log(4R_{\Omega}\sqrt{n})+2D_{\Gamma}\log(4R_{\Gamma}\sqrt{n})}.$ $$

Taking $\lambda=n^{-1/3}$ gives

$$ $\displaystyle\sup_{\theta\in\Theta}\left|\frac{\hat{\eta}_{\theta}}{\hat{\sigma}_{\theta,\lambda}}-\frac{\eta_{\theta}}{\sigma_{\theta}}\right|\leq\,$ $\displaystyle\frac{2\nu_{k}\nu_{l}}{s^{3}n^{1/3}}+\frac{18432\nu^{3}_{k}\nu^{3}_{l}}{s^{2}n^{5/6}}+\left[\frac{8192\nu^{3}_{k}\nu^{3}_{l}}{s^{2}n^{1/3}}+\frac{8\nu_{k}\nu_{l}}{s\sqrt{n}}\right]\left(\frac{L_{k}}{\nu_{k}}+\frac{L_{l}}{\nu_{l}}\right)$ $\displaystyle+\left[\frac{24576\nu^{3}_{k}\nu^{3}_{l}}{s^{2}n^{1/3}}+\frac{16\nu_{k}\nu_{l}}{s\sqrt{n}}\right]\sqrt{2\log\frac{4}{\delta}+2D_{\Omega}\log(4R_{\Omega}\sqrt{n})+2D_{\Gamma}\log(4R_{\Gamma}\sqrt{n})}.$ $$

Using $\nu_{k},\nu_{l}\geq 1$ we can slightly simplify our bound to

$$ $\displaystyle\sup_{\theta\in\Theta}\left|\frac{\hat{\eta}_{\theta}}{\hat{\sigma}_{\theta,\lambda}}-\frac{\eta_{\theta}}{\sigma_{\theta}}\right|\leq\,$ $\displaystyle\left[\frac{24576\nu^{3}_{k}\nu^{3}_{l}}{s^{2}n^{1/3}}+\frac{16\nu_{k}\nu_{l}}{s\sqrt{n}}\right]\left(\frac{L_{k}}{\nu_{k}}+\frac{L_{l}}{\nu_{l}}+\sqrt{2\log\frac{4}{\delta}+2D_{\Omega}\log(4R_{\Omega}\sqrt{n})+2D_{\Gamma}\log(4R_{\Gamma}\sqrt{n})}\right)$ $\displaystyle+\frac{2\nu_{k}\nu_{l}}{s^{3}n^{1/3}}+\frac{18432\nu^{3}_{k}\nu^{3}_{l}}{s^{2}n^{5/6}}$ $$

∎

###### Theorem 6 .

###### Proof.

### A.3 Uniform convergence results

This subsection pertains to uniform convergence results of $\hat{\eta}_{\theta}$ and $\hat{\sigma}_{\theta}$, and are are used in the proof of [Theorem 6](#Thmtheorem6).

###### Proposition 7 .

Under assumptions [(A)](#S3.EGx6) to [(C)](#S1.Ex19), we have that with probability at least $1-\delta$,

$$ $\sup_{\theta\in\Theta}|\hat{\eta}_{\theta}-\eta_{\theta}|\leq\frac{8\nu_{k}\nu_{l}}{\sqrt{n}}\left(\frac{L_{k}}{\nu_{k}}+\frac{L_{l}}{\nu_{l}}+2\sqrt{2\log\frac{2}{\delta}+2D_{\Omega}\log(4R_{\Omega}\sqrt{n})+2D_{\Gamma}\log(4R_{\Gamma}\sqrt{n})}\right).$ $$

###### Proof.

We use $\epsilon$-net arguments on both spaces $\Omega$ and $\Gamma$. Let $\{\omega_{i}\}_{i=1}^{T_{\Omega}}$ be arbitrarily placed centers with radius $\rho_{\Omega}$ such that any point $\omega\in\Omega$ satisfies $\min\|\omega-\omega_{i}\|\leq\rho_{\Omega}$. Similarly, let $\{\gamma_{i}\}_{i=1}^{T_{\Gamma}}$ be centers with radius $\rho_{\Gamma}$ satisfying $\min\|\gamma-\gamma_{i}\|\leq\rho_{\Gamma}$ for any $\gamma\in\Gamma$. Assumption [(B)](#S3.EGx7) ensures this is possible with at most $T_{\Omega}=(4R_{\Omega}/\rho_{\Omega})^{D_{\Omega}}$ and $T_{\Gamma}=(4R_{\Gamma}/\rho_{\Gamma})^{D_{\Gamma}}$ points respectively [, Proposition 5].

We can decompose the convergence bound into simpler components and tackle each component individually

$$ $\displaystyle\sup_{\theta\in\Theta}|\hat{\eta}_{\theta}-\eta_{\theta}|\leq\sup_{\theta}|\hat{\eta}_{\theta}-\hat{\eta}_{\theta^{\prime}}|+\max_{\begin{subarray}{c}\omega^{\prime}\in\{\omega_{1},...,\omega_{T_{\Omega}}\}\\ \gamma^{\prime}\in\{\gamma_{1},...,\gamma_{T_{\Gamma}}\}\end{subarray}}|\hat{\eta}_{\theta^{\prime}}-\eta_{\theta^{\prime}}|+\sup_{\theta}|\eta_{\theta^{\prime}}-\eta_{\theta}|.$ $$

First, let us analyze $|\eta_{\theta}-\eta_{\theta^{\prime}}|$ for any $\theta,\theta^{\prime}\in\Theta$. Recall that $\eta=\operatorname*{\mathbb{E}}[H_{1234}]$ where $H_{1234}=\frac{1}{4!}\sum_{(a,b,c,d)}^{(1,2,3,4)}K_{ab}(L_{ab}+L_{cd}-2L_{ac})$. We have that

$$ $\displaystyle|H^{(\theta)}_{1234}-H^{(\theta^{\prime})}_{1234}|$ $\displaystyle\leq\frac{1}{4!}\sum^{(1234)}_{(abcd)}\left|K^{(\omega)}_{ab}(L^{(\gamma)}_{ab}+L^{(\gamma)}_{cd}-2L^{(\gamma)}_{ac})-K^{(\omega^{\prime})}_{ab}(L^{(\gamma^{\prime})}_{ab}+L^{(\gamma^{\prime})}_{cd}-2L^{(\gamma^{\prime})}_{ac})\right|$ $\displaystyle\leq\frac{1}{4!}\sum^{(1234)}_{(abcd)}\left(\left|K^{(\omega)}_{ab}L^{(\gamma)}_{ab}-K^{(\omega^{\prime})}_{ab}L^{(\gamma^{\prime})}_{ab}\right|+\left|K^{(\omega)}_{ab}L^{(\gamma)}_{cd}-K^{(\omega^{\prime})}_{ab}L^{(\gamma^{\prime})}_{cd}\right|+2\left|K^{(\omega^{\prime})}_{ab}L^{(\gamma^{\prime})}_{ac}-K^{(\omega)}_{ab}L^{(\gamma)}_{ac}\right|\right).$ $$

From Assumption [(A)](#S3.EGx6) we know that $|K_{ab}|\leq\nu_{k}$ and $|L_{ab}|\leq\nu_{l}$, and via Assumption [(C)](#S1.Ex19) we notice that

$$ $\displaystyle\left|K^{(\omega)}_{ab}L^{(\gamma)}_{ab}-K^{(\omega^{\prime})}_{ab}L^{(\gamma^{\prime})}_{ab}\right|$ $\displaystyle=\left|K^{(\omega)}_{ab}L^{(\gamma)}_{ab}-K^{(\omega)}_{ab}L^{(\gamma^{\prime})}_{ab}+K^{(\omega)}_{ab}L^{(\gamma^{\prime})}_{ab}-K^{(\omega^{\prime})}_{ab}L^{(\gamma^{\prime})}_{ab}\right|$ $\displaystyle\leq\left|K^{(\omega)}_{ab}\right|\left|L^{(\gamma)}_{ab}-L^{(\gamma^{\prime})}_{ab}\right|+\left|L^{(\gamma^{\prime})}_{ab}\right|\left|K^{(\omega)}_{ab}-K^{(\omega^{\prime})}_{ab}\right|$ $\displaystyle\leq\nu_{k}L_{l}\|v-v^{\prime}\|+\nu_{l}L_{k}\|\omega-\omega^{\prime}\|$ $\displaystyle\leq\nu_{k}L_{l}\rho_{\Gamma}+\nu_{l}L_{k}\rho_{\Omega}.$ $$

This expression is true for all three components of $|H^{(\theta)}_{1234}-H^{(\theta^{\prime})}_{1234}|$ and so it follows that

$$ $\displaystyle|\eta_{\theta}-\eta_{\theta^{\prime}}|$ $\displaystyle=\left|\operatorname*{\mathbb{E}}[H^{(\theta)}_{1234}]-\operatorname*{\mathbb{E}}[H^{(\theta^{\prime})}_{1234}]\right|\leq\operatorname*{\mathbb{E}}\left|H^{(\theta)}_{1234}-H^{(\theta^{\prime})}_{1234}\right|\leq 4\nu_{k}L_{l}\rho_{\Gamma}+4\nu_{l}L_{k}\rho_{\Omega},$ $\displaystyle|\hat{\eta}_{\theta}-\hat{\eta}_{\theta^{\prime}}|$ $\displaystyle=\left|\frac{1}{(n)_{4}}\sum_{(i,j,q,r)\in{\bf i}^{n}_{4}}H^{(\theta)}_{ijqr}-H^{(\theta^{\prime})}_{ijqr}\right|\leq\frac{1}{(n)_{4}}\sum_{(i,j,q,r)\in{\bf i}^{n}_{4}}\left|H^{(\theta)}_{1234}-H^{(\theta^{\prime})}_{1234}\right|\leq 4\nu_{k}L_{l}\rho_{\Gamma}+4\nu_{l}L_{k}\rho_{\Omega}.$ $$

Now, we study the random error function $\Delta:=\hat{\eta}-\eta$. Note that $\operatorname*{\mathbb{E}}\Delta=0$ since $\hat{\eta}$ is unbiased, and $|H_{ijqr}|\leq 4\nu_{k}\nu_{l}$ via Assumption [(A)](#S3.EGx6). This $\hat{\eta}$, and hence $\Delta$, satisfies bounded differences. Let $F$ denote the kernel tensor $H$ but with sample $(X_{\ell},Y_{\ell})$ replaced by $(X^{\prime}_{\ell},Y^{\prime}_{\ell})$ so that $F$ agrees with $H$ except at indicies $\ell$, and let $\hat{\eta}^{\prime}=\frac{1}{(n)_{4}}\sum_{(i,j,q,r)\in{\bf i}^{n}_{4}}F_{ijqr}$ be it’s HSIC estimator.

For convenience, we denote $(i,j,q,r)\in{\bf i}^{n}_{4}$ simply as $(i,j,q,r)$, and $(i,j,q)\backslash k$ to be the set of 3-tuples drawn without replacement from ${\bf i}^{n}_{3}$ that exclude the number $k$. We can compute the maximal bounded difference $|\Delta-\Delta^{\prime}|=|\hat{\eta}-\hat{\eta}^{\prime}|$ as

$$ $\displaystyle|\hat{\eta}-\hat{\eta}^{\prime}|$ $\displaystyle=\left|\frac{1}{(n)_{4}}\sum_{(i,j,q,r)}H_{ijqr}-F_{ijqr}\right|\leq\frac{1}{(n)_{4}}\sum_{(i,j,q,r)}|H_{ijqr}-F_{ijqr}|$ (10) $\displaystyle=\frac{1}{(n)_{4}}\left(\sum_{(j,q,r)\backslash\ell}\underbrace{|H_{\ell jqr}-F_{\ell jqr}|}_{\leq 8\nu_{k}\nu_{l}}+\sum_{(i,q,r)\backslash\ell}|H_{i\ell qr}-F_{i\ell qr}|+\sum_{(i,j,r)\backslash\ell}|H_{ij\ell r}-F_{ij\ell r}|+\sum_{(i,j,q)\backslash\ell}|H_{ijq\ell}-F_{ijq\ell}|\right)$ $\displaystyle=\frac{1}{(n)_{4}}\Big{(}(n-1)_{3}\cdot 8\nu_{k}\nu_{l}\cdot 4\Big{)}=\frac{32\nu_{k}\nu_{l}}{n}.$ $$

Then, applying McDiarmid’s inequality on $\Delta:=\hat{\eta}-\eta$ followed by a union bound over the $T_{\Omega}T_{\Gamma}$ center pairs gives us, with probability at least $1-\delta$, that

$$ $\displaystyle\max_{\begin{subarray}{c}\omega^{\prime}\in\{\omega_{1},...,\omega_{T_{\Omega}}\}\\ \gamma^{\prime}\in\{\gamma_{1},...,\gamma_{T_{\Gamma}}\}\end{subarray}}|\hat{\eta}_{\theta^{\prime}}-\eta_{\theta^{\prime}}|$ $\displaystyle\leq 32\nu_{k}\nu_{l}\sqrt{\frac{1}{2n}\log\frac{2T_{\Omega}T_{\Gamma}}{\delta}}$ $\displaystyle=\frac{16\nu_{k}\nu_{l}}{\sqrt{n}}\sqrt{2\log\frac{2}{\delta}+2\log T_{\Omega}+2\log T_{\Gamma}}$ $\displaystyle=\frac{16\nu_{k}\nu_{l}}{\sqrt{n}}\sqrt{2\log\frac{2}{\delta}+2D_{\Omega}\log\frac{4R_{\Omega}}{\rho_{\Omega}}+2D_{\Gamma}\log\frac{4R_{\Gamma}}{\rho_{\Gamma}}}.$ $$

Finally, we combine these results to get our uniform convergence bound:

$$ $\displaystyle\sup_{\theta\in\Theta}|\hat{\eta}_{\theta}-\eta_{\theta}|$ $\displaystyle\leq 8\nu_{k}L_{l}\rho_{\Gamma}+8\nu_{l}L_{k}\rho_{\Omega}+\frac{16\nu_{k}\nu_{l}}{\sqrt{n}}\sqrt{2\log\frac{2}{\delta}+2D_{\Omega}\log\frac{4R_{\Omega}}{\rho_{\Omega}}+2D_{\Gamma}\log\frac{4R_{\Gamma}}{\rho_{\Gamma}}}$ $\displaystyle=8\nu_{k}\nu_{l}\left(\frac{L_{k}}{\nu_{k}}\rho_{\Omega}+\frac{L_{l}}{\nu_{l}}\rho_{\Gamma}+\frac{2}{\sqrt{n}}\sqrt{2\log\frac{2}{\delta}+2D_{\Omega}\log\frac{4R_{\Omega}}{\rho_{\Omega}}+2D_{\Gamma}\log\frac{4R_{\Gamma}}{\rho_{\Gamma}}}\right).$ $$

Setting $\rho_{\Omega}=\rho_{\Gamma}=1/\sqrt{n}$ yields the desired result.
∎

###### Proposition 8 .

Under assumptions [(A)](#S3.EGx6) to [(C)](#S1.Ex19), we have that with probability at least $1-\delta$,

$$ $\sup_{\theta\in\Theta}|\hat{\sigma}^{2}_{\theta}-\sigma^{2}_{\theta}|\leq\frac{2048\nu^{2}_{k}\nu^{2}_{l}}{\sqrt{n}}\left(\frac{L_{k}}{\nu_{k}}+\frac{L_{l}}{\nu_{l}}+3\sqrt{2\log\frac{2}{\delta}+2D_{\Omega}\log(4R_{\Omega}\sqrt{n})+2D_{\Gamma}\log(4R_{\Gamma}\sqrt{n})}+\frac{9}{4\sqrt{n}}\right).$ $$

###### Proof.

We use an $\epsilon$-net argument on both spaces $\Omega$ and $\Gamma$. Using the same construction as in [Proposition 7](#Thmtheorem7), we once again decompose our convergence bound:

$$ $\displaystyle\sup_{\theta\in\Theta}|\hat{\sigma}^{2}_{\theta}-\sigma^{2}_{\theta}|\leq\sup_{\theta}|\hat{\sigma}^{2}_{\theta}-\hat{\sigma}^{2}_{\theta^{\prime}}|+\max_{\begin{subarray}{c}\omega^{\prime}\in\{\omega_{1},...,\omega_{T_{\Omega}}\}\\ \gamma^{\prime}\in\{\gamma_{1},...,\gamma_{T_{\Gamma}}\}\end{subarray}}|\hat{\sigma}^{2}_{\theta^{\prime}}-\sigma^{2}_{\theta^{\prime}}|+\sup_{\theta}|\sigma^{2}_{\theta^{\prime}}-\sigma^{2}_{\theta}|.$ $$

First, let us analyze $|\sigma^{2}_{\theta}-\sigma^{2}_{\theta^{\prime}}|$ for any $\theta,\theta^{\prime}\in\Theta$. Recall that $\sigma^{2}=16\left(\operatorname*{\mathbb{E}}[H_{1234}H_{1567}]-\eta^{2}\right)$. It follows that

$$ $\displaystyle|\sigma^{2}_{\theta}-\sigma^{2}_{\theta^{\prime}}|$ $\displaystyle=16\left|\operatorname*{\mathbb{E}}[H^{(\theta)}_{1234}H^{(\theta)}_{1567}-H^{(\theta^{\prime})}_{1234}H^{(\theta^{\prime})}_{1567}]-\operatorname*{\mathbb{E}}[H^{(\theta)}_{1234}H^{(\theta)}_{5678}]+\operatorname*{\mathbb{E}}[H^{(\theta^{\prime})}_{1234}H^{(\theta^{\prime})}_{5678}]]\right|$ $\displaystyle\leq 16\operatorname*{\mathbb{E}}\left|H^{(\theta)}_{1234}H^{(\theta)}_{1567}-H^{(\theta^{\prime})}_{1234}H^{(\theta^{\prime})}_{1567}\right|+16\operatorname*{\mathbb{E}}\left|H^{(\theta)}_{1234}H^{(\theta)}_{5678}-H^{(\theta^{\prime})}_{1234}H^{(\theta^{\prime})}_{5678}\right|.$ $$

Under [Assumptions (A)](#S3.EGx6) and [(C)](#S1.Ex19) we know that $|H_{1234}|\leq 4\nu_{k}\nu_{l}$ and $|H^{(\theta)}_{1234}-H^{(\theta^{\prime})}_{1234}|\leq 4\nu_{k}L_{l}\rho_{\Gamma}+4\nu_{l}L_{k}\rho_{\Omega}$. As such, we have

$$ $\displaystyle|H^{(\theta)}_{1234}H^{(\theta)}_{1567}-H^{(\theta^{\prime})}_{1234}H^{(\theta^{\prime})}_{1567}|$ $\displaystyle\leq|H^{(\theta)}_{1234}H^{(\theta)}_{1567}-H^{(\theta)}_{1234}H^{(\theta^{\prime})}_{1567}|+|H^{(\theta)}_{1234}H^{(\theta^{\prime})}_{1567}-H^{(\theta^{\prime})}_{1234}H^{(\theta^{\prime})}_{1567}|$ $\displaystyle=|H^{(\theta)}_{1234}||H^{(\theta)}_{1567}-H^{(\theta^{\prime})}_{1567}|+|H^{(\theta^{\prime})}_{1567}||H^{(\theta)}_{1234}-H^{(\theta^{\prime})}_{1234}|$ $\displaystyle\leq 32\nu_{k}\nu_{l}(\nu_{k}L_{l}\rho_{\Gamma}+\nu_{l}L_{k}\rho_{\Omega})$ $$

This expression is true for both components of $|\sigma^{2}_{\theta}-\sigma^{2}_{\theta^{\prime}}|$ and so it follows that

$$ $\displaystyle|\sigma^{2}_{\theta}-\sigma^{2}_{\theta^{\prime}}|\leq 1024\nu_{k}\nu_{l}(\nu_{k}L_{l}\rho_{\Gamma}+\nu_{l}L_{k}\rho_{\Omega}).$ (11) $$

Similarly, replacing the expectations $\operatorname*{\mathbb{E}}[H_{1234}H_{1567}]$ and $\operatorname*{\mathbb{E}}[H_{1234}H_{5678}]$ with the respective estimators
$\frac{1}{(n)_{4}(n-1)_{3}}\sum_{(ijqr),(bcd)\backslash i}H_{ijqr}H_{ibcd}$ and $\frac{1}{(n)^{2}_{4}}\sum_{(ijqr),(abcd)}H_{ijqr}H_{abcd}$ give us the same bound

$$ $\displaystyle|\hat{\sigma}^{2}_{\theta}-\hat{\sigma}^{2}_{\theta^{\prime}}|\leq 1024\nu_{k}\nu_{l}(\nu_{k}L_{l}\rho_{\Gamma}+\nu_{l}L_{k}\rho_{\Omega}).$ (12) $$

Next, using [Lemma 9](#Thmtheorem9) and [Lemma 10](#Thmtheorem10) followed by a union bound over the $T_{\Omega}T_{\Gamma}$ center combinations gives us, with probability at least $1-\delta$,

$$ $\displaystyle\max_{\begin{subarray}{c}\omega^{\prime}\in\{\omega_{1},...,\omega_{T_{\Omega}}\}\\ \gamma^{\prime}\in\{\gamma_{1},...,\gamma_{T_{\Gamma}}\}\end{subarray}}|\hat{\sigma}^{2}_{\theta^{\prime}}-\sigma^{2}_{\theta^{\prime}}|$ $\displaystyle\leq 6144\nu^{2}_{k}\nu^{2}_{l}\sqrt{\frac{2}{n}\log\frac{2T_{\Omega}T_{\Gamma}}{\delta}}+\frac{4608\nu^{2}_{k}\nu^{2}_{l}}{n}$ (13) $\displaystyle\leq\frac{2048\nu^{2}_{k}\nu^{2}_{l}}{\sqrt{n}}\left(3\sqrt{2\log\frac{2}{\delta}+2D_{\Omega}\log\frac{4R_{\Omega}}{\rho_{\Omega}}+2D_{\Gamma}\log\frac{4R_{\Gamma}}{\rho_{\Gamma}}}+\frac{9}{4\sqrt{n}}\right).$ $$

Finally, we combine [Equations 11](#S1.E11), [12](#S1.E12) and [13](#S1.E13) to get

$$ $\displaystyle\sup_{\theta\in\Theta}|\hat{\sigma}^{2}_{\theta}-\sigma^{2}_{\theta}|\leq\frac{2048\nu^{2}_{k}\nu^{2}_{l}}{\sqrt{n}}\left(3\sqrt{2\log\frac{2}{\delta}+2D_{\Omega}\log\frac{4R_{\Omega}}{\rho_{\Omega}}+2D_{\Gamma}\log\frac{4R_{\Gamma}}{\rho_{\Gamma}}}+\frac{9}{4\sqrt{n}}+\sqrt{n}\left(\frac{L_{k}}{\nu_{k}}\rho_{\Omega}+\frac{L_{l}}{\nu_{l}}\rho_{\Gamma}\right)\right).$ $$

Setting $\rho_{\Omega}=\rho_{\Gamma}=1/\sqrt{n}$ gives us our desired uniform convergence bound.

∎

###### Lemma 9 .

For any kernels $k$ and $l$ satisfying [Assumption (A)](#S3.EGx6), with probability at least $1-\delta$ we have

$$ $|\hat{\sigma}^{2}-\operatorname*{\mathbb{E}}\hat{\sigma}^{2}|\leq 6144\nu^{2}_{k}\nu^{2}_{l}\sqrt{\frac{2}{n}\log\frac{2}{\delta}}.$ $$

###### Proof.

We apply McDiarmid’s inequality to $\hat{\sigma}^{2}$. First, we show that the variance estimator satisfies bounded differences. For convenience, we denote $(i,j,q,r)\in{\bf i}^{n}_{4}$ simply as $(i,j,q,r)$, and $(i,j,q)\backslash k$ to be the set of 3-tuples drawn without replacement from ${\bf i}^{n}_{3}$ that exclude the number $k$. Recall that

$$ $\hat{\sigma}^{2}=16\left(\frac{1}{(n)_{4}(n-1)_{3}}\sum_{\begin{subarray}{c}(i,j,q,r)\\ (b,c,d)\backslash i\end{subarray}}H_{ijqr}H_{ibcd}-\hat{\eta}^{2}\right).$ $$

Let $F$ denote the kernel tensor $H$ but with sample $(X_{\ell},Y_{\ell})$ replaced by $(X^{\prime}_{\ell},Y^{\prime}_{\ell})$ so that $F$ agrees with $H$ except at indices $\ell$, and let $\hat{\eta}^{\prime}$ and $\hat{\sigma}^{\prime 2}$ denote the HSIC and its variance estimators according to this updated sample set. The deviation is then

$$ $\displaystyle|\hat{\sigma}^{2}-\hat{\sigma}^{\prime 2}|\leq\frac{16}{(n)_{4}(n-1)_{3}}\sum_{\begin{subarray}{c}(i,j,q,r)\\ (b,c,d)\backslash i\end{subarray}}|H_{ijqr}H_{ibcd}-F_{ijqr}F_{ibcd}|+16|\hat{\eta}^{2}-\hat{\eta}^{\prime 2}|.$ $$

We bound the first term by noticing that $\Delta:=H_{ijqr}H_{ibcd}-F_{ijqr}F_{ibcd}$ is zero when none of the indices $\{i,j,q,r,b,c,d\}$ is $\ell$. Let $S:=\{(i,j,q,r,b,c,d):(i,j,q,r)\in{\bf i}^{n}_{4},(b,c,d)\in{\bf i}^{n}_{3}\backslash\{i\},\ell\in\{i,j,q,r,b,c,d\}\}$ be the set of indices where $\Delta$ may be non-zero. By Assumption [(A)](#S3.EGx6) we know that $|\Delta|\leq 32\nu^{2}_{k}\nu^{2}_{l}$. Thus, we can bound the first term by

$$ $\displaystyle\frac{16}{(n)_{4}(n-1)_{3}}\sum_{S}|\Delta|$ $\displaystyle=\frac{512\nu^{2}_{k}\nu^{2}_{l}}{(n)_{4}(n-1)_{3}}|S|$ $\displaystyle=\frac{512\nu^{2}_{k}\nu^{2}_{l}}{(n)_{4}(n-1)_{3}}\left[\underbrace{4(n-1)^{2}_{3}}_{\ell\in\{i,j,q,r\}}+\underbrace{3(n-1)(n-2)_{2}(n-1)_{3}}_{\ell\in\{b,c,d\}}-\underbrace{9(n-1)_{3}(n-2)_{2}}_{\ell\in\{j,q,r\}\text{ and }\ell\in\{b,c,d\}}\right]$ $\displaystyle=512\nu^{2}_{k}\nu^{2}_{l}\left(\frac{16}{n}-\frac{9}{n-1}\right)$ $\displaystyle\leq\frac{8192\nu^{2}_{k}\nu^{2}_{l}}{n}\quad(\forall n>1).$ $$

We can bound the second term using $|\hat{\eta}|\leq 4\nu_{k}\nu_{l}$ ([Assumption (A)](#S3.EGx6)) and the bounded difference result ([10](#S1.E10)) from [Proposition 7](#Thmtheorem7):

$$ $\displaystyle 16|\hat{\eta}^{2}-\hat{\eta}^{\prime 2}|$ $\displaystyle=16|\hat{\eta}+\hat{\eta}^{\prime}||\hat{\eta}-\hat{\eta}^{\prime}|\leq 16\cdot 8\nu_{k}\nu_{l}\cdot\frac{32\nu_{k}\nu_{l}}{n}=\frac{4096\nu^{2}_{k}\nu^{2}_{l}}{n}.$ $$

Combining these two terms, the maximal bounded difference for $\hat{\sigma}^{2}$ is

$$ $\displaystyle|\hat{\sigma}^{2}-\hat{\sigma}^{\prime 2}|\leq\frac{12288\nu^{2}_{k}\nu^{2}_{l}}{n}.$ $$

Finally, applying McDiarmid’s inequality gives us, with probability at least $1-\delta$,

$$ $\displaystyle|\hat{\sigma}^{2}-\operatorname*{\mathbb{E}}\hat{\sigma}^{2}|\leq 6144\nu^{2}_{k}\nu^{2}_{l}\sqrt{\frac{2}{n}\log\frac{2}{\delta}}.$ $$

∎

###### Lemma 10 .

For any kernels $k$ and $l$ satisfying [Assumption (A)](#S3.EGx6), the bias is bounded by

$$ $|\operatorname*{\mathbb{E}}\hat{\sigma}^{2}-\sigma^{2}|\leq\frac{4608\nu^{2}_{k}\nu^{2}_{l}}{n}.$ $$

###### Proof.

The expectation of the variance estimator is

$$ $\operatorname*{\mathbb{E}}\hat{\sigma}^{2}=16\left(\frac{1}{(n)_{4}(n-1)_{3}}\sum_{\begin{subarray}{c}(i,j,q,r)\\ (b,c,d)\backslash i\end{subarray}}\operatorname*{\mathbb{E}}[H_{ijqr}H_{ibcd}]-\frac{1}{(n)^{2}_{4}}\sum_{\begin{subarray}{c}(i,j,q,r)\\ (a,b,c,d)\end{subarray}}\operatorname*{\mathbb{E}}[H_{ijqr}H_{abcd}]\right).$ $$

First, we can break down the left-hand sum into only terms of $\operatorname*{\mathbb{E}}[H_{1234}H_{1567}]$ by considering the cases where $\{i,j,q,r,b,c,d\}$ are unique. Let $S=\{(i,j,q,r,b,c,d):(i,j,q,r)\in{\bf i}^{n}_{4},(b,c,d)\in{\bf i}^{n}_{3}\backslash\{i\}\}$ be the set of all possible indices of our left-hand sum. It follows that

$$ $\displaystyle\sum_{S}\operatorname*{\mathbb{E}}[H_{ijqr}H_{ibcd}]=\sum_{(i,j,q,r,b,c,d)\in{\bf i}^{n}_{7}}\operatorname*{\mathbb{E}}[H_{ijqr}H_{ibcd}]+\sum_{S\backslash{\bf i}^{n}_{7}}\operatorname*{\mathbb{E}}[H_{ijqr}H_{ibcd}].$ $$

If all indices are unique, then the expectation $\operatorname*{\mathbb{E}}[H_{ijqr}H_{ibcd}]$ is equivalent to $\operatorname*{\mathbb{E}}[H_{1234}H_{1567}]$; otherwise, we can bound the expectation by $16\nu^{2}_{k}\nu^{2}_{l}$ via Assumption [(A)](#S3.EGx6). Thus, the bound on the left-hand sum is

$$ $\displaystyle\sum_{\begin{subarray}{c}(i,j,q,r)\\ (b,c,d)\backslash i\end{subarray}}\operatorname*{\mathbb{E}}[H_{ijqr}H_{ibcd}]\leq(n)_{7}\operatorname*{\mathbb{E}}[H_{1234}H_{1567}]+\Big{(}(n)_{4}(n-1)_{3}-(n)_{7}\Big{)}16\nu^{2}_{k}\nu^{2}_{l}.$ $$

Similarly, we can break down the right-hand sum into only terms of $\operatorname*{\mathbb{E}}[H_{1234}H_{5678}]$. Let $R=\{(i,j,q,r,a,b,c,d):(i,j,q,r)\in{\bf i}^{n}_{4},(a,b,c,d)\in{\bf i}^{n}_{4}\}$ be the possible indices of our right-hand sum. We have that

$$ $\displaystyle\sum_{\begin{subarray}{c}(i,j,q,r)\\ (a,b,c,d)\end{subarray}}\operatorname*{\mathbb{E}}[H_{ijqr}H_{abcd}]$ $\displaystyle=\sum_{(i,j,q,r,a,b,c,d)\in{\bf i}^{n}_{8}}\operatorname*{\mathbb{E}}[H_{ijqr}H_{abcd}]+\sum_{R\backslash{\bf i}^{n}_{8}}\operatorname*{\mathbb{E}}[H_{ijqr}H_{abcd}]$ $\displaystyle\leq(n)_{8}\operatorname*{\mathbb{E}}[H_{1234}H_{5678}]+\Big{(}(n)^{2}_{4}-(n)_{8}\Big{)}16\nu^{2}_{k}\nu^{2}_{l}.$ $$

Now, using these two results and Assumption [(A)](#S3.EGx6), we can compute a bound on the desired bias of $\hat{\sigma}^{2}$:

$$ $\displaystyle|\operatorname*{\mathbb{E}}\hat{\sigma}^{2}-\sigma^{2}|$ $\displaystyle=16\left|\frac{1}{(n)_{4}(n-1)_{3}}\sum_{\begin{subarray}{c}(i,j,q,r)\\ (b,c,d)\backslash i\end{subarray}}\operatorname*{\mathbb{E}}[H_{ijqr}H_{ibcd}]-\operatorname*{\mathbb{E}}[H_{1234}H_{1567}]-\frac{1}{(n)^{2}_{4}}\sum_{\begin{subarray}{c}(i,j,q,r)\\ (a,b,c,d)\end{subarray}}\operatorname*{\mathbb{E}}[H_{ijqr}H_{abcd}]+\operatorname*{\mathbb{E}}[H_{1234}H_{5678}]\right|$ $\displaystyle\leq 16\left|\left(1-\frac{(n)_{7}}{(n)_{4}(n-1)_{3}}\right)\left(16\nu^{2}_{k}\nu^{2}_{l}-\underbrace{\operatorname*{\mathbb{E}}[H_{1234}H_{1567}]}_{-16\nu^{2}_{k}\nu^{2}_{l}\leq\cdot\leq 16\nu^{2}_{k}\nu^{2}_{l}}\right)+\left(1-\frac{(n)_{8}}{(n)^{2}_{4}}\right)\left(\underbrace{\operatorname*{\mathbb{E}}[H_{1234}H_{5678}]}_{0\leq\cdot\leq 16\nu^{2}_{k}\nu^{2}_{l}}-16\nu^{2}_{k}\nu^{2}_{l}\right)\right|.$ $\displaystyle\leq 16\left(1-\frac{(n)_{7}}{(n)_{4}(n-1)_{3}}\right)32\nu^{2}_{k}\nu^{2}_{l}<512\nu^{2}_{k}\nu^{2}_{l}\cdot\frac{9}{n}\quad(\forall n\geq 4)$ $\displaystyle=\frac{4608\nu^{2}_{k}\nu^{2}_{l}}{n}.\qed$ $$

###### Proposition 7 .

###### Proof.

###### Proposition 8 .

###### Proof.

###### Lemma 9 .

###### Proof.

###### Lemma 10 .

###### Proof.

## B Additional Experiments

### B.1 High-Dimensional Gaussian Mixture

Figure: (a) HDGM-4
Refer to caption: /html/2409.06890/assets/x21.png

### B.2 Type-I Error

[Table 1](#S2.T1) shows that the type-I error rates for our optimization-based tests are well-controlled.

Table: Table 1: Average type-I error rates under the null distribution over 400 tests. We use $m=512$ samples.

### B.3 Optimizing J vs. HSIC

We examine the trade-off between optimizing the approximate test power $J$ versus just the test statistic HSIC. The results on power versus test size are show in [Figure 6](#S2.F6). Optimizing our proposed objective $J$ significantly outperforms optimizing HSIC for all problems.

Figure: (a) HDGM-10
Refer to caption: /html/2409.06890/assets/x29.png

### B.4 Asymptotic Variance: HSIC vs. MMD

We show estimates of the asymptotic variance of HSIC and MMD along a training trajectory in [Figure 7](#S2.F7). For MMD we consider both a single shuffling of the data (MMD-full), as well as a split shuffling (MMD-split) where we use half the data for our joint distribution sample, and the other half to permute for our product-of-marginals sample.

We note that the initial variance of MMD-split is substantially higher than that of MMD-full, which is much higher than the variance of HSIC. We hypothesize this makes the MMD-based objective harder to optimize, and gives a possible explanation for why MMD-split performs worse than MMD-full. MMD-split/full also exhibit greater final variances, particularly at larger batch sizes.

Figure: (a) $n=128$
Refer to caption: /html/2409.06890/assets/x33.png

### B.5 Training Times

We compare the training times of all methods in [Table 2](#S2.T2). Our method is comparable in speed to C2ST and MMD, and vastly faster than the variational mutual information bound methods InfoNCE and NWJ since both those methods require function evaluations for every possible pair of samples at each training step.

**Table 2: Approximate training times for the HDGM-30 problem over 1,000 epochs on a GeForce RTX 4060. Times are rounded to the nearest minute**
|  | Training Time |
| --- | --- |
| HSIC | 41 m |
| C2ST | 34 m |
| MMD | 44 m |
| InfoNCE | 8 h 38 m |
| NWJ | 8 h 42 m |

## C Experiment Details

Figure: Figure 8: RatInABox simulation environment. The red dot is the current position of the rat and the purple circles indicate the past trajectory over 5 seconds. The box is designed to have only a single protruding wall.
Refer to caption: /html/2409.06890/assets/x36.png

### C.1 Training & Test Details

We design the featurizers $\phi_{\omega}$ and $\phi_{\kappa}$ of our deep kernels $k_{\omega}$ and $l_{\kappa}$ to be neural networks with ReLU activations. We avoid using normalization as it may affect test power. Moreover, we make the Gaussian bandwidth of both $k_{\omega}$ and $l_{\kappa}$ a learnable parameter, as well as the smoothing rate $\epsilon$. To make comparisons as fair as possible, we use similar neural network architectures for each deep learning based method. In general, we let the featurizer of HSIC-D and MMD-D be identical up to a concatenation layer which concatenates $X$ and $Y$ to frame the problem as a two-sample test. We construct the C2ST-S/L classifier as the MMD-D featurizer plus a linear layer classification head with scalar output, and we let both C2ST-S/L and InfoNCE use identical architectures for the classifier and critic. Detailed descriptions of each architecture are demonstrated in the following subsections.

All optimization-based methods (HSIC-D/Dx/O, MMD-D, C2ST-S/L, InfoNCE) are first trained on an identical split of the data, and then tested on the remaining split. In contrast, HSIC-M selects the median bandwidth based on the entire dataset, and is evaluated on the test set. We train our models using the AdamW optimizer with a learning rate of 1e-4 over 1,000 epochs for HDGM and RatInABox and 10,000 epochs for Sinusoid. We use a batch size of 512. All methods are implemented in PyTorch and trained on a NVIDIA A100SXM4 GPU.

Once learned, each methods’ empirical power is evaluated on 100 test sets $(S^{te_{1}}_{Z},...,S^{te_{100}}_{Z})$. Each test set contains $m$ test samples $S^{te_{i}}_{Z}=(Z^{te_{i}}_{1},...,Z^{te_{i}}_{N})$, which are then used to compute the average rejection rate under the null via a permutation test. We use 500 permutations for each test and with a predetermined type-I error rate of 0.05.

### C.2 Architectures

In all experiments we consider deep kernels with Gaussian feature and smoothing kernels $\kappa$ and $q$, where each bandwidth is a trainable parameter randomly initialized around $1.0$. We let the smoothing weight $\epsilon$ also be a learnable parameter initialized to $0.01$. No batch normalization is used and all hidden layers use ReLU activations. Dataset-specific designs are elaborated below.

High-dimensional Gaussian mixture.
We use a feed-forward network for our deep kernel featurizer with latent dimensions $2d,3d,\text{and }2d$. Details of each model is given in [Table 3](#S3.T3).

Sinusoid.
The deep kernel featurizer is taken to be a feed-forward network with widths 1x8x12x8. C2st, infoNCE, and NWJ use a similar architecture –one with widths 2x8x12x8x1– which includes an additional scalar output layer.

RatInABox.
We use a feed-forward featurizer with details given in [Table 4](#S3.T4). Unlike the previous two problems, the sample spaces $\mathcal{X}$ and $\mathcal{Y}$ are not equivalent, and so the deep featurizers for $k$ and $l$ have different architectures.

**Table 3: Featurizer architectures used in deep kernels for HSIC-D, MMD-D, and classifier architecture used for C2ST-S/L on the HDGM problem. Brackets denote a sequence of linear layers with corresponding input and output features.**
| dataset | model | input | featurizer |
| --- | --- | --- | --- |
| HDGM-4 | HSIC-D | X or Y | $\left[\begin{array}[]{c}2\rightarrow 4\rightarrow 6\rightarrow 4\\<br>\end{array}\right]$ |
| MMD-D | [X, Y] | $\left[\begin{array}[]{c}4\rightarrow 8\rightarrow 12\rightarrow 8\\<br>\end{array}\right]$ |  |
| C2ST-S/L | [X, Y] | $\left[\begin{array}[]{c}4\rightarrow 8\rightarrow 12\rightarrow 8\rightarrow 1\\<br>\end{array}\right]$ |  |
| HDGM-8 | HSIC-D | X or Y | $\left[\begin{array}[]{c}4\rightarrow 8\rightarrow 12\rightarrow 8\\<br>\end{array}\right]$ |
| MMD-D | [X, Y] | $\left[\begin{array}[]{c}8\rightarrow 16\rightarrow 24\rightarrow 16\\<br>\end{array}\right]$ |  |
| C2ST-S/L | [X, Y] | $\left[\begin{array}[]{c}8\rightarrow 16\rightarrow 24\rightarrow 16\rightarrow 1\\<br>\end{array}\right]$ |  |
| HDGM-10 | HSIC-D | X or Y | $\left[\begin{array}[]{c}5\rightarrow 10\rightarrow 15\rightarrow 10\\<br>\end{array}\right]$ |
| MMD-D | [X, Y] | $\left[\begin{array}[]{c}10\rightarrow 20\rightarrow 30\rightarrow 20\\<br>\end{array}\right]$ |  |
| C2ST-S/L | [X, Y] | $\left[\begin{array}[]{c}10\rightarrow 20\rightarrow 30\rightarrow 20\rightarrow 1\\<br>\end{array}\right]$ |  |
| HDGM-20 | HSIC-D | X or Y | $\left[\begin{array}[]{c}10\rightarrow 20\rightarrow 30\rightarrow 20\\<br>\end{array}\right]$ |
| MMD-D | [X, Y] | $\left[\begin{array}[]{c}20\rightarrow 40\rightarrow 60\rightarrow 40\\<br>\end{array}\right]$ |  |
| C2ST-S/L | [X, Y] | $\left[\begin{array}[]{c}20\rightarrow 40\rightarrow 60\rightarrow 40\rightarrow 1\\<br>\end{array}\right]$ |  |

**Table 4: Featurizer architectures used in deep kernels for HSIC-D, MMD-D, and classifier architecture used for C2ST-S/L on the RatInABox problem. Brackets denote a sequence of linear layers with corresponding input and output features.**
| method | input | network |
| --- | --- | --- |
| HSIC-D | X | $\left[\begin{array}[]{c}8\rightarrow 32\rightarrow 64\rightarrow 32\\<br>\end{array}\right]$ |
| Y | $\left[\begin{array}[]{c}2\rightarrow 4\rightarrow 8\rightarrow 4\\<br>\end{array}\right]$ |  |
| MMD-D | [X, Y] | $\left[\begin{array}[]{c}10\rightarrow 32\rightarrow 64\rightarrow 32\\<br>\end{array}\right]$ |
| C2ST-S/L | [X, Y] | $\left[\begin{array}[]{c}10\rightarrow 32\rightarrow 64\rightarrow 32\rightarrow 1\\<br>\end{array}\right]$ |

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 5
presentation: 2
questions: Please see the weakness section.
soundness: 3
strengths: The paper is technically well presented, and the claims are well supported, theoretically and experimentally. It introduces novel approaches for learning representations that enhances the power of independence tests, particularly in high-dimensional settings, improving upon traditional methods. Additionally, the work demonstrates theoretical validity and empirical effectiveness of the proposed approaches, complemented by extensive evaluations that reinforce its contributions to the field.
summary: The paper discusses methods for improving independence tests between random variables, focusing on high-dimensional data where traditional tests struggle. It introduces two main approaches: using variational estimators of mutual information (MI), such as InfoNCE and NWJ, and leveraging the Hilbert-Schmidt Independence Criterion (HSIC). The authors demonstrate that variational MI-based tests can be closely related to kernel-based HSIC tests, with both methods aiming to detect dependencies more efficiently. They propose a new method that learns representations to maximize the power of HSIC tests, showing that this approach outperforms standard variational methods in detecting complex dependencies.
weaknesses: - A notable issue is the lack of clarity regarding the aims and achievements of the paper. For instance, while it claims to present two approaches for independence testing, the conclusion suggests that it only studies one approach, leading to confusion about the overall contribution and focus of the work.
- The performance of HSIC-based tests is significantly influenced by the choice of kernel. While the paper proposes methods to optimize kernel selection, identifying the appropriate kernel for specific data types or distributions remains challenging and may not generalize well across diverse datasets. 

- Although the proposed methods enhance sample efficiency compared to traditional tests, they still necessitate a considerable number of samples to effectively detect subtle dependencies in very high-dimensional data, which could limit their practicality in scenarios with restricted data availability. 

- Additionally, the paper does not provide theoretical rates for the convergence of test statistics and lacks a discussion on the asymptotic validity of power in its main content. 

- Furthermore, regarding Algorithm 1, the choice of $ n_{\text{perm}}$ is required; what are the best empirical strategies for selecting this parameter, and how does the choice impact the convergence rate?

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
questions: Questions:
1. Although using permutation tests is easy, it may be less effective and computationally costly. Is it possible to derive a limiting distribution for your test statistics? Also, the computational complexity of Phase Two in Algorithm 1 needs to be discussed.

Small comments:
1. Line 218. If $X \perp Y$ -> if $X \not\perp Y$.
2. Line 223. When HSIC = 0 the statistic is not mean zero.
3. Line 417 and 419. There is no $\lambda$ in $J_\lambda$.
soundness: 3
strengths: This paper is very well written and I was able to read it smoothly. The idea of minimizing the asymptotic power instead of the dependence metric itself is interesting. It is demonstrated in the experiments that their approach improves HSIC's power in high-dimensional problems.
summary: This paper proposes a method to learn representations that facilitates the detection of dependence. The authors instantiate their idea using variational estimators of mutual information and the Hilbert-Schmidt Independence Criterion. Finally, they provide theoretical and empirical evidence to demonstrate the effectiveness of their approach.
weaknesses: A few concerns I have for this paper:

1. It seems to me that the authors didn't highlight their core contribution enough. The InfoNCE approach already exists in the literature and making an independence test out of it based on permutation is not very interesting; applying this trick to HSIC either. The interesting part to me is they propose to minimize the SNR -- a quantity that characterizes the statistical power -- instead of the test statistics themselves. However, this contribution is not supported in the experiments (or is not conveyed clearly). I would like to see a comparison between the InfoNCE (HSIC as well) that minimizes the lower bound (test statistic) versus the one that minimizes the SNR. If the latter shows significant improvement upon the former, I would then say the contribution is solid.

2. The authors didn't mention a closely related topic -- independence tests based on optimal transport (e.g., [1], [2], and references therein). In this line of work, the distance between the joint and the production of marginals is measured by the optimal transport distances.

[1] D. M. Cifarelli and E. Regazzini. On the centennial anniversary of Gini’s theory of statistical relations. Metron, 2017.

[2] L. Liu, S. Pal, Z. Harchaoui. Entropy Regularized Optimal Transport Independence Criterion. AISTATS, 2022.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 2
questions: 1. Can the authors comment on the tightness of the bound presented in proposition 3 (beyond the asymptotics)?
2. Does the MMD in line 255 depend on the function $f$ or function family $\cal{H}$? it is not clear from the notation, while I am pretty sure it does.
3. Could the authors discuss the assumptions of Theorem 4? I would also expect a short discussion on their 'reasonability' i.e., which distributions are captued under such assumptions?
4. Theorem discusses the implication of a unique maximizer. Do we know when such a maximizer exists?
5. Did the authos perform any ablation on the method? It would be specifically interesting to see if the networks size should indeed be linear in the input dimension.
5. The presentation would benefit from clarifying what is K in the InfoNCE notation.
6. "Gaussian kernel with unit bandwidth ... this scheme can perform extremely poorly" - the intro would benefit from a short discussion on such cases, as they form some of the motivation for this work.
soundness: 3
strengths: 1. I like the idea of deriving an optimizer for independence testing directly from the hypothesis test error. This makes its motivation more explicit, rather than information theoretic methods which can be thought of as 'indirect', or inducing independence as a byproduct. The authors directly demonstrate how variational mi bounds are not aiming for P_e minimization.
2. The authors present formal guarantees of the proposed method and connect it to variational bounds of mutual information, which are highly popular methods that are widespread in the community. 
3. When discussing the expression in neural estimation that is effectively required for independence testing - I liked the permutation invariance interpretation. (A note for the authors - this can be also seen from the proof of the DV formula)
summary: The authors introduce two methods for independence testing, which are based on known methodologies. Two approaches are explored:
1. Tests based on variational mutual information estimators, such as InfoNCE, 
2. The HSIC that aims to maximize test power in high-dimensional data.
By recasting the corresponding test as optimization problems over neural nets, the authors demonstrate that HSIC-based tests outperform MI-based ones, particularly for complex dependencies.
weaknesses: 1. I feel that the paper could benefit from better representation, which makes it harder for the reader to understand theciontributions in a bird's view, e.g., some crucial parts are migrated to the appendix, such as the assumptions of the main theorem of the paper, while some parts do not benefit the presentation in my opinion, such as some of the very specific and technical details in Section 5.
2. As the authors utilize 'heavy' machinery such as NN optimization, I would expect at least one experimental results on real world high dimensional data.
3. Literature review: Independence testing has been previously explored in the form of optimization of families of functions. Specifically, classical methods such as canonical correlation analysis (Hoteling 67) and its many extensions tackle the problem of independence testing. In an information theoretic setting there are the Hirschfeld-Gebelein-Renyi coefficient. More recently, max-sliced mutual information (Tsur 24), which was shown useful for independence testing and was connected to constrained variational mutual information estimation.
4. Section 2 lacks a bottom line, and currently feels like it 'stops in the middle'
5. I am not sure the citation formal is valid.

Technical comments:
1. undefined acronyms - InfoNCE, NWJ,DV
2. line 101 - 'method method'
3. line 136 - 'feature'-'features'
4. line 147 - 'by by'

---

confidence: 2
contribution: 3
initial_rating: 6
presentation: 3
questions: - What is the definition of K in Line 130?
- You mention future work on extending this approach to conditional independence testing and causal discovery. Can you provide more details on what specific challenges you foresee in adapting the current method for conditional independence?
soundness: 3
strengths: The paper is well-organized and clearly presented. The authors clearly explain their methods and provide intuitive connections between the variational bounds and kernel-based methods.

Theoretically, the paper is the first to construct a valid independence test based on the variational mutual information estimator in high dimensions. The paper provides theoretical insights into the connection between HSIC and MI tests and shows how to maximize the test power for independence testing using learned representations. 

Although the experiments were all conducted on synthetic datasets, the experimental validation is abundant with various settings. The empirical results show that the performance of proposed HSIC-based tests is nontrivial especially in high dimensions.
summary: The paper introduces new methods to detect dependencies between random variables in high-dimensional data by utilizing variational estimators of MI  and the HSIC. It aims to improve the power of independence tests by leveraging deep learning techniques to learn effective representations of the data, making the dependencies between variables more explicit. The paper also highlights a connection between MI-based tests and HSIC-based tests, showing that they are closely related. The paper introduces the concept of maximizing test power using learned representations, specifically focusing on HSIC-based tests, which are generally found to outperform MI-based ones. The experiments are abundant and support the claims.
weaknesses: One limitation could be the lack of evaluation on real-world datasets despite the authors already presenting three quite informative synthetic datasets. Since the presented tests only focus on testing independence, experiments of causal discovery may not be suitable. But I think testing a pair of variables is still executable. Experiments on any real-world datasets are appreciated, e.g., https://www.cmu.edu/dietrich/causality/projects/causal_learn_benchmarks/. This is just an example, the author does not need to be constrained by this and can use any real-world dataset. If there are any reasons why experiments on real-world datasets may not be feasible or practical, I would be very interested in hearing them and would welcome the explanation with an open mind.

Another limitation, which is also discussed by the authors, is the scope only focusing on the independence test, which is considerably easier compared with the conditional independence test. 

A minor piece of advice: A paragraph or several sentences to highlight the difference between the traditional independence test and the independence test in high dimensions will be appreciated and suitable for readers from different levels.

Meta-review: This work explores two approaches to enhance independence testing between random variables, crucial for machine learning and statistics. The first approach uses variational estimators of mutual information, like InfoNCE and NWJ, to construct powerful finite-sample valid tests, while the second is based on the Hilbert-Schmidt Independence Criterion (HSIC).  The paper demonstrates learning a variational bound in the former is closely related to learning kernels in the latter.  Empirical results show the benefit of learning representations for the HSIC based method in detecting complex structured dependencies.

The paper is well written and has interesting theoretical and practical results.  The empirical evaluation on real world datasets is still quite limited, and more analysis is needed on the sample complexity related to kernel learning.

Decision: Reject