Package: grad-info-opt
GIO: Gradient Information OptimizationGIO is a library that implements Gradient Information Optimization (GIO) at scale, from the paperGIO: Gradient Information Optimization for Training Dataset Selection. GIO is a data selection technique that can
be used to select a subset of training data that gives similar or superior performance to a model trained on full data.Paper AbstractIt is often advantageous to train models on a subset of the available train examples, because the examples are of variable quality or because one would like to train with fewer examples, without sacrificing performance. We present Gradient Information Optimization (GIO), a scalable, task-agnostic approach to this data selection problem that requires only a small set of (unlabeled) examples representing a target distribution. GIO begins from a natural, information-theoretic objective that is intractable in practice. Our contribution is in showing that it can be made highly scalable through a simple relaxation of the objective and a highly efficient implementation. In experiments with machine translation, spelling correction, and image recognition, we show that GIO delivers outstanding results with very small train sets. These findings are robust to different representation models and hyperparameters for GIO itself. GIO is task- and domain-agnostic and can be applied out-of-the-box to new datasets and domains.Features:GIO with quantization using K-means.Sentence embedding script to generate embeddings from data to use in GIOInstallationInstallable via pip:pipinstallgrad-info-optOr install directly form the repository:gitclonegit@github.com:daeveraert/gradient-information-optimization.gitcdgradient-information-optimization
pipinstall-e.Direct installation will require you to install additional dependencies listed below. We welcome contributions to GIO.Requirementsnumpy>=1.21.6jax>=0.3.25pyspark>=2.4.8sentence_transformers>=2.2.2jaxlib>=0.3.2pandas>=1.0.5Quick StartNote:GIO uses a Spark context, or if it can't find one, it will create a local one. You may encounter a Spark error before the algorithm runs complaining it cannot find a free port. In this case, executingexport SPARK_LOCAL_IP="127.0.0.1"should resolve the issue.Here is a simple 2D demonstration of how to use GIO with visualization:fromGIOimportGIOKLimportnumpyasnpimportjax.numpyasjnpimportmatplotlib.pyplotasplt# Create some datadefgetX():mean=[3,4]cov=[[0.5,0],[0,0.5]]np.random.seed(1)x,y=np.random.multivariate_normal(mean,cov,100).Treturnjnp.array([[x[i],y[i]]foriinrange(len(x))])defgetXTest():mean=[3,4]cov=[[0.5,0],[0,0.5]]np.random.seed(5)x,y=np.random.multivariate_normal(mean,cov,100).Treturnjnp.array([[x[i],y[i]]foriinrange(len(x))])X=getX()X_test=getXTest()# Initialize classgio_kl=GIOKL.GIOKL(uniform_low=0,uniform_high=8,uniform_start_size=100,dim=2)# Perform the AlgorithmW,kl_divs,_=gio_kl.fit(X_test,X,normalize=False)W=W[100:]# Remove the uniform start# Plot resultsplt.plot(kl_divs)plt.title("KL Divergence vs. Iterations")plt.xlabel("Iterations")plt.ylabel("KL Divergence")plt.show()plt.clf()plt.scatter([each[0]foreachinW],[each[1]foreachinW],label='Selected Data')plt.scatter([each[0]foreachinX],[each[1]foreachinX],label='Target Data')plt.title("Target Data and Selected Data")plt.xlabel("Dimension 1")plt.ylabel("Dimension 2")plt.legend()plt.show()Here is a more complex example for scale applications, reading and using a CSV that stores embeddings and data, using quantization-explosion, and Spark:fromGIOimportGIOKLimportjax.numpyasjnpimportmatplotlib.pyplotaspltimportpyspark.sql.functionsasF# Initialize classgio_kl=GIOKL.GIOKL(uniform_low=-1,uniform_high=1,uniform_start_size=20,dim=768)# Read datatrain_df,target_df=gio_kl.read_data_from_csv(PATH_TO_TRAIN,PATH_TO_TARGET)# Quantize datamodel_train,model_X,transformed_train,transformed_X=gio_kl.quantize(train_df,target_df)X=jnp.array(model_X.clusterCenters())train=jnp.array(model_train.clusterCenters())centroids_df=gio_kl.spark.createDataFrame(data=[(i,each.tolist())fori,eachinenumerate(model_train.clusterCenters())],schema=["id","centroid"])# Perform the AlgorithmW,kl_divs,_=gio_kl.fit(train,X,max_iter=300,stopping_criterion='sequential_increase_tolerance',v_init='jump')W=W[20:]# Remove the uniform start# Explode back to original data and write resulting datafull_selections_df=gio_kl.explode(W,transformed_train,centroids_df)full_selections_df.select(F.col("_c0"),F.col("_c1")).write.option("delimiter","\t").csv(OUTPUT_PATH)# Plot resultsplt.plot(kl_divs)plt.title("KL Divergence vs. Iterations")plt.xlabel("Iterations")plt.ylabel("KL Divergence")plt.show()Note:For quantization, Spark requires a large rpc message size. It is recommended to placegio_kl.spark.conf.set("spark.rpc.message.maxSize", "500")(or any large number) in the code before calling quantize, if the defaults haven't already been increased.Available OptionsGIOKL.fittakes the following arguments:train: training data as a jnp array (jnp is almost identical to numpy) [M, D] shapeX: target data as a jnp array [N, D] shapeD: initial data as a jnp array, default None. Use None to initialize from 0 (uniform) or a subset of training datak: kth nearest neighbor to use in the KL divergence estimation, default 5max_iter: maximum iterations for the algorithm. One iteration adds one point (cluster)stop_criterion: a string for the stopping criterion, one of the following: 'increase', 'max_resets', 'min_difference', 'sequential_increase_tolerance', 'min_kl', 'data_size'. Default is 'increase'min_difference: the minimum difference between prior and current KL divergence for 'min_difference' stop criterion only. Default is 0resets_allowed: whether if KL divergence increases, resetting G to the full train is allowed (allows the algorithm to pick duplicates). Must be set to true if the stop criterion is 'max_resets'. Default is Falsemax_resets: the number of resets allowed for the 'max_resets' stop criterion only (a reset resets G to the full train set and allows the algorithm to pick duplicates). Default is 2max_data_size: the maximum size of data to be selected for the 'data_size' stop criterion only, as a percentage (of total data) between 0 and 1. Default is 1min_kl: the minimum kl divergence for the 'min_kl' stop criterion only. Default is 0max_sequential_increases: the maximum number of sequential KL divergence increases for the 'sequential_increase_tolerance' stop criterion only. Default is 3random_init_pct: the percent of training data to initialize the algorithm from. Default is 0random_restart_prob: probability at any given iteration to extend the gradient descent iterations by 3x, to find potentially better extrema. Higher values come at the cost of efficiency. Default is 0scale_factor: factor to scale the gradient by, or 'auto'. Default is 'auto', which is recommendedv_init: how to initialize v in gradients descent, one of the following: 'mean', 'prev_opt', 'jump'. Default is 'mean'grad_desc_iter: the number of iterations to use in gradient descent. Default is 50discard_nearest_for_xy: discard nearest in the xy calculation of KL divergence, for use when X and the train set are the same, comes at the cost of efficiency. Default is Falselr: Learning rate for gradient descent. Default is 0.01Citing GIOIf you use GIO in a publication, blog or software project, please cite the paper:@misc{everaert2023gio,
      title={GIO: Gradient Information Optimization for Training Dataset Selection}, 
      author={Dante Everaert and Christopher Potts},
      year={2023},
      eprint={2306.11670},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}