Sharpness-Aware Minimization for Topic Models with High-Quality Document Representations (2025)

Abstract: Recent advanced frameworks in topic models have significantly enhanced the performance compared to conventional probabilistic approaches. Such models, mostly constructed from neural network architecture together with other advanced techniques such as contextual embedding, optimal transport distance and pre-trained language model, etc. have effectively improved the topic quality and document topic distribution. Despite the improvements, these methods lack considerations of effective optimization for complex objective functions that contain log-likelihood and additional regularization terms. In this study, we propose to apply an efficient optimization method to improve the generalization and performance of topic models. Our approach explicitly considers the sharpness of the loss landscape during optimization, which forces the optimizer to choose directions in the parameter space that lead to flatter minima, in which the models are typically more stable and robust to small perturbations in the data. Additionally, we propose an effective strategy to select the flatness region for parameter optimization by leveraging the optimal transport distance between doc-topic distributions and doc-cluster proportions, which can effectively enhance document representation. Experimental results on popular benchmark datasets demonstrate that our method effectively improves the performance of baseline topic models.

## introduction topic models (tms) @xcite @xcite are designed to uncover hidden topic structures within a corpus while also providing topic distributions for individual documents. topic models are utilized across multiple fields in natural * these authors contributed equally to this work. † corresponding author: linhnv@soict.hust.edu.vn language processing @xcite @xcite @xcite . in recent years, several advanced topic models @xcite @xcite @xcite , mostly based on neural networks, have emerged. these models not only enable efficient and flexible parameter inference through automatic gradient back-propagation but also improve the quality of topic-word distributions and document representations. in addition to the application of neural networks in topic modeling, several advanced techniques have been introduced to enhance model performance, such as integrating richer contextual information @xcite @xcite , leveraging contrastive learning strategies @xcite , and applying optimal transport methods @xcite @xcite , among others. while modern neural topic models have successfully improved both the quality of discovered topics and the distribution of topics across documents, they have largely overlooked the issue of model optimization. most recent models @xcite @xcite are built on the variational autoencoder (vae) framework @xcite , which relies on maximizing the log likelihood and regularizing the document-topic distribution. to further enhance model performance, additional objective constraints are often introduced, such as boosting topic diversity through embedding clustering regularization @xcite or improving document representation by maximizing mutual information @xcite . although these techniques result in a more complex final objective function, the training process remains relatively straightforward, using standard gradient back-propagation. in this paper, we apply an effective optimization for topic models that not only minimizes the loss function at specific model parameters, but also enhances the model's robustness to variations in those parameters within a local neighborhood, leading to flatter minima. the relationship between the flatness of minima and generalization has been widely studied from both theoretical and empirical perspectives @xcite @xcite . such sharpness-aware minimization techniques have shown promising results in other areas of machine learning, such as image classification, transfer learning, fine-tuning, and language modeling @xcite @xcite . however, despite the complexity of objective functions in topic models, often involving additional regularization terms @xcite @xcite , there has been little focus on improving optimization methods to boost performance in this domain. our proposed method, sharpness-aware minimization for topic modeling, can be seamlessly integrated into a variety of topic models, significantly improving both generalization and performance in terms of topic quality and document representations. specifically, we introduce a novel local neighborhood region for sharpness-aware minimization that effectively enhances the inference networks of topic models. our method leverages the high clustering accuracy of pre-trained language models @xcite . in detail, we use the optimal transport (ot) distance @xcite between the document-topic distributions from the topic model and the document-cluster proportions from pre-trained clustering to inform the sharpnessaware neighborhood. this approach introduces an awareness of regions with strong document representations during optimization, leading to flatter minima and improved inference performance. we call this method as dream which enhances document representations via sharpness-aware minimization. we summarize the contributions of our study as follows: ## related work topic modeling aims to uncover hidden topics within a corpus of documents. traditionally, this problem has been addressed using graphical probabilistic methods @xcite . beyond these standard approaches, various studies have extended topic modeling to specialized contexts, such as short texts @xcite @xcite and streaming environments @xcite @xcite @xcite . more recently, the focus has shifted toward neural network-based models, which have demonstrated superior generalization and higher performance @xcite @xcite @xcite . most neural topic models are built upon the vae architecture @xcite . in this framework, the encoder (inference) network generates the document's topic distribution, while the decoder (generative) network combines these topic proportions with the topic-word distribution to reconstruct the original data. some approaches have focused on enhancing the inference network by incorporating document embeddings from pretrained language models (plms) @xcite as input @xcite or by imposing additional constraints using plm representations in the objective function @xcite . on the other hand, improvements to the generative process have been made through the use of word embeddings @xcite @xcite , applying conditional transport @xcite , and leveraging optimal transport distance @xcite . another approach in neural topic modeling involves generating topics by clustering document representations directly @xcite @xcite . this method is efficient and yields coherent topics, but determining the topic proportions within a document is not straightforward. additionally, recent research leverages large language models to generate topics as conceptual descriptions @xcite , though these methods struggle to provide word distributions within topics or topic proportions within documents. specifically, @xcite propose a novel topic modeling approach based solely on optimal transport @xcite , capturing the semantic relationships among documents, topics, and word embeddings. ## conclusion in conclusion, this paper presents a novel approach to enhancing topic model performance through an optimization strategy that minimizes both loss value and sharpness. specifically, our proposed optimization, namely dream, conducts sharpnessaware minimization with a constraint with highquality document representations. extensive experiments on benchmark datasets show significant improvements of dream in topic quality and document-topic distribution across various topic models.