Vector-valued Representation is the Key: A Study on Disentanglement and Compositional Generalization

Q1: In our paper, the definition of scalar-value methods is that a factor is encoded by a single dimension, while for vector-value methods, a factor is encoded by several dimensions. SAE utilizes a decoder structure where each layer encodes a unique factor. Consequently, for scalar-value methods, each layer of the decoder is modulated by a single dimension. In contrast, in vector-value method, each layer is modulated by several dimensions.

Therefore, to increase the dimensions in scalar-value methods, we would need to increase the number of layers in the decoder. For each added dimension, an additional layer is required because, by definition, each layer of a scalar-based SAE encodes only one factor. If we were to increase the number of dimensions per layer, the method would transition into a vector-valued method.

Q2: Thank you for your valuable suggestion. To vectorize a method not covered in our paper, we suggest the following systematic set of principles:

1. Identify the basic structure of the model that learns a single factor. Then, extend this basic structure from one dimension to multiple dimensions.
2. Confirm that the inductive bias for disentanglement remains valid after the extension.
3. If the inductive bias still holds after the extension, then we have successfully vectorized the method. If it does not, we need to explore how to modify the inductive bias so that it remains valid for vector representations.

• We are regret that our writting make you misunderstanding that our ideal experiments are to show such difference. We have modicate this part in Section 6.5 (highlight in blue). We would like to highlight: we conduct the ideal experiments to show that in the most optimal situations (i.e., our designed ideal cases), scalar methods can also achieve excellent disentanglement and combinatorial generalization capabilities. Furthermore, we wanted to investigate whether the metrics used in our study prefer vector representations. Our experiments showed that directly mapping scalar-value representations to vector-value based ones did not lead to an increase in CG metrics.

  On the other hand, your suggestion to compare vectorized methods and purely increased dimensionality from a more direct perspective is well received. In response to this, we have conducted additional experiments where the total dimensions were consistent across scalar-based and vector-based methods. The results are as follows:

Please note that FactorVAE (640) has the same total dimension with vec-FactorVAE and β-TCVAE (640) has the same total dimension with vec-β-TCVAE. Our results indicate that under these conditions of consistent latent dimensionality, our compositional generalization still outperforms the new baseline. This might be due to the fact that the disentanglement learning of VAEs is ill-posed when the dimensionality is too large, resulting in poor performance on both disentanglement and compositional generalization.

Thank you for your valuable comments and suggestions. We appreciate your feedback and have made several changes to address your concerns. Please find below our responses to each point you raised.

• The estimation of vector-based total correlation is a very interesting point to research on, although the method is heuristic one, but the method is indeed a good choice for vec-betaTCVAE. Moreover, it's important to note that the total correlation estimation of vec-betaTCVAE method itself is not the primary contribution of our work. Inspired by your comment, we have attempted the following three different methods, and have added these parts into (Section 6.4 and Appendix C highlighted in blue). Note that estimating the total correlation (TC) in vector-valued representations is inherently a difficult problem. We propose three additional ways to approximate TC. These methods are described below, and the experimental results demonstrate that our method performs better in practice than these alternatives.

  1. The total correlation directly derived from multidimensional Gaussian distribution. The specific derivation and calculation methods can be found in the Appendix C of the paper, referred to as "vec-betaTCVAE Sum".
  2. For each training step, we randomly select a variable for each representation vector to calculate TC, referred to as "vec-betaTCVAE Rand".
  3. We calculate TC as the average of multiple runs of the method 2, referred to as "vec-betaTCVAE Rand Sum" (See Appendix C for details)

  The performance of these methods is shown in the table below. Based on our results, our current vectorization method still performs the best, despite the issue raised by the reviewer. We hypothesize that the potential reason for this might be some latent inductive bias in the network, which makes our method even more effective.

For SAE, please refers to Q1.

Dear Reviewer 9DDo,

We would like to thank you for follow-up and increasing the score. We appreciate your valuable feedback and will continue to enhance our work. We fully agree with you that these additions greatly improve the applicability and also sync them into the main text.

• by "basic structure", we mean the structure or loss function (or regularization term) that plays the role of inductive bias. For example, "basic structure" learns a single factor, including: 1. the input of AdaIN-like structure in SAE. 2. each output token in VCT. 3. $q_\phi(z_i)$ in total correlations $KL(q_\phi(z)|\prod_i^m q_\phi(z_i))$. We would like to change "basic structure" into "loss function/architecture inductive bias for disentanglement" for clarity instead. Also, we will include these examples.
• If we misunderstand, please tell us, we are more than happy to answer any remaining questions.
  1. Is it the numbers correct? We checked the code for collecting results and the printed numbers of our experiment and found that the numbers in the Table are correct.
  2. Is it in line with the expectation that vec-beta-TCVAE and vec-SAE do not outperform their scalar-based version on FactorVAE score? Yes, the total correlation used in vec-beta-TCVAE is a heuristic approximation, so the disentanglement performance is expected to be worse than the scalar-valued one with a good estimation. SAE does not leverage any regularization but only the structure itself. There is no explicit constraint on mutual information between the layers. Therefore, the disentanglement performance of vec-SAE is expected to be worse than the scalar-valued SAE with a smaller information bottleneck (limited by scalar).

Thank you for all of your constructive comments and suggestions. We have added the principles to vectorize new methods except those of the paper in conclusion. Please let us know if you have any further questions or concerns, since the discussion stage is coming to an end in thirty minutes.

Q1: The metrics such as accuracy and R2 are indeed applied to the full vector-valued representation in our study. The primary aim of these metrics is to verify the model's ability to generalize to new combinations. In other words, it is meant to assess how well the model can identify the values of these factors under new data combinations.

Q2: We are regret that our writing caused your confusion in Section 6.3. We have revisited and revised the section for clarity. The main content of this section is: In the context of vector-valued representation, we observe a strong positive correlation between the metric of classification and disentanglement, while the metric of regression does not show such a strong correlation. We speculate that the underlying reason might be that regression primarily occurs in factors with multiple values, such as azimuth with 15 values on Shapes3D (For classification, such as scale with 3 values). This leads that the representation of this factor will encode more information. The size of the vector dimensions significantly impacts R2 in such cases.

Q3: We appreciate your insights regarding the strengths and other potential weaknesses of our work. We acknowledge that DCI indeed is also a significant metric. However, we wish to clarify that we do not appreciate to base our conclusions solely on a single metric. In our paper, we have employed multiple measures to evaluate the performance, in line with the widely accepted standards prevalent in the scholarly community. The approach to refer to multiple standards to assess such capabilities has indeed become a consensus within the academic world [1,2,3,4]. On the other hand, our primary focus in the paper is to explore the relationship between disentanglement and generalization. We propose that there exist vector-valued methods that can possess excellent capabilities in both disentanglement and generalization. By vectorizing the methods, we observe an improvement in the combinatorial generalization ability. However, we did not claim that vectorization is a universal technique for improving disentanglement capabilities.

[1] Learning disentangled representation by exploiting pretrained generative models: A contrastive learning view. In International Conference on Learning Representations, 2021.

[2] Structure by architecture: Structured representations without regularization. In The Eleventh International Conference on Learning Representations, 2022.

[3] Infodiffusion: Representation learning using information maximizing diffusion models. international conference on machine learning, 2023.

[4] Visual concepts tokenization. Advances in Neural Information Processing Systems, 2022.

Thank you for your valuable comments and suggestions. We appreciate your feedback and have made several changes to address your concerns. Please find below our responses to each point you raised.

• Thank you for bringing to our attention the omission of the references in our previous work description. We appreciate your insight and have made the necessary changes as per your suggestion (highlight in blue in Section 2). The revised descriptions for the three papers are as follows:
  1. Schott et al. (2021); Montero et al. (2021) go beyond the VAE-based disentanglement studies, though they are still restricted to the scalar-valued representation.
  2. Singh et al. also emphasize the importance of vector-valued representation to compositional generalization.
• Thank you for your detailed feedback and for pointing out the areas of confusion in Section 2.2 of our paper. We are regret that our writing has led to your misunderstanding. We have revised the writtings in Section 2.2, which is highligted in blue. When we mentioned that disentanglement has not been considered in the study of generative models, we were specifically referring to the work by Zhao et al. (2018). We regret that if our wording implied that Xu was the only one who used this approach. Our intention was to emphasize his work along with Montero's but not mentioed that they are the only one. We value your input and will make these modifications in our revised version to avoid these missunderstading.
• We understand your concern about the level of detail provided and we apologize if it was insufficient for a thorough understanding. We appreciate your suggestion to elaborate on why these methods are interesting for testing. We have revised this section in the paper (highlighted in blue in Section 3.2 & 3.3) to provide a clearer and more detailed explanation. We believe that this additional context will help readers better understand our choice of the method for testing.
• We would like to clarify that the results obtained from the use of Beta-VAE and MIG metrics are indeed included in our appendix E. We apologize for any confusion this may have caused, we have highlighted in the main text that the rest of the results are in appendix E.

• 2.2 & 2.3 You are correct in your understanding that our new Total Correlation loss does not impose a penalty on the correlation between different dimensions of different vectors, for instance, between $z_{1,2}$ and $z_{2,1}$. We understand that this might seem counterintuitive, however, our following empirical results speak to the effectiveness of this approach, which suggests that our model still works well even without explicit constraints on correlations of these terms. As for Equation (5), it is a heuristic method. Although we also have actually obtained another three vectorized versions, but found that it did not outperform the orginal version of vec-betaTCVAE in practice. The reasons for this are not entirely clear and require further investigation.

  We think this is a very interesting point and we appreciate the reviewer's suggestion. Inspired by your comment, we have attempted the following three different methods, and have added these parts into our paper (Section 6.4 and Appendix C highlighted in blue). Note that estimating the total correlation (TC) in vector-valued representations is inherently a difficult problem. We propose three additional ways to approximate TC. These methods are described below, and the experimental results demonstrate that our method performs better in practice than these alternatives.

  1. The total correlation directly derived from multidimensional Gaussian distribution. The specific derivation and calculation methods can be found in the Appendix C of the paper, referred to as "vec-betaTCVAE Sum".
  2. For each training step, we randomly select a variable for each representation vector to calculate TC, referred to as "vec-betaTCVAE Rand".
  3. We calculate TC as the average of multiple runs of the method 2, referred to as "vec-betaTCVAE Rand Sum" (See Appendix C for details)

  The performance of these methods is shown in the table below. Based on our results, our current vectorization method still performs the best, despite the issue raised by the reviewer. We hypothesize that the potential reason for this might be some latent inductive bias in the network, which makes our method even more effective.

• 2.4.Thank you for your suggestion, we have add the details of vec-FactorVAE model in the appendix B.

• 3.We appreciate the reviewer bringing these works to our attention. We have now addressed these studies in our paper (highlight in blue) as follows: there is a considerable body of work [1,2,3] exploring how to identify the semantic directions in their latent space, and these directions are precisely represented by vectors. These representations are vector-valued one for semantics.

• 4.The underlying intuition is that essential differences arise between a single-dimensional space and a multi-dimensional one. For instance, the skew lines in three-dimensional space does not exist in a two-dimensional plane. The orthogonality of random vectors in higher-dimensional spaces does not exist in low dimensional space . Therefore, by expanding the single-dimensional representation of a factor in our model to a multi-dimensional one, we hypothesize that the model may be able to learn some fundamental differences. This multi-dimensionality could potentially contribute to achieving better compositional generalization than a scalar one.

• 5.Thank you for your valuable feedback. We have opted to remove the bolding from the values in Table 1 and it was not our intention to mislead our readers.

Q1: Thank you for your inquiry regarding equation (5). As depicted in our new Figure in Appendix C, the term marginal $q(z_j)$ refers to the joint probability distribution of the $j$-th dimension of all the vector representations. We have clarified this point in our paper.

Q2: Thank you for your question regarding the dimensionality matching of our baseline scalar-valued models and vector-valued models in Table 1. We considered two main factors in making this decision. However, in light of your suggestion, we have now included cases with equal total dimensions in Table 1 to provide readers with a more comprehensive understanding. From the perspective of disentangled representation learning, each latent unit is intended to learn a single factor. If the total dimensions were equal, scalar models would have more encoding units in the latent space, leading to an unfair comparison. We took into account models such as SAE and VCT, where the number of model parameters is directly proportional to the number of latent units. More latent units would result in a larger number of parameters for these models, again leading to an unfair comparison.

Thank you for your valuable comments and suggestions. We appreciate your feedback and have made several changes to address your concerns. Please find below our responses to each point you raised.

• We appreciate the reviewer for raising this point, and we fully understand his concern. This indeed might be a question that most readers would have. Therefore, we have added the following experiments, which include results of the two types of VAEs on two datasets regarding disentanglement and compositional generalization across four metrics. Our results indicate that under these conditions of consistent latent dimensionality, our compositional generalization still outperforms the new baseline. This might be due to the fact that the disentanglement learning of VAEs is ill-posed when the dimensionality is too large, resulting in poor performance on compositional generalization. Please note that we have not conducted this experiment for methods like SAE and VCT for the following reasons:
  1. For structural-based methods like VCT and SAE, the number of representation units is related to the number of network layers (SAE) and the size of features (VCT). More representation units lead to a significant increase in the number of parameters, which would result in unfair comparisons.
  2. If we use too many representation units, the computational cost of the model would be quite high. We do not have sufficient computational resources to conduct these set of experiments. | Method | Shapes3D - FactorVAE | Shapes3D - DCI | Shapes3D - R2 | Shapes3D - ACC | MPI3D - FactorVAE | MPI3D - DCI | MPI3D - R2 | MPI3D - ACC |
     |------------------|----------------------|----------------|---------------|----------------|-------------------|-------------|------------|-------------|
     | FactorVAE | 0.83 ± 0.06 | 0.44 ± 0.12 | 0.46 ± 0.18 | 0.39 ± 0.10 | 0.31 ± 0.04 | 0.21 ± 0.01 | 0.30 ± 0.02 | 0.39 ± 0.02 |
     | β-TCVAE | 0.83 ± 0.10 | 0.65 ± 0.16 | 0.45 ± 0.15 | 0.47 ± 0.18 | 0.44 ± 0.05 | 0.27 ± 0.01 | 0.32 ± 0.03 | 0.45 ± 0.03 |
     | SAE | 0.98 ± 0.04 | 0.87 ± 0.12 | 0.72 ± 0.05 | 0.90 ± 0.17 | 0.71 ± 0.04 | 0.47 ± 0.05 | 0.55 ± 0.07 | 0.77 ± 0.02 |
     | VCT | 0.95 ± 0.05 | 0.86 ± 0.02 | 0.56 ± 0.24 | 0.58 ± 0.15 | 0.72 ± 0.04 | 0.47 ± 0.03 | 0.39 ± 0.13 | 0.69 ± 0.09 |
     | FactorVAE (640) | 0.77 ± 0.05 | 0.56 ± 0.12 | 0.77 ± 0.10 | 0.65 ± 0.10 | 0.46 ± 0.03 | 0.43 ± 0.01 | 0.37 ± 0.03 | 0.49 ± 0.02 |
     | β-TCVAE (640) | 0.81 ± 0.11 | 0.74 ± 0.10 | 0.59 ± 0.15 | 0.76 ± 0.18 | 0.43 ± 0.02 | 0.39 ± 0.01 | 0.35 ± 0.02 | 0.44 ± 0.02 | | vec-FactorVAE | 0.93 ± 0.06 | 0.55 ± 0.11 | 0.88 ± 0.05 | 0.96 ± 0.02 | 0.38 ± 0.06 | 0.16 ± 0.05 | 0.53 ± 0.02 | 0.71 ± 0.01 | | vec-β-TCVAE | 0.82 ± 0.08 | 0.31 ± 0.08 | 0.87 ± 0.05 | 0.98 ± 0.01 | 0.42 ± 0.06 | 0.11 ± 0.03 | 0.67 ± 0.02 | 0.78 ± 0.01 |
     | vec-SAE | 0.89 ± 0.08 | 0.63 ± 0.06 | 0.95 ± 0.01 | 0.98 ± 0.01 | 0.62 ± 0.08 | 0.33 ± 0.09 | 0.87 ± 0.03 | 0.88 ± 0.01 |
     | vec-VCT | 0.98 ± 0.04 | 0.85 ± 0.06 | 0.91 ± 0.10 | 0.80 ± 0.09 | 0.70 ± 0.06 | 0.48 ± 0.04 | 0.70 ± 0.07 | 0.77 ± 0.02 |
     | vec-VCT* | 0.97 ± 0.04 | 0.89 ± 0.02 | 0.99 ± 0.02 | 0.90 ± 0.03 | 0.66 ± 0.03 | 0.45 ± 0.06 | 0.85 ± 0.07 | 0.78 ± 0.02 |

2.1 We regret that the notation here is not clear enough, we have modified the text in the paper (highlighted in blue). z_j represents the j’th dimensions of all vectors.

Thank you for your valuable comments and suggestions and the encouraging words. We appreciate the time and effort you took to review our paper. Please find below our responses to your concerns and the changes we have made to address them.

• The way of estiamting the total correlation is a very interesting point and we appreciate the reviewer's suggestion. Inspired by your comment, we have attempted the following three different methods, and have added these parts into our paper (Section 6.4 and Appendix C highlighted in blue). Note that estimating the total correlation (TC) in vector-valued representations is inherently a difficult problem. We propose three additional ways to approximate TC. These methods are described below, and the experimental results demonstrate that our method performs better in practice than these alternatives.

  1. The total correlation directly derived from multidimensional Gaussian distribution. The specific derivation and calculation methods can be found in the Appendix C of the paper, referred to as "vec-betaTCVAE Sum".
  2. For each training step, we randomly select a variable for each representation vector to calculate TC, referred to as "vec-betaTCVAE Rand".
  3. We calculate TC as the average of multiple runs of the method 2, referred to as "vec-betaTCVAE Rand Sum" (See Appendix C for details)

The performance of these methods is shown in the table below. Based on our results, our current vectorization method still performs the best. We hypothesize that the potential reason for this might be some latent inductive bias in the network, which makes our method even more effective.

• Thank you for your insightful comment on investigating the applicability of our learned vector-valued representations to enhance downstream tasks. We conduct experiments in visual reasoning and unfairness. We believe that such an experiment could indeed provide a valuable perspective to our work. We will included these additional experiments in the appendix of our paper and make references to them in the main text. We utilized the settings from disentanglement_lib to evaluate our representations. The experiment results will be updated here.

• Thank you for pointing out the two papers in the field of generalization (Recurrent Independent Mechanisms (RIMs) [1] and Discrete Key-Value Bottleneck [2]). We have discuss them in the updated version. While both these papers emphasize the importance of vector-valued representations for generalization, our work differs significantly in its primary focus. Our research is specifically centered on the aspect of compositional generalization, which is a distinct and critical point of study within the broader context of generalization. Moreover, our key claim is that vector-valued representation is the key simultaneous presence of both disentanglement and compositional generalization abilities, a key feature that we believe sets our work apart from the cited papers. While the cited works have made important strides in understanding the role of vectorization in generalization, our research delves deeper into the relationship between vectorization and compositional generalization, which is also an important perspective.