SlowGuess/NeurIPS_2023_OCR huggingface dataset:
To better understand the effects of counterfactual representation augmentation (first introduced in section 3.3), we ran extra experiments with different  $S$  and  $K$ . We choose the vision problem since a 3D latent space manifests a more obvious difference when physical symmetry is applied. In section 5.1, we show that a larger augmentation factor  $K$  leads to higher sample efficiency. In section 5.2, we visualise the change of learned latent space against training epoch according to different values of  $K$ . In section 5.3, we show that some deliberately incorrect group assumptions  $S$  can also achieve good results.
5.1 Counterfactual representation augmentation improves sample efficiency
Figure 6 shows that a larger factor of counterfactual representation augmentation leads to a lower linear projection loss (the measurement defined in section 4.2.2) of the learned 3D representation. Here,  $K$  is the augmentation factor, and  $K = 0$  means the model is trained without physical symmetry. The comparative study is conducted on 4 training set sizes (256, 512, 1024, and 2048), in which each box plot shows the results of 10 experiments trained with a fixed  $K$  and random initialisation. We see that a larger  $K$  leads to better results and compensates for the lack of training data. E.g., the loss trained on 256 samples with  $K = 4$  is comparable to the loss trained on 1024 samples with  $K = 0$ , and the loss trained on 512 samples with  $K = 4$  is even lower than the loss trained on 2048 samples with  $K = 0$ . Furthermore, when  $K = 0$ , increasing the number of training samples beyond a certain point does not further shrink the error, but increasing  $K$  still helps.
5.2 Counterfactual representation augmentation improves interpretability
Figure 7 visualises the latent space during different stages of model training, and we see that a larger  $K$  leads to a better enforcement of interpretability. The horizontal axis shows the training epoch. Three experiments with different  $K$  values ( $\times 0, \times 4, \times 16$ ) are stacked vertically. Each experiment is
Figure 6: A comparison of linear projection MSEs among different augmentation factors  $(K)$  and training set sizes, which shows that counterfactual representation augmentation improves sample efficiency.
Figure 7: A visualisation of the learned latent space against training epoch, in which we see that a larger  $K$  leads to a stronger enforcement on learning a linear latent space. Here, we plot how the encoder projects an equidistant 3D grid of true Cartesian coordinates onto the  $z$  space. Different colours denote respective axes in the true coordinates.
trained twice with random initialisation. Each subplot shows the orthogonal projection of the  $z$  space onto the plane spanned by  $z_{1}$  and  $z_{3}$ , therefore hiding most of the  $y$ -axis (i.e. ball height) wherever a linear disentanglement is fulfilled. During training, the role of physical symmetry is to "straighten" the encoded grid and a larger  $K$  yields a stronger effect.
5.3 Counterfactual representation augmentation with deliberately incorrect group assumptions
Additionally, we test SPS with deliberately incorrect group assumptions. The motivation is as follows. In real applications, researchers may incorrectly specify the symmetry constraint when the data are complex or the symmetry is not known a priori. SPS is more useful if it works with various groups assumptions close to the truth. In our analysis, we are surprised to find that SPS still learns interpretable representations under alternate group assumptions via perturbing the correct one.
Figure 8 shows our results with the vision task (on the bouncing ball dataset). The  $x$  tick labels show the augmentation method. Its syntax follows section 3.3, e.g.,  $(\mathbb{R}^1, +) \times \mathrm{SO}(2)$  denotes augmenting representations by 1D translations and 2D rotations. The  $y$  axis of the plot is still linear projection loss (as discussed in section A.5.3) that evaluates the interpretability of the learned representation. As is shown by the boxplot, five out of five perturbed group assumptions yield better results than the "w/o Symmetry" baseline. Particularly,  $(\mathbb{R}^3, +) \times \mathrm{SO}(2)$  and  $(\mathbb{R}^2, +) \times \mathrm{SO}(3)$  learn significantly more linear representations, showing that some symmetry assumptions are "less incorrect" than others, and that SPS can achieve good results under a multitude of group assumptions.
Figure 8: Evaluation on various group assumptions, which shows that physical symmetry is a very robust inductive bias as even the incorrect symmetry assumptions lead to better results than baseline. Here, the  $y$  axis is linear projection loss between the learned location factor and the true coordinates, so a lower value means better interpretability of representations. The leftmost box shows the baseline without symmetry constraint. The next five boxes show five deliberately incorrect group assumptions, trained with  $K = 4$ . The rightmost box shows the correct group assumption.
The idea of using a predictive model for better self-supervised learning has been well established [Oord et al., 2018; Chung et al., 2015; LeCun, 2022]. In terms of model architecture, our model is very similar to VRNN [Chung et al., 2015]. In addition, our model can be seen as a variation of joint-embedding predictive architecture (JEPA) in [LeCun, 2022] if we eliminate the reconstruction losses on the observation. In fact, we see the network topology of a model as the "hardware" and see the learning strategy (e.g., contrastive method, regularised method, or a mixed one) as the "software". The main contribution of this study lies in the learning strategy — to use physical symmetry to limit the complexity of the prior model, and to use counterfactual representation augmentation to increase sample efficiency.
The existing notation of "symmetry" as in [Higgins et al., 2018; Bronstein et al., 2021] is very different from physical symmetry as an inductive for representation learning. Most current symmetry-based methods care about the relation between observation  $x$  and latent  $z$  [Sanghi, 2020; Quessard et al., 2020; Dupont et al., 2020; Huang et al., 2021]. E.g., when a certain transformation is applied to  $x$ ,  $z$  should simply keep invariant or follow a same/similar transformation. Such an assumption inevitably requires some knowledge in the domain of  $x$ . In contrast, physical symmetry focuses solely on the dynamics of  $z$ , and therefore we only have to make assumptions about the underlying group transformation in the latent space. We see two most relevant works in the field of reinforcement learning [Mondal et al., 2022; Dupont et al., 2020], which apply an equivariant assumption similar to the physical symmetry used in this paper. The major differences are twofold. First, to disentangle the basic factors, our method requires no interactions with the environment. Second, our method is much more concise; it needs no other tailored components or other inductive biases such as symmetric embeddings network and contrastive loss used in [Dupont et al., 2020] or MDP homomorphism applied in [Mondal et al., 2022].
We have identified several limitations in the generality and soundness of SPS. Firstly, when the underlying concept following physical symmetry only contains partial information of the time series and cannot fully reconstruct the inputs, SPS may not function properly. We hypothesize that this issue is connected to content-style disentanglement, and present some preliminary results in appendix A.4 and A.5. Secondly, the current model lacks the ability to distill concepts from multibody systems. For example, it is unable to learn the concept of pitch from polyphonic music or understand 3D space from videos featuring multiple moving objects. Lastly, it is essential to develop a formalized theory for quantifying the impact of counterfactual representation augmentation in future work. This would involve measuring the degree of freedom in the latent space with and without physical symmetry, and explaining why incorrect symmetry assumptions can still result in a correct and interpretable latent space.
In this paper, we use physical symmetry as a novel inductive bias to learn interpretable and low-dimensional representations from time-series data. Experiments show that physical symmetry effectively distills an interpretable linear pitch concept, which agrees with human music perception, from music audios without any labels. With the same method, we can learn the concept of 3D Cartesian space from monocular videos of bouncing ball shot from a fixed perspective. In addition, a robust training technique, counterfactual representation augmentation, is developed to enforce physical symmetry during training. Analysis shows that counterfactual representation augmentation leads to higher sample efficiency and better latent-space interpretability, and it stays effective even when the symmetry assumption is incorrect. Last but not least, we see that with physical symmetry, our sequential representation learning model can drop the the Gaussian prior regulation on the latent space. Such a result empirically indicates that physical symmetry, as a causal (counterfactual) inductive bias, might be more essential compared to the Gaussian prior as a purely statistical regularization.
We'd like to thank to Dr. Zhao Yang and Dr. Maigo Wang for inspiring discussion on physical symmetry. Also, we'd like to like to extend our sincere thanks to Junyan Jiang and Yixiao Zhang for their useful and practical suggestions.
Taketo Akama. Controlling symbolic music generation based on concept learning from domain knowledge. In Arthur Flexer, Geoffroy Peeters, Julián Urbano, and Anja Volk, editors, Proceedings of the 20th International Society for Music Information Retrieval Conference, ISMIR 2019, pages 816-823, 2019.
Junwen Bai, Weiran Wang, and Carla P Gomes. Contrastively disentangled sequential variational autoencoder. Advances in Neural Information Processing Systems, 34:10105-10118, 2021.
Samy Bengio, Oriol Vinyals, Navdeep Jaitly, and Noam Shazeer. Scheduled sampling for sequence prediction with recurrent neural networks. Advances in Neural Information Processing Systems, 28, 2015.
Michael M Bronstein, Joan Bruna, Taco Cohen, and Petar Velickovic. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. arXiv preprint arXiv:2104.13478, 2021.
Kyunghyun Cho, Bart van Merrienboer, Dzmitry Bahdanau, and Yoshua Bengio. On the properties of neural machine translation: Encoder-decoder approaches. In Proceedings of SSST-8, Eighth Workshop on Syntax, Semantics and Structure in Statistical Translation, pages 103-111, Doha, Qatar, October 2014. Association for Computational Linguistics.
Collins Chris. GeneralUser GS, 2017.
Junyoung Chung, Kyle Kastner, Laurent Dinh, Kratarth Goel, Aaron C Courville, and Yoshua Bengio. A recurrent latent variable model for sequential data. In C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 28. Curran Associates, Inc., 2015.
Alain De Cheveigné and Hideki Kawahara. Yin, a fundamental frequency estimator for speech and music. The Journal of the Acoustical Society of America, 111(4):1917-1930, 2002.
Emilien Dupont, Miguel Bautista Martin, Alex Colburn, Aditya Sankar, Josh Susskind, and Qi Shan. Equivariant neural rendering. In International Conference on Machine Learning, pages 2761-2770. PMLR, 2020.
Jesse Engel, Rigel Swavely, Lamtharn Hanoi Hantrakul, Adam Roberts, and Curtis Hawthorne. Self-supervised pitch detection by inverse audio synthesis. In ICML 2020 Workshop on Self-supervision in Audio and Speech, 2020.
The appendix is structured into three main parts.
The first part (section A.1, A.2) provides additional details about SPS. Section A.1 focuses on implementation-related aspects, while section A.2 presents experimental results concerning the reconstruction and prediction loss.
The second part (section A.3 to A.5) introduces an extended version of SPS called  $\mathrm{SPS + }$ . Section A.3 describes the capabilities of  $\mathrm{SPS + }$  in achieving content-style disentanglement along with interpretable learning. The related experiments are presented in section A.4 and section A.5.
The third part (section A.6) presents two additional complex experiments conducted separately using SPS+ and SPS, respectively.
A.1 SPS implementation details
A.1.1 Architecture details
Our models for both tasks share the following architecture. The encoder first uses a 2D-CNN with ReLU activation to shrink the input down to an  $8 \times 8$  middle layer, and then a linear layer to obtain  $z$ . If the encoder is in a VAE (instead of an AE), two linear layers characterises the posterior, one for the mean and the other for the log-variance. The prior model is a vanilla RNN of one layer with 256 hidden units and one linear layer projection head. The decoder consists of a small fully-connected network followed by 2D transposed convolution layers mirroring the CNN in the encoder. Its output is then passed through a sigmoid function. We use no batch normalisation or dropout layers.
Minor variations exist between the models for the two tasks. In the audio task, we use three convolution layers in the encoder, with three linear and three 2D transposed convolution layers in the decoder. In the vision task, as the data are more complex, we use four convolution layers in the encoder, with four linear and four 2D transposed convolution layers in the decoder.
A.1.2 Training details
For both tasks, we use the Adam optimiser with learning rate  $= 10^{-3}$ . The training batch size is 32 across all of our experiments. For all VAE-based models, including  $\mathrm{SPS}_{\mathrm{VAE}}$  (ours/ablation) and  $\beta$ -VAE (baseline), we set  $\beta$  (i.e.,  $\lambda_3$  in Equation (1)) to 0.01, with  $\lambda_1 = 1$  and  $\lambda_2 = 2$ . All BCE and MSE loss functions are calculated in sum instead of mean.  $K = 4$  for all SPS models except for those discussed in section 5 where we analyse the influence of different  $K$ .
The RNN predicts  $z_{n+1:T}$  given the first  $n$  embeddings  $z_{1:n}$ . We choose  $n = 3$  for the audio task and  $n = 5$  for the vision task. We adopt scheduled sampling [Bengio et al., 2015] during the training stage, where we gradually reduce the guidance from teacher forcing. After around 50000 batch iterations, the RNN relies solely on the given  $z_{1:T}$  and predicts auto-regressively.
A.2 SPS reconstruction and prior prediction results
We investigate the reconstruction and prediction capacities of our model and show that they are not harmed by adding symmetry constraints. For the music task, we compare our model, our model abating symmetry constraints, and a  $\beta$ -VAE trained solely for the reconstruction of power spectrogram. Table 3 reports per-pixel BCE of the reconstructed sequences from the original input frames (Self-recon) and from the RNN predictions (Image-pred). We also include  $\mathcal{L}_{\mathrm{prior}}$ , the MSE loss on the RNN-predicted  $\hat{z}$  as defined in section 3.2. The results show that our models slightly surpasses the ablation and baseline models in all three metrics.
Similarly, Table 4 displays the reconstruction and prediction losses on the test set for the video task. Results show that adding symmetry constraints does not significantly hurt the prediction losses. Frame-wise self-reconstruction is significantly lower for the SPS models, but only by a small margin.
A.3  $\mathbf{SPS}+$
SPS can use physical symmetry to learn interpretable factors that evolve over time. We call those factors content representation. However, many problems can not be represented by content represent-
Table 3: Reconstruction and prediction results on the audio task.
Methods	Self-recon ↓	Image-pred ↓	Lprior ↓
SPSVAE, K=4 (Ours)	0.0292±0.0003	0.0296±0.0005	0.0012±0.0006
SPSAE, K=4 (Ours)	0.0292±0.0002	0.0296±0.0002	0.0012±0.0003
SPSVAE, K=0 (Ablation)	0.0291±0.0003	0.0295±0.0004	0.0030±0.0033
SPSAE, K=0 (Ablation)	0.0291±0.0002	0.0295±0.0004	0.0087±0.0212
β-VAE	0.0303±0.0008	N/A	N/A
Table 4: Reconstruction and prediction losses of the video task. Two outliers are removed from the 50 runs.
Method	Self-recon ↓	Image-pred ↓	Lprior ↓
SPSVAE, K=4 (Ours)	0.64382 ± 9e-05	0.6456 ± 4e-04	0.14 ± 0.05
SPSAE, K=4 (Ours)	0.64386 ± 7e-05	0.6458 ± 3e-04	0.17 ± 0.07
SPSVAE, K=0 (Ablation)	0.64372 ± 4e-05	0.6459 ± 2e-04	0.19 ± 0.10
SPSAE, K=0 (Ablation)	0.64367 ± 5e-05	0.6456 ± 1e-04	0.11 ± 0.03
β-VAE	0.64345 ± 5e-05	N/A	N/A
tation alone. For example, the bouncing balls can have different colours and the pitch scales can be generated by different instruments. If the colour of a ball or the timbre of a sound scale are constant within a trajectory, those latent spaces are hard to constrain by physical symmetry. We call such invariant factors style representation. In order to deal with these problems, we combine SPS with a simple content-style disentanglement technique:  $\mathrm{SPS}+$ , a more general framework of SPS. We use random pooling to constrain the style factors, and use physical symmetry to constrain the content representation in the same way as SPS in Section 3.1
A.3.1 Model
Figure 9 shows the design of  $\mathrm{SPS}+$ , which belongs to the family of disentangled sequential autoencoders [Bai et al., 2021; Hsu et al., 2017; Vowels et al., 2021; Yingzhen and Mandt, 2018; Zhu et al., 2020]. During the training process, the temporal data input  $\mathbf{x}_{1:T}$  is first fed into the encoder  $E$  to obtain the corresponding representation  $\mathbf{z}_{1:T}$ .  $\mathbf{z}_{1:T}$  is then split into two parts: the style factor  $\mathbf{z}_{1:T,s}$  and the content factor  $\mathbf{z}_{1:T,c}$ . The style factor  $\mathbf{z}_{1:T,s}$  is passed through the random-pooling module  $P$ , where one element  $z_{\tau,s}$  is randomly picked. The content factor  $\mathbf{z}_{1:T,c}$  is fed into three branches, then combined with  $z_{\tau,s}$  to reconstruct. For random pooling in the training stage, one style vector is randomly selected from all time steps (i.e., 15 for the music task and 20 for the vision task) of the sequence to represent  $z_s$ . In the testing stage, only the first 5 (vision task) or 3 (music task) frames are given, and  $z_s$  will be selected from them.
A.3.2 Training objective
The following loss functions in  $\mathrm{SPS + }$  slightly vary from those in SPS. For  $\mathrm{SPS + }$ ,  $\mathcal{L}_{\mathrm{prior}}$  and  $\mathcal{L}_{\mathrm{sym}}$  work on the content part of latent variables only. Other loss functions are exactly the same as those defined in section 3.2
$$
\mathcal {L} _ {\text {p r i o r}} = \ell_ {2} \left(\hat {\mathbf {z}} _ {2: T, c}, \mathbf {z} _ {2: T, c}\right), \tag {8}
$$
$$
\mathcal {L} _ {\mathrm {s y m}} = \ell_ {2} \left(\tilde {\mathbf {z}} _ {2: T, c}, \hat {\mathbf {z}} _ {2: T, c}\right) + \ell_ {2} \left(\tilde {\mathbf {z}} _ {2: T, c}, \mathbf {z} _ {2: T, c}\right). \tag {9}
$$
$$
\ell_ {2} \left(\tilde {\mathbf {z}} _ {2: T, c}, \mathbf {z} _ {2: T, c}\right) = \frac {1}{K} \sum_ {k = 1} ^ {K} \ell_ {2} \left(S _ {k} ^ {- 1} \left(R \left(S _ {k} \left(\mathbf {z} _ {1: T - 1, c}\right)\right), \mathbf {z} _ {2: T, c}\right), \right. \tag {10}
$$
Figure 9: An overview of our model.  $\mathbf{x}_{1:T}$  is fed into the encoder  $E$  to obtain the corresponding representation  $\mathbf{z}_{1:T}$ , which is then split into two parts: the style factor  $\mathbf{z}_{1:T,s}$  and the content factor  $\mathbf{z}_{1:T,c}$ . The style factor is passed through the random-pooling layer  $P$ , where an element  $z_{\tau,s}$  is randomly selected. The content factor is fed into three different branches and combined with  $z_{\tau,s}$  to reconstruct three outputs respectively:  $\mathbf{x}_{1:T}^{\prime}, \hat{\mathbf{x}}_{2:T+1}$  and  $\tilde{\mathbf{x}}_{2:T+1}$ . Here,  $R$  is the prior model and  $S$  is the symmetric operation. The inductive bias of physical symmetry enforces  $R$  to be equivariant w.r.t. to  $S$ , so  $\tilde{\mathbf{z}}$  and  $\hat{\mathbf{z}}$  should be close to each other and so are  $\tilde{\mathbf{x}}$  and  $\hat{\mathbf{x}}$ .
$$
\ell_ {2} \left(\tilde {\mathbf {z}} _ {2: T, c}, \hat {\mathbf {z}} _ {2: T, c}\right) = \frac {1}{K} \sum_ {k = 1} ^ {K} \ell_ {2} \left(S _ {k} ^ {- 1} \left(R \left(S _ {k} \left(\mathbf {z} _ {1: T - 1, c}\right)\right), \hat {\mathbf {z}} _ {2: T, c}\right), \right. \tag {11}
$$
$$
\mathcal {L} _ {\mathrm {B C E}} \left(\tilde {\mathbf {x}} _ {2: T}, \mathbf {x} _ {2: T}\right) = \frac {1}{K} \sum_ {k = 1} ^ {K} \mathcal {L} _ {\mathrm {B C E}} \left(D \left(S _ {k} ^ {- 1} \left(R \left(S _ {k} \left(\mathbf {z} _ {1: T - 1, c}\right)\right), z _ {\tau , s}\right), \mathbf {x} _ {2: T}\right). \right. \tag {12}
$$
A.4 SPS+ on learning pitch & timber factors from audios of multiple instruments
A.4.1 Dataset and setups
We synthesise a dataset that contains around 2400 audio clips played by multiple instruments. Similar to the dataset in section 4.1.1, each clip contains 15 notes in major scales with the first 8 notes ascending and the last 8 notes descending. Each note has the same volume and duration. The interval between every two notes is equal. We vary the starting pitch such that every MIDI pitch in the range C2 to C7 is present in the dataset. For each note sequence, we synthesise it using 53 different instruments, yielding 2376 audio clips. Specifically, two soundfont are used to render those audio clips respectively: FluidR3_GM [Wen, 2013] for the train set and GeneralUser GS v1.471 [Chris, 2017] for the test set. The pitch ranges for different instruments vary, so we limit each instrument to its common pitch range (See Table 14).
We assume  $z_{c}\in \mathcal{R}$  and  $z_{s}\in \mathcal{R}^{2}$ , and use random  $S\in G\cong (\mathbb{R}, + )$  to augment  $z_{c}$  with  $K = 4$
A.4.2 Results on pitch-timbre disentanglement
We evaluate the content-style disentanglement using factor-wise data augmentation following [Yang et al., 2019]. Namely, we change (i.e., augment) the instrument (i.e., style) of notes while keeping their pitch the same, and then measure the effects on the encoded  $z_{c}$  and  $z_{s}$ . We compare the normalised  $z_{c}$  and  $z_{s}$ , ensuring they have the same dynamic range. Ideally, the change of  $z_{s}$  should be much more significant than  $z_{c}$ . Here, we compare four approaches: 1) our model (SPS+), 2) our model without splitting for  $z_{s}$  (SPS with  $z \in \mathcal{R}^3$  and  $S \in G \cong (\mathbb{R}, + )$ ) as an ablation, 3) GMVAE [Luo et al., 2019], a domain-specific pitch-timbre disentanglement model trained with explicit pitch labels, and 4) TS-DSA E [Luo et al., 2022], a recent unsupervised pitch-timbre disentanglement model based on Disentangled Sequential Autoencoder (DSA E).
Figure 10 presents the changes in normalised  $z_{c}$  and  $z_{s}$  measured by L2 distance when we change the instrument of an anchor note whose pitch is D3 and synthesised by accordion. Table 5 provides a more quantitative version by aggregating all possible instrument combinations and all different
Figure 10: Comparisons for  $\Delta z_{c}$  and  $\Delta z_{s}$  for different instruments against accordion, with pitch kept constant at MIDI pitch D3.  $\Delta z_{c}$  and  $\Delta z_{s}$  are changes in normalised  $z_{c}$  and  $z_{s}$ , so that higher black bars relative to white bars means better results. All results are evaluated on the test set.
Table 5: Mean ratios of changes in normalised  $z_{c}$  and  $z_{s}$  under timbre augmentation across all possible instrument combinations under different constant pitches in the test set.
Methods	||Δzc||2/||Δzs||2 ↓
SPSVAE+ (Ours)	0.49
SPSVAE (Ablation)	2.20
GMVAE (Baseline)	0.67
TS-DSAE (Baseline)	0.65