Topic-aware transformer with hierarchical prompting learning for multi-label image classification

Topic model

Topic models are a class of probabilistic models designed to uncover latent semantic structures within a collection of documents by representing each document as a mixture of hidden topics. In our work, a topic model is employed to capture the latent structure within the label combinations associated with each image. By utilizing a probabilistic topic modeling approach with a predefined number of topics L, we derive a $\mathbf{D}-\mathbf{L}$ matrix that represents the probability distribution of documents over the identified topics.

Vision transformer

For a standard ViT model, the input image $x_i\in {\mathbb{R}}^{C\times H\times W}$ is initially divided into N patches $x_p\in {\mathbb{R}}^{N\times \left(P^2\times C\right)}$. Here, H and W are the height and width of the input images, C is the number of channels, and P denotes the patch size. Each patch is flattened into a 1D vector and passed through a trainable linear projection layer, mapping it into a latent space of dimension $d$, with position encodings added. This process is represented as ${\mathbf{x}}_p=\mathrm{Embed}\left(x_p\right)$. A learnable class token $x_{cls}^0\in {\mathbb{R}}^d$ is then prepended to the patch embeddings as input to the transformer blocks. The $i$-th transformer block can be formulated as:

(1) $$\left[{\mathbf{x}}_{cls}^i,{\mathbf{X}}^i\right]=B_i\left(\left[{\mathbf{x}}_{cls}^{i-1},{\mathbf{X}}^{i-1}\right]\right),i=1,2,\dots ,B.$$

Here ${\mathbf{x}}_{cls}^i$, ${\mathbf{X}}^i$ represent the class token and patch embeddings as the outputs of the $i$-th transformer block. Finally, the final block’s class token ${\mathbf{x}}_{cls}^B$ is used to predict the probability of each label appearing in the image. The class token is passed into the classification head and activated by the sigmoid function, represented by the formula:

(2) $$r=Sigmoid\left(W\left(x_{cls}^B\right)\right)$$where $r\in {\mathbb{R}}^N$ represents predicted probabilities for N categories. $W\in {\mathbb{R}}^{d\times N}$ denotes the classification head parameters.

Multi-granularity topic information extraction

Multi-label datasets encapsulate not only profound visual content but also intricate semantic structures. Effectively extracting and leveraging these semantic relationships can enhance the model’s capacity to accurately perform multi-label classification. Topic models are particularly well-suited for this purpose, as they can uncover latent co-occurrence patterns and provide semantic representations from discrete label data. To this end, we employ a topic model to identify sample-level topic information, which serves as auxiliary signals to augment the transformer, thereby improving its classification performance.

Given a multi-label dataset consisting of M samples, we represent the combination of labels associated with each sample as a document $d_i$. The collection of all such documents from the training set forms a corpus $D=\left[d_1,d_2,...,d_M\right]$, which serves as the input to the topic model. By specifying the number of topics as L, the topic model learns a latent semantic structure and outputs a document-topic distribution matrix of size M $\times$ L. In this matrix, the $i$-th row represents the probability distribution of the document $d_i$ across the L topics. Since each document corresponds to a unique image sample, we interpret this distribution as the probability of the $i$-th image sample being associated with each of the L topics, thereby providing an enriched semantic representation for subsequent classification tasks. Furthermore, we assign each image to the topic with the highest probability. To further refine this representation, we extract topic distributions at multiple granularities by varying the number of topics L. Coarse-grained topic distributions (with smaller L) capture broader semantic structures, while fine-grained distributions (with larger L) provide more detailed label dependencies.

Topic-aware prompt pool

Topic models are capable of extracting latent semantic structures from multi-label datasets, offering valuable guidance for multi-label classification tasks. However, directly aligning these semantic topic representations with visual features remains a nontrivial challenge. To bridge this gap, we introduce an auxiliary task that incorporates a learnable prompt token into the transformer architecture. The prompt token is designed to capture topic-level semantics and facilitates the integration of external semantic knowledge into the model’s internal representation space. While the naive way of learning a single shared prompt for multi-granularity topic structures is parameter efficient and promotes the learning of global semantic representation, it still suffers from severe supervision entanglement and insufficient capacity to capture the semantic distinctions across different levels of granularity. Therefore, we adopt a prompt pool to store topic information, where each prompt is responsible for modeling topic information at a specific level of granularity. The prompt pool is denoted as $P=\left[p^1,p^2,\dots ,p^K\right]$, where $p^k\in {\mathbb{R}}^d$, $d$ represents the embedding size and K is the number of prompt tokens as well as the number of prompt blocks. Correspondingly, for each prompt token $p^k$, we define a learnable linear projection layer $t^k:{\mathbb{R}}^d⟶{\mathbb{R}}^{C_k}$ serving as a classification head mapping the learned prompt token to one of the $C_k$ topic categories associated with the image. The collection of these classifiers is denoted as $T=\left[t^1,t^2,\dots ,t^K\right]$.

Hierarchical prompt of multi-granularity topic information

Recent advances have demonstrated the effectiveness of prompt-based tuning in adapting pre-trained Transformer models to downstream tasks. In this paradigm, a prompt token is prepended to the input sequence, enabling task-specific conditioning while keeping the backbone frozen. Accordingly, the computation of the transformer encoder turns into:

(3) $$\left[{\mathbf{x}}_{cls}^B,{\mathbf{X}}^B,\mathbf{z}\right]=B\left(\left[{\mathbf{x}}_{cls},\mathbf{X},\mathbf{p}\right]\right).$$

Motivated by this idea, we adopt ViT as the backbone model due to its self-attention mechanism, which not only enables the class token to absorb the semantic information injected by the prompt token, thereby enhancing the performance of the multi-label classification task, but also allows the prompt token to attend to visual features from the image, further optimizing the auxiliary task of topic classification. Moreover, we propose a hierarchical prompt strategy that integrates topic information at different levels of granularity into selected transformer blocks, thereby effectively incorporating the semantic structures captured by the topic model. According to the coarse-to-fine multi-level topic hierarchy, our proposed TATHPL first selects multiple transformer blocks and then reshapes them into prompt blocks. Specifically, for each prompt block, a corresponding prompt token that encodes topic information at a particular level of granularity is inserted into the input sequence of that layer, which can be expressed as

(4) $$\left[{\mathbf{x}}_{cls}^i,{\mathbf{X}}^i,{\mathbf{z}}^k\right]=B_i\left(\left[{\mathbf{x}}_{cls}^{i-1},{\mathbf{X}}^{i-1},{\mathbf{p}}^k\right]\right)i=1,2,\dots ,B,k=1,2,\dots ,K,k\le i.$$

Here, ${\mathbf{p}}^k$ represents the $k$-th prompt, which indicates that the $i$-th block is the $k$-th selected prompt block. To clarify the hierarchical prompt injection mechanism, we provide an example based on the Corel5k dataset. In this setting, we insert prompts into the first and last blocks of the Transformer encoder, which consists of 12 blocks in total. The corresponding insertion operations at these two blocks are formulated as follows:

(5) $$\left[{\mathbf{x}}_{cls}^1,{\mathbf{X}}^1,{\mathbf{z}}^1\right]=B_1\left(\left[{\mathbf{x}}_{cls}^0,{\mathbf{X}}^0,{\mathbf{p}}^1\right]\right)$$

(6) $$\left[{\mathbf{x}}_{cls}^{12},{\mathbf{X}}^{12},{\mathbf{z}}^2\right]=B_{12}\left(\left[{\mathbf{x}}_{cls}^{11},{\mathbf{X}}^{11},{\mathbf{p}}^2\right]\right).$$

After the computation by the prompt block, the prompt ${\mathbf{z}}^k$ is no longer passed to the subsequent block. Instead, it is used directly for the prediction of topic categories. To derive the similarity score, we project the prompt ${\mathbf{z}}^k$ into a logit value through the linear projection layer ${\mathbf{t}}^k$ followed by a softmax function. Finally, we compute the loss for the topic classification task generated by the $k$-th prompt layer using the cross-entropy loss function as follows:

(7) $${\mathcal{L}}_{\mathtt{t}\mathtt{o}\mathtt{p}\mathtt{i}\mathtt{c}}^k=-log\frac{\exp \left(Softmax\left({\mathbf{z}}_y{\mathbf{t}}_y\right)\right)}{{\sum }_i\exp \left(Softmax\left({\mathbf{z}}_i{\mathbf{t}}_i\right)\right)}$$where $y$ represents that the input image belongs to the $y$-th topic. For simplicity, we omit the superscript $k$. Correspondingly, the overall loss function of the topic classification task is formulated as follows:

(8) $${\mathcal{L}}_{topic}=\sum _k{\lambda }_k{\mathcal{L}}_{topic}^k$$where ${\lambda }_k$ is the balance parameter that controls the relative contribution of the $k$-th prompt branch to the total loss. The balance parameters are introduced to facilitate stable learning across multiple levels of topic granularity, mitigating the risk of overemphasis on any individual prompt-level supervision during training.

Multi-label prediction

After the hierarchical integration of multi-granularity topic prompts into the transformer encoder, we obtain the final class token ${\mathbf{x}}_{cls}^B$ which encodes both enriched visual representations and external topic-aware semantics. Then we treat each label prediction as a binary classification task. Specifically, the final class token ${\mathbf{x}}_{cls}^B$ is projected into a logit using a linear projection layer, followed by a sigmoid activation to estimate the confidence score for the presence of this label:

(9) $$r_j=Sigmoid\left(W_j\left(x_{cls}^B\right)\right)$$where $r_j\in R^d$, $r=\left[r_1,\dots ,r_N\right]\in {\mathbb{R}}^{d\times N}$ denotes the confidence score for the presence of each label, $W_j\in {\mathbb{R}}^d$, $W=\left[W_1,\dots ,W_N\right]\in {\mathbb{R}}^{d\times N}$ represents the parameters of the linear projection layer.

Performance on MS-COCO

MS-COCO is a large-scale image dataset designed for object detection, image segmentation, and multi-label classification. It contains 122,218 images, with 82,081 images in the training set and 40,137 images in the test set used for model evaluation. The dataset includes 80 object categories commonly found in scenes, with each image annotated with an average of 2.9 labels. This dataset datasets can be found at https://cocodataset.org. On the MSCOCO dataset, we select the first and second blocks as our prompt blocks and insert topic prompts containing 2 and 6 topics, respectively. The balance parameter $\lambda$ is set to $0.15$ for all experiments.

As shown in Table 2, the results indicate that our model achieves an mAP score of $81.9\mathrm{\%}$ at a resolution of $224\times 224$, outperforming the sub-optimal model Spatial Regularization Network (SRN). Furthermore, it demonstrates improvements in other evaluation metrics, such as CF1 and OF1, indicating that our model not only improves overall classification accuracy but also balances precision and recall more effectively, demonstrating stronger robustness in multi-label classification. The key reason our model achieves these results lies in its ability to effectively capture the latent semantic structure of label combinations, referred to as “topic information.” By introducing hierarchical prompt learning, our ViT model learns label correlations at multiple granularities. In contrast, traditional methods often overlook this aspect. Specifically, our model guides attention to image regions aligned with topic information at different levels, from coarse-grained global semantics to fine-grained task-relevant areas, ensuring that the model effectively understands and distinguishes the latent relationships between labels at every level.

Performance on NUS-WIDE

NUS-WIDE is a large-scale real-world web-based image dataset containing 269,648 images and 81 visual concepts. We trained our model on 161,789 images and evaluated it on 107,859 images. This dataset datasets can be found at https://www.kaggle.com/datasets/7cbbf047bc9c47b4f2c00e83531d3376ab8887bb0deed2ce2ee1596fe96aa94d?select=NUSWIDE. We insert a topic prompt containing 2 topics into the last block of the model and set the balance parameter $\lambda$ to $1$. The experimental results are shown in Table 3. We used three metrics—mAP, CF1, and OF1—to evaluate the performance of our model and compared it with five previously proposed state-of-the-art methods. In all experiments, the image resolution was $224\times 224$. Our model outperforms other state-of-the-art and baseline models.

Performance on Corel5k

Corel5k is a widely used multi-label dataset consisting of 4,999 images, with 4,500 images designated for the training set and 499 images for the validation set. The dataset is annotated with 260 categories. This dataset can be found at https://github.com/watersink/Corel5K. In this dataset, we select the first and last blocks as the prompt blocks and insert topic prompts with topic numbers of 2 and 3, respectively. And both of the balance parameters $\lambda$ are set to be $0.1$. As shown in Table 4, our method demonstrates superior performance compared to other state-of-the-art methods and the baseline, which exhibits strong transfer learning capabilities.

Comparison with baseline methods

To prove the efficiency of the refinements put forward in our work, we compare the TATHPL with the baseline model ViT on MS-COCO. The results, presented in Table 5, show that our model consistently outperforms the baseline under various evaluation metrics. As expected, topic prompts facilitate a better understanding of label dependencies.

Effect of model components

Our proposed framework incorporates two essential components: topic extraction and the hierarchical prompt. In this subsection, we conduct an ablation study to evaluate the contribution of each module. The results are shown in Table 6. Compared with the baseline backbone network, introducing the Hierarchical Prompt alone yields a slight improvement in performance, increasing the mAP from $62.4\mathrm{\%}$ to $62.5\mathrm{\%}$ on the Corel5k dataset. When the Topic Extraction module is further incorporated, the performance is significantly enhanced, achieving an mAP of $63.5\mathrm{\%}$. While the Hierarchical Prompt alone contributes modest gains, the more substantial improvement observed with the addition of Topic Extraction indicates that the extracted topic information plays a valuable role in boosting multi-label classification performance.

Statistical significance analysis

We conducted an ablation study and statistical significance analysis to compare the performance of our proposed model with the baseline model ViT on the Corel5k dataset. Specifically, for each test sample, we generated multi-label classification probability predictions from both models. We then compared the differences between the two models using the paired t-test and Wilcoxon Signed-Rank Test. The paired t-test yielded a p-value of 0.028, while the Wilcoxon Signed-Rank Test produced a p-value of 0.022. Both p-values are below the significance threshold of 0.05, indicating that the performance differences between the two models are statistically significant.

The results show that our model outperforms the baseline model, confirming the effectiveness and advantages of our approach.

Effect of the number of topics L

For the topic extraction module, the number of topics L is a critical hyperparameter. We investigate its impact by injecting topic information with varying values of L (ranging from 2 to 8) into both the first and the last encoder layers of a standard ViT-B/16 model. The results are illustrated in Fig. 3. As shown, when topics are injected into the first layer, the best performance is achieved with $L=2$, while increasing L gradually leads to performance degradation. Based on this observation, we typically choose $L=2$ for shallow layers. In contrast, when topic information is injected into the final encoder layer, smaller values of L result in suboptimal performance. When $L=8$, the performance improved to some extent. These results suggest that as the number of topics increases, the topic prompts should be injected into deeper layers of the transformer for optimal performance.

Mean average precision (mAP)

For a single class, given N samples sorted by prediction scores in descending order, the average precision (AP) is defined as:

(13) $$AP=\frac{1}{N_{+}}\sum _{k=1}^{N_{+}}P@k$$where: $N_{+}$ is the total number of positive samples in the class and $P@k$ is the precision at the position of the $k$-th positive sample in the ranked list AP measures the average precision at all positions where a positive sample is found, providing a single score that combines both ranking quality and classification accuracy. For predictions with C classes, the mean average precision is computed as:

$$\mathrm{m}\mathrm{A}\mathrm{P}=\frac{1}{C}\sum _{k=1}^{C}AP(k)\times 100\mathrm{\%}.$$

Overall metrics

(20) $$\mathrm{O}\mathrm{P}=\frac{{\sum }_{i=1}^{N}T{P}_i}{{\sum }_{i=1}^N(T{P}_i+F{P}_i)}$$

(21) $$\mathrm{O}\mathrm{R}=\frac{{\sum }_{i=1}^{N}T{P}_i}{{\sum }_{i=1}^N(T{P}_i+F{N}_i)}$$

(22) $$\mathrm{O}\mathrm{F}1=\frac{2\cdot \mathrm{O}\mathrm{P}\cdot \mathrm{O}\mathrm{R}}{\mathrm{O}\mathrm{P}+\mathrm{O}\mathrm{R}}.$$