Package: coreax
Coreax© Crown Copyright GCHQCoreax is a library forcoreset algorithms, written inJAXfor fast execution and GPU support.For $n$ points in $d$ dimensions, a coreset algorithm takes an $n \times d$ data set and
reduces it to $m \ll n$ points whilst attempting to preserve the statistical properties
of the full data set. The algorithm maintains the dimension of the original data set.
Thus the $m$ points, referred to as thecoreset, are also $d$-dimensional.The $m$ points need not be in the original data set. We refer to the special case where
all selected points are in the original data set as acoresubset.Some algorithms return the $m$ points with weights, so that importance can be
attributed to each point in the coreset. The weights, $w_i$ for $i=1,...,m$, are often
chosen from the simplex. In this case, they are non-negative and sum to 1:
$w_i >0$ $\forall i$ and $\sum_{i} w_i =1$.Quick exampleConsider $n=10,000$ points drawn from six $2$-D multivariate Gaussian distributions. We
wish to reduce this to only 100 points, whilst maintaining underlying statistical
properties. We achieve this by generating a coreset, setting $m=100$.
We plot the underlying data (blue) as-well as the coreset
points (red), which are plotted sequentially based on the order the algorithm selects
them in. The coreset points are weighted (size of point) to optimally reconstruct the
underlying distribution. Runexamples/herding_stein_weighted.pyto replicate.We compare the coreset to the full original dataset by calculating the maximum
mean discrepancy (MMD).
This key property is an integral probability metric, measuring
the distance between the empirical distributions of the full dataset and the coreset.
A good coreset algorithm produces a coreset that has significantly smaller MMD
than randomly sampling the same number of points from the original data, as is the case
in the example below.Kernel herdingRandom sampleExample applicationsChoosing pixels from an image: In the example below, we reduce the original 180x215
pixel image (38,700 pixels in total) to a coreset approximately 20% of this size.
(Left) original image.
(Centre) 8,000 coreset points chosen using Stein kernel herding, with point size a
function of weight.
(Right) 8,000 points chosen randomly.
Runexamples/david_map_reduce_weighted.pyto  replicate.Video event detection: Here we identify representative frames such that most of the
useful information in a video is preserved.
Runexamples/pounce.pyto replicate.OriginalCoresetSetupBefore installing coreax, make sure JAX is installed. Be sure to install the preferred
version of JAX for your system.InstallJAX, noting that there are (currently) different setup paths for CPU and GPU use.Install coreax by cloning the repo and then runningpip install .from your local coreax directory.To install additional optional dependencies required to run the examples inexamplesusepip install .[test]instead.A how-to guideHere are some of the most commonly used classes and methods in the library.Kernel herdingKernel herding is one (greedy) approach to coreset construction.
Acoreax.coresubset.KernelHerdingobject is created by supplying acoreax.kernel.Kernelobject, such as aSquaredExponentialKernel. A coreset is
generated by calling thefitmethod on the kernel herding object.Note that, throughout the codebase, there are block versions of herding for fitting
within memory constraints. These methods partition the data into blocks before carrying
out the coreset algorithm, restricting the maximum size of variables handled in the process.fromsklearn.datasetsimportmake_blobsimportnumpyasnpimportjax.randomfromcoreaximport(ArrayData,KernelHerding,SizeReduce,SquaredExponentialKernel,)fromcoreax.kernelimportmedian_heuristic# Generate some datanum_data_points=10_000num_features=2num_cluster_centers=6random_seed=1989x,_=make_blobs(num_data_points,n_features=num_features,centers=num_cluster_centers,random_state=random_seed,)# Request 100 coreset pointscoreset_size=100# Setup the original data objectdata=ArrayData.load(x)# Set the bandwidth parameter of the kernel using a median heuristic derived from# at most 1000 random samples in the data.num_samples_length_scale=min(num_data_points,1_000)generator=np.random.default_rng(random_seed)idx=generator.choice(num_data_points,num_samples_length_scale,replace=False)length_scale=median_heuristic(x[idx])# Compute a coreset using kernel herding with a squared exponential kernel.herding_key=jax.random.key(random_seed)herding_object=KernelHerding(herding_key,kernel=SquaredExponentialKernel(length_scale=length_scale))herding_object.fit(original_data=data,strategy=SizeReduce(coreset_size=coreset_size))# The herding object now has the coreset, and the indices of the original data# that makeup the coreset as populated attributesprint(herding_object.coreset)print(herding_object.coreset_indices)Kernel herding with weightingA coreset can be weighted, a so-calledweighted coreset, to attribute importance to
each point and to better approximate the underlying data distribution.
Optimal weights can be determined by implementing acoreax.weights.WeightsOptimiser, such as theMMDWeightsOptimiser.fromcoreaximport(KernelHerding,SizeReduce,SquaredExponentialKernel,)fromcoreax.weightsimportMMDasMMDWeightsOptimiser# Define a kernelkernel=SquaredExponentialKernel(length_scale=length_scale)# Define a weights optimiser to learn optimal weights for the coreset after creationweights_optimiser=MMDWeightsOptimiser(kernel=kernel)# Compute a coreset using kernel herding with a squared exponential kernel.herding_object=KernelHerding(herding_key,kernel=kernel,weights_optimiser=weights_optimiser)herding_object.fit(original_data=data,strategy=SizeReduce(coreset_size=coreset_size))# Determine optimal weights for the coresetherding_weights=herding_object.solve_weights()Kernel herding with refineTo improve the quality of a coreset, arefinestep can be added.
These functions work by substituting points from the coreset with points from the
original dataset such that the MMD decreases. This improves the
coreset quality because the refined coreset better captures the
underlying distribution of the original data, as measured by the reduced MMD.There are several different approaches to refining a coreset, which can be found in the
classes and methods incoreax.refine. In the example below, we create a refiner object,
pass it to the herding object, and then call the refine method.fromcoreaximport(KernelHerding,SizeReduce,SquaredExponentialKernel,)fromcoreax.refineimportRefineRegular# Define a refinement objectrefiner=RefineRegular()# Compute a coreset using kernel herding with a squared exponential kernel.herding_object=KernelHerding(herding_key,kernel=SquaredExponentialKernel(length_scale=length_scale),refine_method=refiner)herding_object.fit(original_data=data,strategy=SizeReduce(coreset_size=coreset_size))# Refine the coreset to improve qualityherding_object.refine()# The herding object now has the refined coreset, and the indices of the original# data that makeup the refined coreset as populated attributesprint(herding_object.coreset)print(herding_object.coreset_indices)Scalable herdingFor large $n$ or $d$, you may run into time or memory issues. The classcoreax.reduction.MapReduceuses partitioning to tractably compute an approximate
coreset in reasonable time.
There is a necessary impact on coreset quality, for a dramatic improvement in computation time.
These methods can be used by simply replacingcoreax.reduction.SizeReducein the
previous examples withMapReduceand setting the parameterleaf_sizein line with
memory requirements.fromcoreax.coresubsetimportKernelHerdingfromcoreax.kernelimportSquaredExponentialKernelfromcoreax.reductionimportMapReduce# Compute a coreset using kernel herding with a squared exponential kernel.herding_object=KernelHerding(herding_key,kernel=SquaredExponentialKernel(length_scale=length_scale),)herding_object.fit(original_data=data,strategy=MapReduce(coreset_size=coreset_size,leaf_size=200))For large $d$, it is usually worth reducing dimensionality using PCA. Seeexamples/pounce_map_reduce.pyfor an example.Stein kernel herdingWe have implemented a version of kernel herding that uses aStein kernel, which
targetskernelised Stein discrepancy (KSD)rather than MMD.
This can often give better integration error in practice, but it can be slower than
using a simpler kernel targeting MMD.
To use Stein kernel herding, we have to define a
continuous approximation to the discrete measure, e.g. using kernel density estimation (KDE),
or an estimate the score function $\nabla \log f_X(\mathbf{x})$ of a continuous PDF from
a finite set of samples.
In this example, we use a Stein kernel with a squared exponential base
kernel, computing the score function explicitly.importnumpyasnpfromcoreaximport(SquaredExponentialKernel,SteinKernel,KernelDensityMatching,)# Select indices to form a subset of data for learning score functiongenerator=np.random.default_rng(random_seed)idx=generator.choice(len(data),subset_size,replace=False)data_subset=data[idx,:]# Learn a score function from the subset of the data, through a kernel density# estimation applied to a subset of the data.kernel_density_score_matcher=KernelDensityMatching(length_scale=length_scale,kde_data=data_subset)score_function=kernel_density_score_matcher.match()# Define a kernel to use for herdingherding_kernel=SteinKernel(SquaredExponentialKernel(length_scale=length_scale),score_function=score_function,)# Compute a coreset using kernel herding with a Stein kernelherding_object=KernelHerding(herding_key,kernel=herding_kernel)herding_object.fit(original_data=data,strategy=SizeReduce(coreset_size=coreset_size))Score matchingThe score function, $\nabla \log f_X(\mathbf{x})$, of a distribution is the derivative
of the log-density function. This function is required when evaluating Stein kernels.
However, it can be difficult to specify analytically in practice.To resolve this, we have implemented an approximation of the score function using a
neural network as inSong et al. (2019).
This approximate score function can then be passed directly to a Stein kernel, removing
any requirement for analytical derivation. More details on score matching methods
implemented are found incoreax.score_matching.importnumpyasnpfromcoreaximport(SteinKernel,SlicedScoreMatching,)fromcoreax.kernelimportPCIMQKernel# Select indices to form a subset of data for learning score functiongenerator=np.random.default_rng(random_seed)idx=generator.choice(len(data),subset_size,replace=False)data_subset=data[idx,:]# Learn a score function from a subset of the data, through approximation using a neural# network applied to a subset of the datascore_key=jax.random.key(random_seed)sliced_score_matcher=SlicedScoreMatching(score_key,random_generator=jax.random.rademacher,use_analytic=True,num_epochs=10,num_random_vectors=1,sigma=1.0,gamma=0.95,)score_function=sliced_score_matcher.match(data_subset)# Define a kernel to use for herdingherding_kernel=SteinKernel(PCIMQKernel(length_scale=length_scale),score_function=score_function,)Release cycleWe anticipate two release types: feature releases and security releases. Security
releases will be issued as needed in accordance with thesecurity policy. Feature releases will
be issued as appropriate, dependent on the feature pipeline and development priorities.Coming soonSome features coming soon include:Coordinate bootstrapping for high-dimensional data.Other coreset-style algorithms, including kernel thinning and recombination, as means
to reducing a large dataset whilst maintaining properties of the underlying distribution.