Learning the RoPEs: Better 2D and 3D Position Encodings with STRING

We would like to sincerely thank the Reviewer for the comments. We provide detailed answers below.

$**Additional tests, e.g. 3D-based imitation learning methods from [1,2]**$:

Thank you very much for the comment. The main goal of the 3D bi-arm KUKA experiments was to combine the best techniques to obtain new SOTA. We have already established STRING as providing gains on the top of the regular 3D algorithms, that in turn improved upon their 2D counterparts in Sec. 4.2.2., where we focused on 3D object localization (see: Table 4). We explicitly describe this goal, putting our plan for bi-arm KUKA experiments in the context of our previous findings, at the beginning of Sec. 4.4 (l.372-375: $*“Establishing STRING as superior to other methods on previous tasks, ...”*$).

Having said that, we still included in the bi-arm KUKA section 3D-based imitation learning methods that do not leverage STRING. We trained via imitation learning a policy directly using depth to produce normal maps (paragraph: $*“Implicit depth via normal maps”*$ at the beginning of Sec. 4.4.2). This is based on well established methods for incorporating depth into robotics tasks from [1]. As we show in Fig. 5, STRING provides ~$**11**$% improvement accuracy-wise, as compared to that baseline.

We would like to note one more thing. The 2D policies we compared against in the bi-arm KUKA section were strong enough to outperform various 3D counterparts given similar computational budget (e.g. point-cloud-based methods that were used in the first paper pointed by the Reviewer; we will explicitly clarify it in the camera-ready version). This was the case since point cloud (PC) approaches suffer from the very high computational footprint. The average number of points in the bi-arm KUKA PCs was of order 5K+. Training/deploying policies directly using all the points was infeasible and various tokenization techniques to reduce the number of 3D tokens had a detrimental effect on the performance. Furthermore, base 2D policies applied in that section used vision encoders pre-trained on massive amount of vision data. Trained-from-scratch 3D encoders were not a match for them. Thus it was important for us to build on the 2D policies in that section. The approach we used there with the positional encoding mechanism in the form of STRING applied on RGB-D images via the 3D patch-lifting (see paragraph: $*“Lifting patches to 3D for STRING”*$, p.7) enabled us to successfully do it. First of all: it did not require working with PCs and retained the same number of patches as its 2D counterpart. That was critical for computational efficiency. Secondly, it could easily reuse 2D-pretrained checkpoints, leveraging vision input pre-processing from the 2D baseline as is, only to add depth signal to modulate attention computation. The STRING variant re-used all pre-trained parameters of its 2D counterpart, introducing a negligible amount of additional 3D-specific trainable parameters. Finally, STRING variant could be set up in such a way that at the very beginning of its training it was equivalent to the 2D variant. This was achieved by setting up a depth scaling parameter for the patch-depth (obtained by averaging over depth values of its constituting pixels) as zero at the beginning of training. We commented on some of these in the paper (e.g. a paragraph “From 2D to 3D with STRING”, p. 8), but will provide more detailed discussion in the camera-ready version.

We also would like to note that we added several additional experiments for the rebuttal, confirming the effectiveness of STRING:

1. We re-ran the evaluations for the $*BowlOnRack*$ task (we chose one task, given rebuttal period time constraints) from the Aloha-sim portfolio, by increasing the number of trials $**from 10 to 100**$. The ranking of methods in that setting stays exactly the same as reported before with STRING outperforming others.
2. We run experiments for 2D and 3D object detection using the SoViT-400m models which surpass ViT-H and are competitive in performance with ViT-G. As compared to baseline, STRING provides >$**4.7**$% relative improvement on COCO AP and >$**5.1**$% relative improvement on LVIS AP and improves also upon regular RoPE in the OWL-ViT-2D setting. For 3D object detection with SoViT-400m models, using STRING led to additional accuracy gains as compared to the baseline ($**1.2**$ of the baseline improvement coming from using a larger model).
3. We also conducted additional scaling experiments, where we increased the image resolution from 256 x 256 to 512 x 512. STRING maintained an advantage over the baseline for the object detection task from Sec. 4.2.2 (3D bounding boxes): $**6**$%+ for the 2D variant and almost $**10**$% for the 3D variant.

$\underline{References}$:

[1] Early or late fusion matters: Efficient rgb-d fusion in vision transformers for 3d object recognition; Tziafas et al.; IROS 2023.

We would like to sincerely thank the Reviewer for the comments. We provide detailed answers below.

$**Discussion on trade-offs regarding using Cayley-STRING and Circulant-STRING**$:

Thank you very much for an excellent comment. With token dimensionality per head $d_{\mathrm{head}}$, Cayley-STRING introduces $d_{\mathrm{head}}(d_{\mathrm{head}} - 1) / 2+ d_{\mathrm{head}} / 2$ parameters per head while Circulant-STRING adds $d_{\mathrm{head}}$ parameters per head. This is the case since Cayley-STRING trains a full skew-symmetric matrix and d / 2 frequencies, while Circulant-STRING uses a learnable circulant matrix which is fully determined by its first row. Furthermore, multiplication with Cayley-STRING matrices takes quadratic time in $d_{head}$ during training, whereas multiplication with Circulant-STRING matrices takes time $O(d_{\mathrm{head}} \log(d_{\mathrm{head}}))$ via Fast Fourier Transform (FFT). During inference, quadratic cost can be in principle absorbed into existing q/k projections at no extra cost (for the vanilla attention mechanism). Both approaches improve upon regular RoPE, yet Cayley-STRING in general leads to larger improvements. Thus we have here a classic quality-speed tradeoff. For applications with strict memory constraints, we recommend to use Circulant-STRING due to its very compact computational footprint, whereas for other applications Cayley-STRING is recommended. We will add the above discussion to the camera-ready version of the paper.

===========================================================================================================================================================

$**The impact of discretization in the sensory data on the performance of the system**:$

Thank you for your question. We conducted our experiments with images of various resolutions, seeing consistent gains across the board, as compared to the baselines. The statement that STRING is effective even if the sensory data is not accurate (due to discretization or noise) is also supported by our 3D experiments with the bi-arm KUKA robot. In all those experiments, depth sensors provide a relatively noisy signal, yet STRING is capable of leveraging depth sensors to provide $**8**$%+ accuracy improvements in regular in-distribution tasks and from $**15**$% to $**42**$% improvements for out-of-distribution tasks. We will add this discussion to the camera-ready version of the paper.

We would like to sincerely thank the Reviewer for the comments. We provide detailed answers below.

$**Theorem 3.5**$:

The proof relies on the fact that the computation of $exp(\mathbf{M})$ for a given matrix $\mathbf{M}$ can be re-written as: $\exp(\mathbf{M}) = \Sigma \exp(\mathbf{D}) \Sigma^{-1}$ for a given factorization: $\mathbf{M}=\Sigma \mathbf{D} \Sigma^{-1}$ (from the Taylor-series definition of matrix exponentiation). The next observation is that, for the matrices of interest in the proof, the factorization can be done efficiently with $\Sigma$ and $\mathbf{D}$ being: (1) a DFT matrix and (2) a diagonal matrix with the main diagonal given as the action of $\Sigma$ on an easy-to-compute vector, respectively. The final observation is that the DFT matrix supports fast matrix-vector multiplication (via the Fast Fourier Transform) and the diagonal matrix exponentiation is obtained by element-wise exponentiation. We will clarify it in the final version of the paper by providing a sketch of the proof in the main body of the paper.

Thank you very much for an excellent question regarding the backpropagation pass. Even though the analysis in the paper was made for the forward pass, it can be extended to the backward pass. This comes from the fact that the computation of the gradient in the backpropagation over networks leveraging Circulant-STRING matrices also involves multiplications with DFT matrices. The analysis of the backpropagation is more technical, but basically follows steps from Sec. 4.1 of [1]. We will provide a paragraph to discuss computational gains in backpropagation in the camera-ready version of the paper.

$\underline{References}$:

[1] An Exploration of Parameter Redundancy in Deep Networks with Circulant Projections; Cheng et al.; ICCV 2015.

===========================================================================================================================================================

$**Statistics or discussions on the number of parameters, memory consumption, or training/inference time**$:

Thank you very much for the comment. STRING introduces a negligible amount of extra trainable parameters. With token dimensionality per head $d_{\mathrm{head}}$, Cayley-STRING introduces $d_{\mathrm{head}}(d_{\mathrm{head}} - 1) / 2 + d_{\mathrm{head}} / 2$ parameters per head while Circulant-STRING adds $d_{\mathrm{head}}$ parameters per head. Thus for instance for the: (1) Cayley-STRING variant with no parameter sharing across heads, (2) Cayley-STRING with parameters shared across heads and (3) Circulant-STRING with no parameter sharing across heads, parameter count increase is marginal: (1) 0.34%, (2) 0.028% and (3) 0.016% respectively. In terms of training time and memory footprint, the performance of Cayley-STRING and Circulant-STRING is similar to that of the RoPE-M model (~\pm 7 \%). We will include these numbers in the supplementary material of the camera-ready version of the paper.

===========================================================================================================================================================

$**Theorem 3.3 and 3.4**$:

We thank the Reviewer for the comment, and will ensure that our notation is consistent and clear. Theorem 3.3 shows that RoPE is a special case of STRING for a particular choice of an antisymmetric generator. Meanwhile, Theorem 3.4 shows that RoPE is a special case of STRING with basis transformation matrix P = I. Both theoretical results show that our new method is more general than that previous SOTA, but from complementary perspectives. As shown in the paper, this increased generality unlocks better performance in experiments. We will make sure this is clear in the text and highlight the differences between these important, closely-related results.

===========================================================================================================================================================

$**Clearer introduction to the notation at the beginning**$:

Thank you for the question. $d$ refers to the dimensionality of the latent representations in the Transformer, i.e. the length of queries, keys and values. Meanwhile, $d_c$ refers to the dimensionality of the position coordinates encoded using STRING, i.e. 1 for sequences (text), 2 for 2D data (images), and 3 for 3D data (images with depth). We will make sure this is clear in the text.

We would like to sincerely thank the Reviewer for the comments. We provide detailed answers below.

$**Increasing the number of trials in testing (currently 10 trials / task for Aloha-sim and 50 trials for bi-arm KUKA)**:$

We sincerely thank the Reviewer for the comment. We would like to note that 10 trials / task is considered normal for the evaluation of the robotics controllers, both in sim and on hardware (e.g. [1]: 5 repetitions as stated in Sec 6.1, $*"This is repeated **5** times"*$ and [2]: 10 repetitions in $*“We ran each of the learned controller **10** times.”*$ at the beginning of Sec. B (“Results and Analysis”)).

In fact it is also considered normal for testing several regular machine learning algorithms. For instance, most papers evaluating Random Fourier Features and related methods, such as [3] ($*“All experiments are repeated **ten** times”*$, p.7), [4] (caption of Fig. 2 $*“**10** independent trials are executed”*$) that usually can be quickly evaluated with very limited computational resources, apply 10 repetitions to calculate empirical MSEs.

Having said that, following Reviewer’s comment for the rebuttal, we re-ran the evaluations for the $*BowlOnRack*$ task (we chose one task, given rebuttal period time constraints) from the Aloha simulation portfolio, by increasing the number of trials $**from 10 to 100**$. The ranking of different methods in that setting $**stays exactly the same**$ as reported before: 1. STRING (83% accuracy) 2. regular ViT baseline (80% accuracy) and 3. RoPE (75%). We also would like to emphasize that all reported KUKA experiments are $**on-robot**$ (and thus much more time-consuming) and in the paper we use many more trials for the corresponding evals.

$\underline{References}$:

[1] Scalable Deep Reinforcement Learning for Vision-Based Robotic Manipulation; Kalashnikov et al.; CoRL 2018; $**Best Systems Paper, ~1.8K citations**$.

[2] Deep Spatial Autoencoders for Visuomotor Learning; Finn et al., ICRA 2016; $**764 citations**$.

[3] Nystrom Method vs Random Fourier Features: A Theoretical and Empirical Comparison; Yang et al.; NeurIPS 2012; 450+ citations).

[4] Quasi-Monte Carlo Feature Maps for Shift-Invariant Kernels; Avron et al.; ICML 2014

===========================================================================================================================================================

$**Scaling Experiments to show that STRING still performs well**:$

Thank you very much for an interesting question about scaling to larger models. Following the Reviewer’s suggestion, we have run such experiments using the SoViT-400m class of models which surpass ViT-H and are competitive in performance with ViT-G [1]. As compared to baseline, STRING provides $**>4.7**$% relative improvement on COCO AP and $**>5.1**$% relative improvement on LVIS AP and improves also upon regular RoPE in the OWL-ViT-2D setting. For the task of 3D object detection with SoViT-400m models, using STRING led to additional accuracy gains as compared to the baseline ($**1.2**$ of the baseline improvement coming from using a larger model). Thus, we tested that STRING does perform well also on larger models.

We have also conducted additional scaling experiments, where we increased the resolution of images from 256 x 256 to 512 x 512. In the higher-resolution setting, STRING maintained an advantage over the baseline for the object detection task from Sec. 4.2.2 (3D bounding boxes): $**6**$%+ for the 2D variant and almost $**10**$% for the 3D variant.

We will include all these results in the camera-ready version of the paper.

$\underline{References}$:

[1] Getting ViT in Shape: Scaling Laws for Compute-Optimal Model Design; Alabdulmohsin et al.; NeurIPS 2023

===========================================================================================================================================================

$**Improvements Over Baselines**$:

Thank you very much for your question. The STRING variant provides highest average accuracy over all Aloha-sim tasks. However, what is as important and clearly seen in Fig. 4, it also provides significant convergence improvements in training, achieving accuracy level of regular RoPE in $**1.6x**$ shorter time. We explicitly commented on that in l.367-368. This clearly distinguished STRING from regular RoPE. Furthermore, our 3D object detection experiments clearly show improvements coming from STRING (from approximately $**1**$% to $**1.5**$%), in both: 2D and 3D settings, as compared to $**the best**$ RoPE variants (Table 4). These improvements are actually visible to the naked eye. We show it in Fig. 2, where predicted small 3D boxes (in blue) are clearly misaligned with ground truth bounding boxes (in green) for baseline and RoPE, while pretty well aligned for STRING.