SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Xuanjie Liu<sup>1,2</sup>
Daniel Chin<sup>1,2</sup>
Yichen Huang
Gus Xia1,2
<sup>1</sup>Mohamed bin Zayed University of Artificial Intelligence
$^{2}$ New York University Shanghai
{Xuanjie.Liu, Nanfeng.Qin, Yichen.Huang, Gus.Xia} @mbzuai.ac.ae
We have recently seen great progress in learning interpretable music representations, ranging from basic factors, such as pitch and timbre, to high-level concepts, such as chord and texture. However, most methods rely heavily on music domain knowledge. It remains an open question what general computational principles give rise to interpretable representations, especially low-dim factors that agree with human perception. In this study, we take inspiration from modern physics and use physical symmetry as a self-consistency constraint for the latent space of time-series data. Specifically, it requires the prior model that characterises the dynamics of the latent states to be equivariant with respect to certain group transformations. We show that physical symmetry leads the model to learn a linear pitch factor from unlabelled monophonic music audio in a self-supervised fashion. In addition, the same methodology can be applied to computer vision, learning a 3D Cartesian space from videos of a simple moving object without labels. Furthermore, physical symmetry naturally leads to counterfactual representation augmentation, a new technique which improves sample efficiency.
Interpretable representation-learning models have achieved great progress for various types of time-series data. Taking the music domain as an example, tailored models [Ji et al., 2020] have been developed to learn pitch, timbre, melody contour, chord progression, texture, etc. from music audio. These human-interpretable representations have greatly improved the performance of generative algorithms in various music creation tasks, including inpainting [Wei et al., 2022], harmonization [Yi et al., 2022], (re-)arrangement, and performance rendering [Jeong et al., 2019].
However, most representation learning models still rely heavily on domain-specific knowledge. For example, to use pitch scales or instrument labels for learning pitch and timbre representations [Luo et al., 2020, 2019; Engel et al., 2020; Lin et al., 2021; Esling et al., 2018] and to use chords and rhythm labels for learning higher-level representations [Akama, 2019; Yang et al., 2019; Wang et al., 2020; Wei and Xia, 2021]. Such an approach is presumably very different from human learning; even without formal music training, we see that many people can learn pitch, a fundamental music concept, simply from the experience of listening to music. Hence, it remains an open question how to learn interpretable pitch factor using inductive biases that are more general. In other words, what general computational principle gives rise to the concept of pitch.
We see a similar issue in other domains. For instance, various computer-vision models [McCarthy and Ahmed, 2020; Trevithick and Yang, 2021; Mescheder et al., 2019; Riegler et al., 2017] can learn
3D representations of human faces or a particular scene by using domain knowledge (e.g., labelling of meshes and voxels, 3D convolution, etc.) But when these domain setups are absent, it remains a non-trivial task to learn the 3D location of a simple moving object in a self-supervised fashion.
Figure 1: An illustration of physical symmetry as our inductive bias.
Inspired by modern physics, we explore to use physical symmetry (i.e., symmetry of physical laws) as a weak self-consistency constraint for the learned latent  $z$  space of time-series data  $x$ . As indicated in Figure 1, this general inductive bias requires the learned prior model  $R$ , which is the induced physical law describing the temporal flow of the latent states, to be equivariant to a certain transformation  $S$  (e.g., translation or rotation). Formally,  $z_{t+1} = R(z_t)$  if and only if  $z_{t+1}^S = R(z_t^S)$ , where  $z^S = S(z)$ . In other words,  $R$  and  $S$  are commutable for  $z$ , i.e.,  $R(S(z)) = S(R(z))$ . Note that our equivariant assumption applies only to the latent  $z$  space. This is fundamentally different from most existing symmetry-informed models [Bronstein et al., 2021], in which the equivariant property also imposes assumptions on the raw data space.
Specifically, we design self-supervised learning with physical symmetry  $(\mathbf{SPS})^1$ , a method that adopts an encoder-decoder framework and applies physical symmetry to the prior model. We show that SPS learns a linear pitch factor (that agree with human music perception) from monophonic music audio without any domain-specific knowledge about pitch scales, f0, or harmonic series. The same methodology can be applied to the computer vision domain, learning 3D Cartesian space from monocular videos of a bouncing ball shot from a fixed perspective. In particular, we see four desired properties of SPS as a self-supervised algorithm for interpretability:
- Conciseness: SPS does not require contrastive samples, bach normalization, or a large batch size. We can even drop the Gaussian prior regularization, which is usually required to learn a meaningful latent space. (See Section 3.2 - Section 4.)
- Sample efficiency: SPS learns low-dimensional representations of a dynamic system in a very sample-efficient way. (See Section 4 - Section 5.2.)
- Robustness: Even with an incorrect symmetry assumption, SPS can still learn more interpretable representations than baselines. (See Section 5.3.)
- Extendability: SPS can be easily combined with other learning techniques. For example, if we further assume an extra global invariant style code, the model becomes a disentangled sequential autoencoder, capable of learning content-style disentanglement from temporal signals. (See appendix.)
2 Intuition
The idea of using physical symmetry for representation learning comes from modern physics. In classical physics, scientists usually first induce physical laws from observations and then discover symmetry properties of the law. (E.g., Newton's law of gravitation, which was induced from planetary orbits, is symmetric with respect to Galilean transformations.) In contrast, in modern physics, scientists often start from a symmetry assumption, based on which they derive the corresponding law and predict the properties (representations) of fundamental particles. (E.g., general relativity was developed based on a firm assumption of symmetry with respect to Lorentz transformations).
<sup>1</sup>The source code is publicly available at https://github.com/XuanjieLiu/Self-supervised-learning-via-Physical-Symmetry. The demo page is available at https://xuanjieliu.github.io/SPS_demo/
Analogously, we use physical symmetry as an inductive bias of our representation learning model, which helps us learn a regularised prior and an interpretable low-dim latent space. If it is a belief of many physicists that symmetry in physical law is a major design principle of nature, we regard symmetry in physical law as a general inductive bias of perception. In other words, if physical symmetry leads an AI agent to learn human-aligned concepts in a self-supervised fashion, we believe that it could also provide insights into the ways that human minds perceive the world.
The introduction of physical symmetry naturally leads to counterfactual representation augmentation, a novel learning technique which helps improve sample efficiency. Representation augmentation means to "imagine" extra pairs of  $z_{t}^{S}$  and  $z_{t + 1}^{S}$  as training samples for the prior model  $R$ . Through the lens of causality, this augmentation can be seen as a counterfactual inductive bias of the prediction on the representation level – what if the prior model makes predictions based on transformed latent codes? As indicated in Figure 1, it requires the prediction of the  $z$  sequence to be equivariant to certain group transformations,  $S$ . This regularisation also constrains the encoder and decoder indirectly through the prior model since the network is trained in an end-to-end fashion.
With physical symmetry, we aim to learn an interpretable low-dimensional representation  $z_{i}$  of each high-dimensional sample  $x_{i}$  from time-series  $\mathbf{x}_{1:T}$ . We focus on two problems in this paper: 1) to learn a 1D linear pitch factor of music notes from music audio, where each  $x_{i}$  is a spectrogram of a note, and 2) to learn 3D Cartesian location factors of a simple moving object (a bouncing ball) from its trajectory shot by a fixed, single camera, where each  $x_{i}$  is an image.
3.1 Model
Figure 2 shows the model design of SPS. 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}$ . Then it is fed into three branches. In the first branch (the green line),  $\mathbf{z}_{1:T}$  is decoded directly by the decoder  $D$  to reconstruct  $\mathbf{x}_{1:T}^{\prime}$ . In the second branch (the orange line),  $\mathbf{z}_{1:T}$  is passed through the prior model  $R$  to predict its next timestep,  $\hat{\mathbf{z}}_{2:T+1}$ , which is then decoded to reconstruct  $\hat{\mathbf{x}}_{2:T+1}$ . In the third branch (the blue line), we transform  $\mathbf{z}_{1:T}$  with  $S$ , pass it through  $R$ , and transform it back using the inverse transformation  $S^{-1}$  to predict another version of the next timestep  $\tilde{\mathbf{z}}_{2:T+1}$ , and finally decode it to  $\tilde{\mathbf{x}}_{2:T+1}$ . We get three outputs from the model:  $\mathbf{x}_{1:T}^{\prime}, \hat{\mathbf{x}}_{2:T+1}$ , and  $\tilde{\mathbf{x}}_{2:T+1}$ .
The underlying idea of physical symmetry is that the dynamics of latent factor and its transformed version follow the same physical law characterised by  $R$ . Therefore,  $\tilde{\mathbf{z}}$  and  $\hat{\mathbf{z}}$  should be close to each other and so are  $\tilde{\mathbf{x}}$  and  $\hat{\mathbf{x}}$ , assuming  $S$  is a proper transformation. This self-consistency constraint helps the network learn a more regularised latent space.
Figure 2: 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 fed into three different branches yielding 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 with respect 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}}$ .
The total loss contains four terms: reconstruction loss  $\mathcal{L}_{\mathrm{rec}}$ , prior prediction loss  $\mathcal{L}_{\mathrm{prior}}$ , symmetry-based loss  $\mathcal{L}_{\mathrm{sym}}$ , and KL divergence loss  $\mathcal{L}_{\mathrm{KLD}}$ . Formally,
$$
\mathcal {L} = \mathcal {L} _ {\text {r e c}} + \lambda_ {1} \mathcal {L} _ {\text {p r i o r}} + \lambda_ {2} \mathcal {L} _ {\text {s y m}} + \lambda_ {3} \mathcal {L} _ {\text {K L D}}, \tag {1}
$$
where  $\lambda_1, \lambda_2$  and  $\lambda_3$  are weighting parameters. By referring to the notations in section 3.1,
$$
\mathcal {L} _ {\text {r e c}} = \mathcal {L} _ {\mathrm {B C E}} \left(\mathbf {x} _ {1: T} ^ {\prime}, \mathbf {x} _ {1: T}\right) + \mathcal {L} _ {\mathrm {B C E}} \left(\hat {\mathbf {x}} _ {2: T}, \mathbf {x} _ {2: T}\right) + \mathcal {L} _ {\mathrm {B C E}} \left(\tilde {\mathbf {x}} _ {2: T}, \mathbf {x} _ {2: T}\right) \tag {2}
$$
$$
\mathcal {L} _ {\text {p r i o r}} = \ell_ {2} \left(\hat {\mathbf {z}} _ {2: T}, \mathbf {z} _ {2: T}\right), \tag {3}
$$
$$
\mathcal {L} _ {\mathrm {s y m}} = \ell_ {2} \left(\tilde {\mathbf {z}} _ {2: T}, \hat {\mathbf {z}} _ {2: T}\right) + \ell_ {2} \left(\tilde {\mathbf {z}} _ {2: T}, \mathbf {z} _ {2: T}\right). \tag {4}
$$
$\mathcal{L}_{\mathrm{KLD}}$  is the Kulback-Leibler divergence loss between the prior distribution of  $z_{i}$  and a standard Gaussian. Lastly, we build two versions of SPS:  $\mathrm{SPS}_{\mathrm{VAE}}$  and  $\mathrm{SPS}_{\mathrm{AE}}$ , with the latter replacing the VAE with an AE (and trivially doesn't have  $\mathcal{L}_{\mathrm{KLD}}$ ).
3.3 Symmetry-based counterfactual representation augmentation
During training,  $S$  is the counterfactual representation augmentation since it creates extra imaginary sequences of  $z$  (i.e., imaginary experience) to help train the prior. In practice, for each batch we apply  $K$  different transformations  $S_{1:K}$  to  $\mathbf{z}$  and yield  $K$  imaginary sequences. Thus, the two terms of symmetry-based loss can be specified as:
$$
\ell_ {2} \left(\tilde {\mathbf {z}} _ {2: T}, \mathbf {z} _ {2: T}\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}\right)\right), \mathbf {z} _ {2: T}\right) \right. \tag {5}
$$
$$
\ell_ {2} \left(\tilde {\mathbf {z}} _ {2: T}, \hat {\mathbf {z}} _ {2: T}\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}\right)\right), \hat {\mathbf {z}} _ {2: T}\right) \right. \tag {6}
$$
where the lower case  $k$  denotes the index of a specific transformation and we refer to  $K$  as the augmentation factor. Likewise, the last term of reconstruction loss can be specified as:
$$
\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}\right)\right)\right), \mathbf {x} _ {2: T}\right) \right. \tag {7}
$$
Each  $S$  applied to each sequence  $\mathbf{z}_{1:T}$  belongs to a certain group, and different groups are used for different problems. For the music problem, we assume  $\mathbf{z}_i$  be to 1D and use random  $S\in G\cong (\mathbb{R},+)$ . In other words, we add a random scalar to the latent codes. As for the video problem, we assume  $\mathbf{z}_i$  be to 3D and use random  $S\in G\cong (\mathbb{R}^2, + )\times \mathrm{SO}(2)$ . In other words, random rotation and translation are applied on two dimensions of  $\mathbf{z}_i$ .
We test SPS under two modalities of temporal signals: music (section 4.1) and video (section 4.2). Each model is executed with 10 random initialisations, and evaluated on the test set. The highlight of this section is that SPS effectively learns interpretable low-dimensional factors that align with human perceptions. Also, by utilizing small training sets, we show the high sampling efficiency of our model. In the appendix, we further show that SPS also maintains accuracy in both reconstruction and prediction tasks (section A.2). Additionally, we present supplementary results trained on more complicated datasets and a more advanced configuration of our model, called  $\mathrm{SPS + }$ , which enables content-style disentanglement in addition to interpretable learning. The findings from these more complex scenarios align closely with those observed in the simpler cases presented in this section.
4.1 Learning linear pitch factor from music audio
4.1.1 Dataset and training setups
We synthesise a training dataset of 27 audio clips, each containing 15 notes in major scales with the first 8 notes ascending and the last 8 notes descending. We vary the starting pitch by integer numbers of MIDI pitch such that every MIDI pitch in the range A#4 to C7 is present in the training set. Only the accordion is used to generate the clips. For each clip in the training set, we uniformly randomly shift its pitch upwards by a decimal between 0 and 1, in this way generate the test set for evaluation.
We convert each audio clip into a sequence of image segments for processing. First, we run STFT (with sample rate  $= 16000 / s$ , window length  $= 1024$ , hop length  $= 512$ , with no padding and no logarithmic frequency scale) over each audio clip to obtain a power spectrogram. After normalising the energy to the range [0, 1], we slice the power spectrogram into fifteen image segments, each containing one note. The CNN encoder, in each timestep, takes one segment as input. For the latent space, we assume  $\mathbf{z}_i \in \mathcal{R}$  and sample counterfactual representation augmentation  $S \sim \mathcal{U}([-1,1])$  where  $S \in G \cong (\mathbb{R}, +)$ . Note this assumption does not indicate any domain-specific inductive bias of music, such as the logarithmic relationship between pitch and frequency or the relationship between F0 and harmonics.
4.1.2 Results on interpretable pitch space
Figure 3: A visualisation of the mapping between the 1D learned factor  $\mathbf{z}$  and the true pitch, in which a straight lines indicate a better result. In the upper row, models encode notes in the test set to  $\mathbf{z}$ . The  $x$  axis shows the true pitch and the  $y$  axis shows the learned pitch factor. In the lower row, the  $x$  axis traverses the  $\mathbf{z}$  space. The models decode  $\mathbf{z}$  to audio clips. We apply YIN to the audio clips to detect the pitch, which is shown by the  $y$  axis. In both rows, a linear, noiseless mapping is ideal, and our method performs the best.
Figure 3 demonstrates that the 1D pitch factor learned by our model exhibits a linear relationship with the conventional numerical ordering used to represent pitch by humans (e.g. MIDI pitch numbers). The plot shows the mappings of two tasks and six models. In the embedding task (the first row), the  $x$ -axis is the true pitch and the  $y$ -axis is embedded  $\mathbf{z}$ . In the synthesis task (the second row), the  $x$ -axis is  $\mathbf{z}$  and the  $y$ -axis is the detected pitch (by YIN algorithm, a standard pitch-estimation method by [De Cheveigné and Kawahara, 2002]) of decoded (synthesised) notes. The first two models are SPS based on VAE and AE, respectively, trained with counterfactual representation augmentation factor  $K = 4$ . The third and fourth models are trained without constraints of physical symmetry ( $K = 0$ ), serving as our ablations. The fifth one is a vanilla  $\beta$ -VAE, trained only to reconstruct, not to predict. The last one is SPICE [Gfeller et al., 2020], a SOTA unsupervised pitch estimator with strong domain knowledge on how pitch linearity is reflected in log-frequency spectrograms. As the figure shows, 1) without explicit knowledge of pitch, our model learns a more interpretable pitch factor than  $\beta$ -VAE, and the result is comparable to SPICE, and 2) without the Gaussian prior assumption of latent variable distribution, our model  $\mathrm{SPS}_{\mathrm{AE}}$  also learns a continuous representation space.
Table 1: The linearity of learned pitch factor and synthesized sound pitch evaluated by  $R^2$ .
Methods	Learned factor R2↑	Synthesis R2↑
SPSVAE, K=4 (Ours)	0.999±0.001	0.986±0.025
SPSAE, K=4 (Ours)	0.998±0.001	0.986±0.025
SPSVAE, K=0 (Ablation)	0.997±0.002	0.910±0.040
SPSAE, K=0 (Ablation)	0.993±0.006	0.832±0.129
β-VAE	0.772±0.333	0.534±0.275
SPICE	1.000	N/A
Table 1 shows a more quantitative analysis using  $R^2$  as the metric to evaluate the linearity of the pitch against  $z$  mapping from the encoder and the decoder. All models except SPICE are trained with 10 random initializations.
4.2 Learning object 3D coordinates from videos of a moving object
4.2.1 Dataset and training setups
We run physical simulations of a bouncing ball in a 3D space. The ball is randomly thrown and affected by gravity and the bouncing force (elastic force). A fixed camera records a 20-frame video of each 4-second simulation to obtain one trajectory (see Figure 4). The ball's size, gravity, and proportion of energy loss per bounce are constant across all trajectories. For each trajectory, the ball's initial location and initial velocity are randomly sampled. We utilize 512 trajectories for training, and an additional 512 trajectories for evaluation.
For the latent space, we set the dimension of the latent space to 3, but only constrain 2 of them by augmenting the representations with  $S \in G \cong (\mathbb{R}^2, +) \times \mathrm{SO}(2)$ . Those two dimensions are intended to span the horizontal plane. The third one, which is the unaugmented latent dimension, is intended to encode the vertical height.
Figure 4: Two example trajectories from the bouncing ball dataset.
4.2.2 Results on interpretable 3D representation
Figure 5 visually evaluates the interpretability of the learned location factors by traversing the  $z$  space, one dimension at a time, and using the learned decoder to synthesise images. If the learned factors are linear w.r.t. the 3D Cartesian coordinates, the synthesised ball should display linear motions as we change  $z$  linearly. In brief, SPS learns an a more interpretable and linear  $z$  space. Here, subplot (a) depicts the results of  $\mathrm{SPS}_{\mathrm{VAE}}$  with  $K = 4$ . We see that the un-augmented dimension,  $z_{2}$ , controls the height of the ball, while  $z_{1}$  and  $z_{3}$  move the ball along the horizontal (ground) plane. Each axis is much more linear than in (b) and (c). Subplot (b) evaluates  $\mathrm{SPS}_{\mathrm{VAE}}$  with counterfactual representation augmentation  $K = 0$ , essentially turning off SPS. As  $z_{i}$  varies, the ball seems to travel along curves in the 3D space, showing the ablation learns some continuity w.r.t. the 3D space, but is obviously far from linear. In (c), the  $\beta$ -VAE fails to give consistent meaning to any axis.
Table 2 further shows quantitative evaluations on the linearity of the learned location factor, in which we see that SPS outperforms other models by a large margin. To measure linearity, we fit a linear regression from  $z$  to the true 3D location over the test set and then compute the Mean Square Errors (MSE). Therefore, a smaller MSE indicates a better fit. To give an intuitive example, the MSEs of (a), (b) and (c) in Figure 5 are 0.09, 0.58 and 0.62 respectively. Here, we also include the results of  $\mathrm{SPS}_{\mathrm{AE}}$ . Very similar to the music experiment in session 4.1, we again see that even without the Gaussian prior assumption, our model  $\mathrm{SPS}_{\mathrm{AE}}$  learns an interpretable latent space comparable to  $\mathrm{SPS}_{\mathrm{VAE}}$ .
(a)  $\mathrm{SPS}_{\mathrm{VAE}}$ $K = 4$  (Ours)
(b)  $\mathrm{SPS}_{\mathrm{VAE}}$ $K = 0$  (Ablation)
(c)  $\beta$ -VAE (Baseline)
Figure 5: A visualisation of latent-space traversal performed on three models: (a) ours, (b) ablation, and (c) baseline, in which we see (a) achieves better linearity and interpretability. Here, row  $i$  shows the generated images when changing  $z_{i}$  and keeping  $z_{\neq i} = 0$ , where the  $x$  axis varies  $z_{i}$  from  $-2\sigma_{z}$  to  $+2\sigma_{z}$ . We center and normalise  $z$ , so that the latent space from different runs is aligned for fair comparison. Specifically, in (a), changing  $z_{2}$  controls the ball's height, and changing  $z_{1}, z_{3}$  moves the ball parallel to the ground plane. In contrast, the behavior in (b) and (c) are less interpretable.
Table 2: Linear fits between the true location and the learned location factor. We run the encoder on the test set to obtain data pairs in the form of (location factor, true coordinates). We then run a linear fit on the data pairs to evaluate factor interpretability. Two outliers are removed from the 50 runs.
Method	x axis MSE ↓	y axis MSE ↓	z axis MSE ↓	MSE ↓
SPSVAE, K=4 (Ours)	0.11 ± 0.09	0.31 ± 0.34	0.26 ± 0.34	0.26 ± 0.30
SPSAE, K=4 (Ours)	0.13 ± 0.07	0.39 ± 0.33	0.21 ± 0.17	0.24 ± 0.17
SPSVAE, K=0 (Ablation)	0.33 ± 0.10	0.80 ± 0.18	0.75 ± 0.17	0.62 ± 0.14
SPSAE, K=0 (Ablation)	0.26 ± 0.09	0.44 ± 0.27	0.55 ± 0.17	0.42 ± 0.15
β-VAE	0.36 ± 0.03	0.70 ± 0.01	0.68 ± 0.03	0.58 ± 0.01