# ConceptVAE: Self-Supervised Fine-Grained Concept Disentanglement from 2D Echocardiographies
## 3 ConceptVAE
### Model Architecture  
Figure 2 displays the detailed architecture of ConceptVAE. A simple auto-encoder operates independently (in terms of gradients) from the rest of the model. It comprises an _Encoder Stem_ that generates features \(x_{stem}\) at a 4\(\times\) output stride, and an _Image Decoder_ that reconstructs the original input. After a stop-gradient operation, an _Encoder Middle_ block applies a series of residual convolutional blocks starting from the encoder stem's features, projecting the features to concepts. The projections are used by a _Concept Discretizer_ classification head, with \(x_{middle}\) having a 16\(\times\) output stride. For each spatial location, a Softmax activation creates a probability distribution over \(C\) concepts. Using the Gumbel-Softmax trick [38] with hard sampling and gradient pass-through, a grid of one-hot vectors is sampled from the concept probabilities grid. This one-hot vector grid indexes a learned matrix of concept embeddings to produce a 2D concept map \(x_{concept}\). Subsequently, \(x_{middle}\) and \(x_{concept}\) are concatenated along the channel axis and passed into a _Concept Stylizer_ block. This block generates a 2D grid \(x_{style}\) of \(S\) channels capturing the style properties of each concept. At this point, each location within the \(16\times\)-stride grid has an identified concept and an associated style vector. The channel-wise concatenation of \(x_{concept}\) and \(x_{style}\) constitutes the model's latent space (\(x_{latent}\)). Notably, \(x_{concept}\) is derived from discrete embeddings, using a shared learnable embedding matrix for all input samples. In contrast, \(x_{style}\) is a continuous tensor computed based on local features \(x_{middle}\) and the sampled discrete concepts \(x_{concept}\). Consequently, \(x_{style}\) is specific to the sampled \(x_{concept}\), meaning that sampling a different concept at location \(i,j\) will result in a different style vector \(x_{style}^{ij}\). A _Feature Decoder_ projects the latent space to reconstruct the lower \(4\times\)-stride features of the _Encoder Stem_, denoted as \(x_{stem}^{rec}\). Lastly, the _EMA Decoder_ is employed to recover the original input image from the latent space. This reconstruction is core to ConceptVAE, as it guides the model to learn how to decompose an input into fine-grained concepts with associated styles, and reconstruct the input from concepts alone or from concepts and  
Figure 2: ConceptVAE model architecture and training setup, where the EMA blocks represent the exponential moving average mirrors of regular blocks. Loss components are shown in colored ellipses, and _s.g._ denotes _stop-gradient_. Solid arrows indicate tensor flows within the model, while dashed arrows represent tensors involved in loss functions. associated styles. Using the _EMA Decoder_ for the reconstruction ensures there is no mode collapse for the concepts or styles. Architecturally, the _Encoder Stem_ module is designed as a simple sequence of convolutional, instance normalization, max-pooling, and Leaky ReLU stages. The final layer is a normalization layer that ensures channel-wise zero mean and unit standard deviation, helping to prevent potential feature collapse. This module contains three convolutional layers with \(3\times 3\) kernels and strides 2, 1, 1 respectively, and one max-pooling layer with \(2\times 2\) kernel and stride 2, yielding a field of view size of 17 px. The _Image Decoder_ block maintains this simplicity, consisting of 2 upsampling stages based on \(3\times 3\) transposed convolution layers with stride 2. Regular \(1\times 1\) convolutions, normalization, and Leaky ReLU layers are inter-twined between the two up-sampling stages to improve the module's decoding capacity. The _Encoder Middle_ block employs a residual architecture. As in the _Image Decoder_ block, the first layer is a Leaky ReLU activation, as the input to this block comes from the normalized convolutional output of the Encoder Stem.