提出高效训练软最大化混合模型的新方法,解决高维场景下参数识别难题。
Learning large softmax mixtures with warm start EM
- 基于新潜层矩估计法构造算法初始值,提升收敛效率
- 在高维下实现近参数级的混合原子恢复精度
- 适合处理大数量类别、高维特征的大型语言模型优化
软最大化混合模型(SMMs)是用于建模从 $p$ 个候选项中选择属性 $x_j \in \RR^L$ 的离散 $K$-混合模型,广泛应用于异质人群建模与大语言模型末层。本文对高维下 SMM 的 EM 算法进行系统分析,首次证明:在通用特征下具局部可识别性;在随机特征且 $p$ 足够大时具完全可识别性($L > 1$)。理论揭示了算法收敛所需的初始化半径,指导样本级温启动构造。通过新提出的混合测度矩估计方法,构建了方法矩(MoM)估计器,其可实现参数级收敛速度,但质量随 $K$ 指数下降。当 $K$ 较大时,推荐多次随机初始化运行 EM。利用该矩估计法还可估计混合原子的 $K$-维子空间,显著减少采样次数。
原文摘要 · Abstract (English)
Softmax mixture models (SMMs) are discrete $K$-mixtures introduced to model the probability of choosing an attribute $x_j \in \RR^L$ from $p$ candidates, in heterogeneous populations. They have been known as mixed multinomial logits in the econometrics literature, and are gaining traction in the LLM literature, where single softmax models are routinely used in the final layer of a neural network. This paper provides a comprehensive analysis of the EM algorithm for SMMs in high dimensions. Its population-level theoretical analysis forms the basis for proving (i) local identifiability, in SSMs with generic features and, further, via a stochastic argument, (ii) full identifiability in SSMs with random features, when $p$ is large enough. These are the first results in this direction for SSMs with $L > 1$. The population-level EM analysis characterizes the initialization radius for algorithmic convergence. This also guides the construction of warm starts of the sample level EM. Under suitable initialization, the EM algorithm is shown to recover the mixture atoms of the SSM at near-parametric rate. We provide two main directions for warm start construction, both based on a new method for estimating the moments of the mixing measure underlying an SSM with random design. First, we construct a method of moments (MoM) estimator of the mixture parameters, and provide its first theoretical analysis. While MoM can enjoy parametric rates of convergence, and thus can serve as a warm-start, the estimator's quality degrades exponentially in $K$. Our recommendation, when $K$ is not small, is to run the EM algorithm several times with random initializations. We again make use of the novel latent moments estimation method to estimate the $K$-dimensional subspace of the mixture atoms. Sampling from this subspace reduces substantially the number of required draws.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。