# ConceptVAE: Self-Supervised Fine-Grained Concept Disentanglement from 2D Echocardiographies
## 3 ConceptVAE
### Training Objectives  
To train ConceptVAE, we devise a series of loss terms inspired by classical (discrete) VAE formulations, but adapted to guide the learning process towards identifying and personalizing concepts. We employ two types of reconstruction losses, illustrated in blue in Figure 2: an image-based loss \(\mathcal{L}_{img}\), which uses Mean Squared Error (MSE) over pixel values, and a feature-based loss \(\mathcal{L}_{feat}\), which uses MSE over low-level feature tensors. The simple auto-encoder is trained using \(\mathcal{L}_{img}\) between the original input image \(\mathcal{I}_{orig}\) and the reconstructed image based on the \(4\times\)-stride feature map. The EMA version of the _Encoder Stem_ is used to compute the target for the tensor produced by the _Feature Decoder_ block, while the _EMA Decoder_ is used to compute the reconstructed image from \(x_{latent}\). The use of both pixel- and feature-level reconstruction losses has been previously employed in VAE/GAN setups [34; 39], to boost both training stability and image generation fidelity. The feature decoder takes both \(x_{concept}\) and \(x_{style}\) as inputs. While \(x_{concept}\) is generated by sampling from a discrete concept codebook, \(x_{style}\) is computed directly as a (continuous) function of \(x_{middle}\) and \(x_{concept}\). Consequently, the network could potentially exploit this setup by minimizing the influence of \(x_{concept}\) and relying more heavily on the more direct path of \(x_{style}\), effectively reducing its operation to that of a simple auto-encoder. In this scenario, \(x_{concept}\) would lose its semantic significance, and \(x_{style}\) would function as a rich bottleneck representation rather than a style characteristic of a concept. To address this undesired behavior, an image/feature reconstruction is performed where the style components of \(x_{latent}\) are explicitly zeroed out. The _EMA Decoder_ is reused to obtain a reconstructed version of the input image, relying solely on \(x_{concept}\), without the style component \(x_{style}\). The target of this reconstruction is a blurred version of the input image, with blurring serving as an approximation for removing fine details and textures, thereby partially eliminating the notion of style. Both pixel- and feature-based losses are employed to evaluate the reconstruction quality when using only the spatial distribution of concepts. This approach guides the _Feature Decoder_ block to focus on the concept component of \(x_{latent}\) and also encourages the _Encoder Middle_ to learn to detect relevant concepts within input images. Another key aspect of concept detection is its invariance to specific styles. This means that two different (augmented) views of the same medical image should produce the same concept maps, despite variations in their visual appearances. Pixel-level and texture differences should be captured by \(x_{style}\), while more complex anatomical structures should be encoded in \(x_{concept}\). To guide this behavior during training, we introduced a _Concept consistency loss_, illustrated with orange in Figure 2. The _Concept Discretizer_ block first computes a grid of concept probabilities, from which it generates a spatial grid of sampled concept indices. Following this, the concept maps from augmented views should be equivalent, even if the augmentations involve translations, rotations, or other spatial shifts (We use _equivalent_ instead of _identical_ because augmentations like translations, rotations, and shearing can spatially shift the placement of concepts within the image. Nevertheless, the correspondences between the initial and shifted locations are known, and they can be used to enforce similarity between \(p(c)|\mathcal{I}_{orig}\) and \(p(c)|\mathcal{I}_{augm}\)). The _EMA Encoder Stem_, _EMA Encoder Middle_, and the _EMA Concept Discretizer_ are used to compute the target probability distributions \(p_{ema}(c)\) for the concept consistency loss: \(\mathcal{L}_{cc}=-p_{ema}(c)\ln p(c)\).