Category-aware EEG Image Generation Based on Wavelet Transform and Contrast Semantic Loss

1 Introduction

In recent years, the reconstruction of visual stimuli from electroencephalogram (EEG) has emerged as a highly promising research area within the domain of brain-computer interface (BCI), by extracting visually relevant features from EEG and ultimately reconstructing visual stimuli [Bai et al., 2023; Kavasidis et al., 2017]. This technology has the ability to convert neural signals into images, thereby establishing a crucial link between brain activities and the external world and deepening our comprehension of the intricate relationship between brain activities and perception. It holds immense potential, especially for individuals with severe disabilities. By enabling them to convey their thoughts and intentions through visual representations, it has the potential to revolutionize assistive communication methods.

∗Corresponding Author

With the powerful image-generation capabilities of generative networks, such as adversarial generative networks (GAN) [Goodfellow et al., 2014], variational auto-encoder networks (VAE) [Kingma, 2013], and denoising diffusion probabilistic models (DDPM) [Zhang et al., 2025; Ho et al., 2020], which make it possible to perceive brain visual stimuli from EEG and reconstruct visual stimuli. However, existing related methods face substantial challenges. First, reconstructing visual stimuli from EEG at the pixel level is difficult and unnecessary. To this end, reconstructing semantically consistent images is of significant importance in BCI, leading to the inability to utilize traditional objective image evaluation indicators, such as structural similarity (SSIM). Therefore, the need arises for quantitatively evaluating the quality of visual stimuli to be reconstructed using an objective semantic-based score. In addition, EEG are time-series and noisy, which complicates the extraction of visually relevant features from them. Moreover, the training of advanced models, such as DreamDiffusion [Bai et al., 2023] and MinDvis [Chen et al., 2023], requires an excessive amount of computational resources, severely restricting their widespread application.

To process noisy time-series EEG signals efficiently, the integration of traditional signal analysis methods, such as the discrete wavelet transform (DWT) [Chen et al., 2017], into deep-learning modules has been employed [Zeng et al., 2024]. By incorporating DWT-based modules into EEG encoder, we can leverage both the spatial and frequency characteristics of EEG, enhancing the feature-extraction process. In addition, with the continuous evolution of deep-learning models, the gate mechanisms in Mamba [Gu and Dao, 2023] have been proven effective in neural networks for selectively controlling information flow.

Motivated by the success of gated attention mechanisms in Mamba and DWT in signal processing, in this work, we integrated a DWT module with the gated attention mechanism within a well-designed EEG encoder, extracting meaningful features from EEG by effectively capturing both spatial and frequency-domain information while selectively focusing on relevant features. To reconstruct semantically consistent images without pixel-level ground truth, a category-aware clustering loss is utilized to cluster samples of the same category and separate different categories in a high-dimensional space, thereby improving the reconstruction accuracy and zero-shot generalization ability of the model. After the EEG feature extraction, we employ a pre-trained diffusion-based model to generate high-quality images. To ensure that the reconstructed images accurately reflect the semantic meaning of the input EEG, we incorporate an image classification model as a performance validation step. Subsequently, we propose a related evaluation metric of semantic-based score, enabling a quantitative assessment of the semantic consistency between the input EEG and the reconstructed images. The novelty of this work is three-fold:

1. To the best of our knowledge, it is the first time to use category-aware clustering loss to better extract the EEG features associated with visual stimuli, with the assistance of the CLIP loss [Radford et al., 2021] to align image and EEG features.
2. We propose a novel semantic-based evaluation metric using pre-trained classifiers to quantitatively assess EEG-to-image reconstruction performance.
3. To address EEG's noisy time-series nature, we design a specialized encoder combining DWT and gated attention, achieving strong classification performance and supporting downstream image generation.

2 RelatedWork

Reconstructing visual stimuli from brain signals, such as functional magnetic resonance imaging(fMRI) and EEG, has achieved remarkable results in previous studies. For instance, conditional Generative Adversarial Networks and VAEs have been applied to encode fMRI signals and further reconstruct the visual stimuli from the features of brain signals. Studies such as Shen et al. [Shen et al., 2019] and Takagi et al. [Takagi and Nishimoto, 2023] demonstrated that fMRI data could be translated into semantically correct and high-quality images. In addition, EEG-based reconstruction algorithms, despite the challenges of noise and lower spatial resolution, have also exhibited promising results. Techniques integrating convolutional neural networks with GANs, as in [Song et al., 2021; Yang et al., 2021], have enabled the generation of visual representations corresponding to real-time brain activity. To simplify the application of BCI and make the technology more low-cost and convenient for use, EEG is a more practical signal entry than fMRI. Recent outstanding results, such as DreamDiffusion [Bai et al., 2023] and ATM-S [Li et al., 2024], have shown that, with the assistance of the powerful image generation ability of diffusion model, deep learning models can reconstruct visual stimuli from EEG, capturing some semantic information embedded in neural activity. Moreover, multi-modal approaches combining EEG and fMRI have enhanced image quality and robustness by leveraging complementary features from both modalities. These advancements emphasize the growing potential of brain-signal-based image generation for applications.

3 Methods

3.1 Overall Architecture

Fig. 1 depicts the overall flowchart of the proposed model. In this study, the proposed model comprises three main parts: EEG embedding with a well-designed EEG encoder seen in Fig. 2, the part of the downstream task (image reconstruction and EEG classification), and the semantic evaluation. Within the EEG embedding part, several well-designed modules are employed, including the DWT block, gated attention mechanisms, and a feature-fusion module. During network training, to achieve optimal performance in terms of feature alignment and category guidance, three distinct losses are combined: the CLIP loss, the Mean Square Error (MSE) loss, and the category-aware clustering loss. In the image reconstruction of the downstream task branch, the approach adopted is similar to that of the ATM-S method [Li et al., 2024].

3.2 DWT Block

To effectively capture both the temporal- and frequency- domain characteristics of EEG, we propose a DWT-based feature extraction module utilizing the PyTorch Wavelet package [Cotter, 2020]. In this block, the EEG embedding input can be written as E ∈ RB×C×T , where B represents the batch size, C denotes the number of channels and T is the length of the signal. Using DWT module, the model can synthesize both frequency and spatial features at different stages, enabling the model to learn more complex features.

Channel-wise DWT Decomposition: Given the timeseries characteristics and high noise levels inherent in EEG, in this study, the Daubechies-1 (db1) wavelet known as the Haar wavelet, with a single-level decomposition is used for the implementation of our DWT module. The rationale behind this choice lies in the optimal localization properties of the db1 wavelet within the time domain, demonstrating remarkable effectiveness in capturing abrupt signal changes [Ocak, 2009]. Furthermore, single-level decomposition strikes an appropriate balance between frequency resolution and computational efficiency, which can balance the accuracy and computational demands for EEG processing. First, we apply one-dimensional (1D) DWT decomposition to each EEG channel independently:

where E{b,c,:} denotes the temporal sequence for batch b and channel c. cAc1 ∈ RT/2 denotes the approximation coefficients of the low-frequency components. which can capture the overall trends and low-frequency patterns of the EEG. Concurrently, cDc1 ∈ RT/2 represents the detail coefficients of the high-frequency components, which retain the rapid changes and high-frequency characteristics of the signals.

Inverse DWT Module: We employ a 1D inverse DWT module (iDWT) to map the proposed EEG features back to the temporal domain according to the following equation:

where ψ represents the wavelet reconstruction parameters. This reconstruction process preserves not only the learned frequency-domain patterns but also temporal coherence. To integrate both the original temporal features and the optimized ones, we propose a feature fusion module, expressed mathematically as:

where Wf and bf are learnable parameters. F represents the proposed features derived from [cAc1 , cDc1 ], and Conv denotes the convolution operations.

3.3 Gated AttentionMechanism

To process EEG embeddings from DWT modules while maintaining channel characteristics, we propose a gateenhanced multi-head attention mechanism. This approach integrates the traditional gating mechanism within the attention of the transformer for selective feature processing. Our implementation features an adaptive gating integration mechanism for EEG embeddings, bridging local convolutional and global transformer features.

Local Feature Integration: The local feature extraction process begins with channel-wise temporal convolution, which processes each EEG channel independently, thus preserving their temporal relationships and individual patterns. Subsequently, we applied a cross-channel spatial convolution, integrating information from neighboring electrodes to effectively capture the spatial dependencies within the EEG structure. The integration of cross-channel features is realized through a comprehensive spatial processing pipeline that operates on wavelet-transformed signals. By applying two-dimensional convolution operations across the channel dimension of the wavelet coefficients, we effectively capture the intricate spatial relationships among different EEG channels within the frequency domain. This spatial integration process concatenates the wavelet coefficients from all channels before applying a spatial convolution kernel, enabling the model to learn meaningful cross-channel patterns. The convolution output then goes through batch normalization and non-linear activation, yielding a refined feature representation that simultaneously maintains the channel-specific frequency characteristics from the wavelet transform and the spatial relationships between channels.

Global Frequency Domain Attention: For global feature processing, we utilize a transformer-based encoder that employs multi-head self-attention mechanisms to capture longrange dependencies throughout the entire sequence length while preserving the channel-specific information structure. The integrated spatiotemporal features are then processed by a specialized channel-wise attention mechanism module. By applying a dedicated attention operation to each channel's spatiotemporal components, this module learns to selectively emphasize the most relevant patterns while simultaneously maintaining the unique characteristics of individual EEG channels. The use of multi-head attention assigns adaptive weights to different components based on their importance, enabling the model to focus on discriminative spectral features essential for downstream tasks while suppressing less informative frequency content.

Dynamic Weighted Projection: The integration process constructs a dynamic weighted combination model, where the local convolution features and the global transformer features are balanced based on the learned channel-specific weights. This method allows the model to automatically adjust the contribution ratio of different feature types according to the characteristics of the input signal and the specific requirements of each channel. As intelligent selectors, gate mechanisms excel in EEG processing by learning to balance local temporal patterns with global dependencies. The combination of local patterns and global dependencies has proven particularly valuable in different applications, including EEG processing [Song et al., 2022; Lai et al., 2023], since the functional role of different channels, determined by their spatial location and measured brain activity, often requires different processing strategies. Therefore, the relationship between local patterns and global signal dependence can vary considerably depending on different regions of the brain and cognitive states [Song et al., 2022; Du et al., 2023]. The adaptive specialty of our proposed gate mechanism ensures that the model can handle these changes while preserving the integrity of the signal. In addition, the combination of gating mechanisms with multi-head attention mechanisms offers a robust framework for EEG processing, enabling the model to concentrate on the most relevant patterns while maintaining computational efficiency.

3.4 Multi-Branch Feature ExtractionModule

Due to the proposed multi-branch architecture, we effectively extract and fuse temporal and spatial features from EEG via parallel processing streams, capturing both timedomain characteristics and inter-channel relationships present in EEG. The temporal branch focuses on extracting timedomain characteristics from EEG by employing a series of well-designed convolution operations based on ShallowConvNet [Schirrmeister et al., 2017]. First, a temporal convolution is applied across the time dimension while processing each channel independently, thereby capturing local temporal patterns during the task. Subsequently, the temporal features undergo dimensionality reduction through average pooling, and then a point-wise convolution is performed, which adjusts the feature representation while maintaining temporal relationships.

The spatial branch is designed to capture inter-channel relationships and spatial patterns across the EEG electrode array. We implement this by initially performing a spatial convolution across the channel dimension, enabling the network to learn local spatial patterns between neighboring electrodes. This is followed by depth-wise separable convolutions, which efficiently expand the receptive field while maintaining computational efficiency.

The fusion mechanism employs an attention-based approach to adaptively combine temporal and spatial features. After concatenating features from both branches, we apply a channel attention mechanism that generates dynamic weights for each feature channel. The attention module consists of channel-wise average pooling followed by two convolution layers with a bottleneck structure, producing attention weights through a sigmoid activation. The final stage of our architecture integrates the attention-weighted features through a fusion branch. A point-wise convolution combines the weighted features, followed by batch normalization and nonlinear activation. The output then undergoes a final adaptive pooling operation to ensure consistent dimensionality.

3.5 Loss Function

In our approach, we introduce a dual-loss mechanism that integrates three complementary components: CLIP loss and MSE loss for general feature alignment, and a category-aware clustering loss for enhanced clustering. By employing labelfree category-aware clustering loss, we endow the model with potential zero-shot discrimination capabilities, particularly relevant in cases where the label space may differ between the training and testing phases, as observed in our dataset where the label categories of the training and testing sets are inconsistent. During network training, these two types of losses are combined using different weights as hyperparameters. We calculate the CLIP losses, including the loss between EEG features and image features and the loss between EEG features and text features, to ensure that the EEG features can reflect relevant information. Subsequently, we calculate the category-aware clustering loss for the classification results. The CLIP loss, Lalign and the MSE loss, LMSE, between image features, FI , and EEG features, FE, can be written as follow:

where k is the index of training sample.

Category-aware Clustering Loss: A key innovation of our approach lies in implementing contrastive learning without relying on absolute class labels through class-aware clustering loss. We dynamically construct relationships within each batch by examining the relative similarities among samples. Instead of depending on predefined classes, we adaptively create positive and negative pairs based on the data in the current batch, enabling the model to learn feature representations through comparison. This similarity-based clustering operation is independent of the total number of classes in the dataset, making it particularly effective in scenarios where class distributions may vary between training and testing phases. The category-aware clustering loss can be formulated as follow:

In the formulas above, by element-wise multiplication, the Sim refers to the similarity between the feature, Qkj , which belongs to the kth sample of the j th class, and the center of class j, Cj . J is the amount of categories in the current batch. ϵ is used as a constant to make sure the numerical stability during calculating. And the constant M is used as the threshold. When the dissimilarity achieves M, the loss assumes that these two samples could be completely separated. The final loss function can be written as follow:

where K is the number of samples in the current batch. λ1, λ2 and λ3 are the weighting parameters to balance these three loss functions.

3.6 Dataset

We conducted our experiments using the THINGS-EEG [Gifford et al., 2022] dataset, which is a large-scale collection of EEG recordings designed to capture neural signals to visual stimuli. The THINGS-EEG dataset was developed as part of the THINGS initiative 1 , aiming at exploring the neural representation of object concepts. It comprises EEG recorded from participants who were exposed to a diverse set of real-world object images. These images were meticulously selected from the broader THINGS dataset [Hebart et al., 2019], which encompasses over 26,000 object concepts spanning a wide semantic and perceptual range. In this work, MNE python package was applied to conduct the preprocessing. [Gramfort et al., 2013]. The continuous EEG data were segmented into trials from 200ms before stimulus onset to 800ms after stimulus onset. Baseline correction was implemented by subtracting the mean pre-stimulus interval for each trial and channel. Subsequently, the data were downsampled to 100Hz. Trials containing target stimuli were excluded from our experiment. For our experiments, we employed the same training and testing datasets from the THINGS-EEG dataset as those used in previous research.

3.7 Evaluation

In our experiments, we primarily used the classification accuracy including Top-1 accuracy and Top-5 accuracy to evaluate the capability of EEG encoder. For the reconstruction of visual stimuli, traditional quantitative metrics, such as MSE and SSIM, are not appropriate. This is because we only need to reconstruct images with consistent semantics and do not require pixel-level reconstruction recovery. Therefore, to assess the semantic information of the generated images, we creatively propose a semantic-based score as an evaluation metric for the quality of the generated images, based on the image classification model pre-trained on ImageNet, which, to the best of our knowledge, is the first time to be proposed. Specifically, we utilize the pre-trained weights of the ConvNext model fromMeta [Liu et al., 2022] to classify the generated images. For any generated image, we adopt the following rules:

1. If the result of Top-1 from the ConvNext model contains the same label, or the THINGS-EEG labels are included in the ImageNet labels, the score of the image is set to 1.
2. If the results of Top-5 from the ConvNext model contain the same label, the score of the image is set as the sum of the probabilities of the classes that contain the same label.

1 https://things-initiative.org/

1. If the results of Top-5 from the ConvNext model do not contain the exact same label, we will use WordNet [Miller, 1995] to obtain the label's classes.

Then we calculate the score using the following formula:

Here, Pi represents the probability of each class obtained from the pre-trained ConvNext model, and Sim is the WuPalmer Similarity [Wu and Palmer, 1994].

Here, we carefully selected several representative methods, as comparison methods, including ATM-S, ATM-E [Li et al., 2024], NERV [Chen, 2024], NICE [Song et al., 2023], EITResNet [Zheng et al., 2024; He et al., 2016]. Also, in the following part, we compared the generated images in the same subject, which is the subject "sub-08" to assess whether the generated images are able to reflect the semantic information successfully.

3.8 Implementation Details

In this work, we well-trained the proposed model with a maximum of 40 epochs, a learning rate of 3×10-4 , and a batch size of 64. Early stopping was applied, and training was halted if the top-1 classification accuracy did not improve for 10 consecutive epochs. To complement the evaluation, we also recorded the accuracy of the top 5 classification to capture the five most probable class predictions, which were further applied to analyze semantic consistency. In the whole training process, we also applied pre-trained models in different parts to reduce the computation cost and increase the generality. First, we applied the public CLIP weights to generate image embedding from the testing and training dataset of THINGSEEG. Then, we applied the pre-trained diffusion models including SDXL-Turbo to generate images [Podell et al., 2023; Sauer et al., 2025; Ye et al., 2023]. Finally, we used the pretrained ConvNext trained with ImageNet [Deng et al., 2009] and wordNet [Miller, 1995] to calculate the semantic scores based on the classification results.

4 Experimental Results

4.1 Visual Inspection of Image Reconstruction

To evaluate the ability of our framework to reconstruct highquality images from EEG, we conducted relevant imagegeneration tasks on the THINGS-EEG dataset using all relevant model reconstructions. Fig. 3 exhibits the representative reconstructed images from subject-08. The first column presents the visual stimuli corresponding to the EEG signals. In terms of the semantic quality of the generated images, significant differences exist among the results produced by different methods shown in the subsequent columns. The second column displays the images generated by our method. From Fig. 3, we can see that When faced with the visual stimulus of "antelope", both our method and ATM-S retain, to a certain extent, the key features related to "antelope". Although pixel-level precise reconstruction may not be achieved, the overall outlines and main structures are highly recognizable. For example, the shape of the antelope can be roughly outlined. However, the images reconstructed by other methods resemble either cows or horses. For another instance, considering the visual stimulus of "bird" in the third row, only our method can, to some degree, make the generated image recognizable as a bird. Other methods fail to do so. In particular, the image reconstructed by ATM-E is not an animal-related image at all, showing a huge semantic discrepancy.

4.2 Quantitative Semantic Analysis

From Fig. 3, it can be observed that although our method cannot reconstruct visual stimuli at the pixel level, it is capable of generating images with similar semantics. To quantitatively compare the semantic similarity of images generated by different methods, in this work, we propose a quantitative semantic-based score. We categorize the scores of all generated images into three types based on the mean and standard deviation. For scores above (mean + std), the images are classified as good. For scores between (mean + std) and (mean - std), the images are classified as intermediate. For scores below (mean - std), the images are classified as bad. From Fig. 4, we can see that, to some extent, the higher the score, the closer the semantics of the reconstructed image is to the visual stimuli corresponding to the EEG signals. By comparing the magnitudes of these scores, we can, to a certain extent, conduct a horizontal comparison of the quality of images reconstructed by different methods, seen in Table 1.

4.3 Zero-shot EEGClassification

In the process of our EEG reconstruction of visual stimuli, a powerful EEG feature extractor is required to better extract corresponding semantic information from EEG. Evaluation of the ability of EEG feature extraction to extract relevant semantic information from EEG can be used to assess the ability of each model to reconstruct EEG visual stimuli. Therefore, we conducted a zero-shot classification task using the THINGS-EEG testing dataset. By aligning the EEG features with image features through the CLIP loss, we simultaneously optimize the EEG features via category-aware clustering loss, enabling the separation of different representations in the high-dimensional space while clustering similar ones. Subsequently, these features are processed by a lightweight Multi-Layer Perceptron (MLP) classifier.

The results, as presented in Table ??, demonstrate that our model achieves superior performance compared to other baseline methods. To be specific, although our method does not achieve the best Top-1 and Top-5 accuracy across all subjects, overall, it demonstrates the best performance, which can be clearly observed from the average precision. This suggests that our method is more stable and effective in handling the task in general, rather than excelling in only a few specific cases.

4.4 Ablation Study

To explore the impact of each component on the proposed model, several ablation studies were conducted. First, based on our model, we removed the feature alignment loss. As shown in Table 2, Lalign is the most crucial loss for training our model. Without Lalign, the accuracy of EEG classification becomes extremely low, indicating that the model cannot fit the data under such circumstances. Furthermore, we removed the DWT module, the local branch, and the global branch separately to train the classification task for each subject. The accuracy of the classification task decreased significantly across all subjects. A similar phenomenon also occurred in another ablation experiment: when we changed the gated attention from the global branch to the local branch. Regarding the other two loss functions LMSE and LC during the training process, we found that removing either one of them led to a slight decline in the performance of our model.

5 Discussion

EEG, as an indispensable tool in BCI, enables non-invasive acquisition of relevant brain information. EEG can be utilized to understand brain-related activities and guide clinical treatment and assisted living for relevant individuals. Among them, the reconstruction of relevant visual stimuli from EEG can make EEG more tangible, which is conducive to connecting individuals with the real world.

In this work, we present a novel framework for EEGto-image reconstruction, which enhances the extraction of meaningful features from EEG. To be specific, the proposed model integrates advanced components such as the DWT block and gated attention mechanisms, as well as a novel mixed loss function that combines CLIP loss for feature alignment with a label-free category-aware clustering loss, aiming to improve classification accuracy and semantic alignment. Furthermore, we innovatively utilize WordNet-based classification accuracy and semantic similarity measures and creatively propose a semantic-based score to objectively evaluate the semantic information of generated images, addressing limitations in existing evaluation methods and providing a scalable and quantitative way to assess the semantic consistency of reconstructed images.

Compared with previous results, EIT [Zheng et al., 2024] applied mature deep learning tools such as ResNet as the EEG feature extractor for image reconstruction. However, despite its advance in efficiency and computation cost, the nature of such computer vision models may not be suitable for processing complex time-series data, especially for EEG data and the complicated downstream tasks such as high-level EEG feature extraction and feature alignment with images. Compared to NERV [Chen, 2024] and ATM-S [Li et al., 2024], our model successfully introduced the DWT module as a frequency-domain feature extractor. Also, the application of category-aware clustering loss leads to further performance improvements.

Based on our analysis, there is a strong correlation between the classification accuracy and the semantic score of the generated images. Specifically, EEG that are correctly classified by the model typically result in images with higher semantic similarity scores. In contrast, EEG that are not correctly classified usually produce images with significantly lower semantic similarity scores. This finding emphasizes the importance of accurate classification in maintaining semantic fidelity during the EEG-to-image conversion process. Challenges found when generating images for specific categories: Some categories completely fail to generate semantically meaningful images. We have identified several potential reasons for this limitation: (1) Ambiguous or incorrect labels: Categories with unclear or incorrect labels, such as "bator4", pose significant challenges. Even when converted into semantic vectors through CLIP, these labels lack meaningful semantic relationships, leading to image generation failures. (2) Diffusion model constraints: In this paper, the directly used pre-trained diffusion model for reconstructing visual stimuli may lack corresponding visual representations for certain categories, limiting its ability to generate relevant images.

Despite numerous challenges, this study highlights the potential of EEG-based image generation as a research frontier. Future work could explore the use of interpretable models or traditional EEG analysis methods to extract more robust features, thereby enhancing the interpretability of the process. Moreover, improving the consistency between EEG features and semantic representations through more sophisticated architectures or loss functions could further advance the field. Developing customized pre-training strategies for diffusion models to include a broader range of categories is another promising avenue for improving image generation results.

In conclusion, our findings demonstrate the feasibility and value of EEG-based image generation, offering insights into the semantic encoding of neural signals and paving the way for more interpretable and accessible brain-computer interface technologies.

Ethical Statement

There are no ethical issues.

Acknowledgments

This work was partially supported by the Beijing Natural Science Foundation (L245015), the National Key R&D Program Program of China (2023YFC2706400), the Chongqing Natural Science Foundation (CSTB2024NSCQ-MSX0451), and the Medical Innovation Program(MIP) in Kyoto University. The authors also would like to show a great appreciation towards Dr.Yoshifumi Mori, who now works in Netherlands Institute for Neuroscience, for his valuable suggestions and insights in psychiatric diseases.