# ConceptVAE: Self-Supervised Fine-Grained Concept Disentanglement from 2D Echocardiographies
## 3 ConceptVAE
### Training Objectives
The EMA concept probability map \(p_{ema}(c)\) is computed on an augmented view of the initial input image which incorporates transformations such as rotations, translations, shearings, zooming, gamma contrast changing and Gaussian blurring. Since these operations can alter positions, we must account for the spatial mapping between \(p(c)\) and \(p_{ema}(c)\). To simplify this and avoid optimization noise due to imperfect mapping, each augmentation procedure selects a random location uniformly, and all image operations are performed relative to this point. The result includes a tuple of the augmented input image \(\mathcal{I}_{augm}\), an initial location \(l_{ij}\), and the equivalent location \(l_{i^{\prime}j^{\prime}}\) after all operations. In our implementation of \(\mathcal{L}_{cc}\) we indexed only the grid positions of the spatial locations \(l_{ij}\) and \(l_{ff^{\prime}}\) from \(p(c)\) and \(p_{ema}(c)\), respectively. Therefore, only one pair of grid locations (containing the concept probability distributions) is used per each sample inside a training batch. We use the EMA blocks instead of the model blocks to prevent feedback loops that could lead to collapsing concept probabilities (e.g., always detecting the same concept). An additional constraint \(\mathcal{L}_{style}\) was imposed on \(x_{style}\) to ensure it has unit covariance and zero mean along the channel (style) dimension (illustrated with green in Figure 2). Specifically, when \(x_{style}\) is flattened across across batches (\(B\)), height (\(H\)) and width (\(W\)) if forms a matrix of shape shape \((S,BHW)\). This matrix must have a row-wise mean of 0, a row-wise standard deviation of 1, and zero correlation between rows. This constraint ensures that \(x_{style}\) has independent components with a known range of values, discussed in details in Section 5.5. To control the deviation of \(p(c)|\mathcal{I}\) from \(p_{0}(c)\), we use two priors. Without enforcing these priors during training, the entropy of \(p_{ij}(c)\) would be minimized, canceling the effect of concept sampling and reducing the model's operation to a deterministic auto-encoder. Consequently, the concept probability grid \(p(c)|\mathcal{I}\) would lose much of its semantic significance, reverting to a regular discrete latent variable instead of encoding high-level semantics into a fixed set of concept probabilities. This, in turn, would constrain the functionality of the concept consistency loss. We employ two types of priors: at the grid-location level and at image level. Since we are modeling echocardiographies, these images typically feature an ultrasound cone centered within a surrounding black background. The grid-location level prior is computed as follows: for grid locations inside the ultrasound cone, the prior is a uniform distribution over the last \(C-1\) concepts, with the first concept having zero mass (as we always designate the first concept to model the background). For grid locations outside the cone, the prior assigns all probability mass to the first concept. The KL-divergence \(\mathcal{D}_{KL}(p(c)|\mathcal{I}\parallel p_{0}(c))\) is computed at all grid locations and averaged across the \((B,H,W)\) dimensions. For the image-level prior loss it is assumed that only the first concept should be detected outside the cone, with a uniform spread of concepts inside the cone across all samples in the current batch. Therefore, the concept probability vectors of all grid locations inside and outside the echo cones are averaged across all samples in the batch to obtain two image-level concept prevalence vectors: \(d_{cone}(c)\) for the cone region and \(d_{bg}(c)\) for the background. The KL-divergence loss with the same priors is used for these concept prevalence vectors. Equation (2) formalizes the final prior loss \(\mathcal{L}_{prior}\), where \(\mathbf{1}_{c}(b,i,j)\) is an indicator function that equals 1 if location \(i,j\) in sample \(b\) of the current batch pertains to an ultrasound cone.