Many machine learning tasks require sampling a subset of items from a collection based on a parameterized distribution. The Gumbel-softmax trick can be used to sample a single item, and allows for low-variance reparameterized gradients with respect to the parameters of the underlying distribution. However, stochastic optimization involving subset sampling is typically not reparameterizable. To overcome this limitation, we define a continuous relaxation of subset sampling that provides reparameterization gradients by generalizing the Gumbel-max trick. We use this approach to sample subsets of features in an instance-wise feature selection task for model interpretability, subsets of neighbors to implement a deep stochastic k-nearest neighbors model, and sub-sequences of neighbors to implement parametric t-SNE by directly comparing the identities of local neighbors. We improve performance in all these tasks by incorporating subset sampling in end-to-end training.
翻译:许多机器学习任务要求从一个基于参数分布的收藏中取样一个子项。 Gumbel- softmax 戏法可以用来对单个物品进行取样,并允许对基本分布参数进行低差的重新校准梯度。然而,涉及子抽样的随机优化通常无法重新校准。为了克服这一限制,我们定义了子抽样的连续放松,通过对 Gumbel-max 戏法进行概括化,提供再度梯度。我们用这个方法对实例特征选择任务中的特征组进行取样,用于模型解释性、邻居子集,以实施深孔式K-最近邻模型,以及邻居子序列,通过直接比较当地邻居的身份来实施参数 t-SNE。我们通过在终端到终端培训中纳入子取样,改进所有这些任务的绩效。