Exploring and Improving the Spatial Reasoning Abilities of Large Language Models

As discussed in [15] the ability of an LLM to pattern match is driven by in-context learning on the provided numerical tokens, which can be formulated as the problem of using the context $s_{1:k}=(s_{1},...,s_{k})$, where each $s_{i}$ is a symbol and using it to autoregressively predict $s_{k+1}$ by using the factorized conditional probability $p(s_{1:k})=\Pi_{i=1}^{n}p(s_{i}|s_{1},...,s_{i-1})$. Usual in-context learning examples segment the prompt into continuations of multiple examples, each a variable length sequence: $x_{1:k}=(x_{1},...,x_{k})$ where each $x_{i}=(s_{1}^{i},s_{2}^{i},...,s_{m^{i}}^{i})$.

We adapt this paradigm for directional labeling by having $x_{1}$ state the model’s expertise in spatial analysis and prompt it to generate direction labels for a newly provided sequence given the examples. Then each $x_{i}$ from $i=2$ onwards is an example, an input-output pair $D_{i};L_{i}$ where $D_{i}$ is the sequence of symbol aggregates $(d_{1}^{i},d_{2}^{i},...d_{j}^{i})$ of length $j$, with each $d_{k}^{i}$ further being separated out into the symbols $(d_{k_{x}}^{i},d_{k_{y}}^{i})$, representing a coordinate in 2D Cartesian space. $L_{i}$ is similarly a sequence of words $(l_{1}^{i},l_{2}^{i},...l_{j-1}^{i})$ of length $j-1$, where each $l_{k}^{i}$ is a word representing the direction that describes the movement from $d_{k}^{i}$ to $d_{k+1}^{i}$ and is one of [left, right, up, down]. There are thus $j$ points and $j-1$ segment labels (see Fig.  1).

Using a similar prompt framework to the one described above, $x_{1}$ states the model’s expertise in spatial analysis and prompts it to identify the overall shape of the movement represented by the list of 2D coordinates; for example, "moves along a path that mirrors the pattern of a checkmark" or "moves in a circular path". Since the dataset size is fairly limited, no in-context learning is applied and the model is evaluated zero-shot by specifying $x_{2}$ as the input sequence of $(x,y)$ coordinates $(d_{1}^{i},d_{2}^{i},...d_{j}^{i})$ (see Fig.  1).

To provide a baseline for the various prompting mechanisms probed, we initially implement zero-shot prompting, in which no exemplars are imparted to the model. Therefore, analogous to 2D experiments, $x_{1}$ states the model’s expertise in spatial analysis and prompts it to classify the type of motion exemplified by the sequence into one of $N$ categories, but there are no further sequences (see Fig.  2). For our set of experiments, we designate three classes that correspond to meaningfully and spatially disparate motions - lifting, rotating and sliding.

In-context Learning (ICL) [1] supplies a few samples that the model can generalize from. $x_{1}$ declares that the model that it is an expert in 3D trajectory labeling and prompts the model to classify the type of motion exemplified by the sequence into one of $N$ categories. Then each $x_{i}$ from $i=2$ onwards is an input-output pair $D_{i};l_{i}$ where $D_{i}$ is the sequence of symbol aggregates $(d_{1}^{i},d_{2}^{i},...d_{j}^{i})$, with each one representing a coordinate in 3D Cartesian space $(d_{k_{x}}^{i},d_{k_{y}}^{i}),d_{k_{z}}^{i})$. $l_{i}$ is a single word describing the motion and belongs to one of the $N$ classes (see Fig.  2).

Chain-of-Thought Prompting (CoT) [21], in which the few-shot examples are augmented by a step-by-step reasoning, has superseded In-context Learning on an array of textual and mathematical reasoning tasks. We extend it to our task as follows. The starting guideline for $x_{1}$ is the same as in-context learning, with the declaration that the model is an expert in 3D trajectory labeling and a prompt for the model to classify the type of motion exemplified by the sequence into one of $N$ categories. However in CoT, successive pairs $x_{2},x_{3},...,x_{2t},x_{2t+1}$ for some $t$, represent the few-shot examples. The even $x_{2t}$ is the representative sequence of 3D coordinates and the odd $x_{2t+1}$ is its associated answer, which is the motion-type label and accompanying reasoning steps (such as the fact that back-and-forth changes in the $x$ and $y$ coordinates can hint at a rotating motion, see Fig.  2).

Anticipating the challenge with generalizing from irregular examples and bolstered by the model’s performance on simpler 2D data, we also propose a new method of prompting called "Spatial Prefix-Prompting" (SPP). The method draws from a prior selection of fixed questions that instigate the model to first ponder a tangentially related spatial problem (e.g. identifying the single direction an object is moving in or checking whether a point is in the center of a circle), and then use the "knowledge gained" to answer an adjacent question, such as labeling a new, more complex 3D trajectory (see Fig.  2). This technique does not necessitate the more intensive CoT-style curation of step-by-step examples, and we hypothesize that it may build upon the more fundamental spatial concepts a model is trained on to generalize better than few-shot learning (wherein the selected examples may not be representative of all the trajectories).

There aren’t many datasets for elementary 2D shapes and directions, and given the ease of generating such data, we decided to autogenerate the datasets. For the direction labeling task, a dataset of size 30 is generated with 10 short-horizon sequences of 2D coordinates (of length 6-8 segments), 10 long-horizon long sequences (length 35 - 40), and 10 short floating-point sequences using Python’s NumPy package, with each segment’s size and the directions randomly chosen based on a fixed seed. We experiment with both 1) scaling the values to integers between [0, 100] (to use fewer tokens to represent a single number) 2) scaling the values to double-digit fractions between [0, 100] (to test whether the higher token representation translates to higher precision). Only Zero-shot and In-context Learning are applied for Shape and Direction Labeling respectively. For the shape identification, we use some previously collected hand-gesture data in which human demonstrators sign a variety of shapes, including circles and check marks, and 2D positions of the finger are recorded; the dataset tests the model on the inherent noise from human demonstrations. A subset of the dataset is "cleaned" for use by having an expert to remove extraneous points that don’t belong to the shape, and it is normalized to the range [0, 100] (both integer and floating point). The size of the dataset is 13.

We use the CALVIN benchmark [13], a dataset for learning long-horizon language-conditioned tasks for robotics. It includes 3D end-effector positions (can be extracted from the low-dimensional state vector) and associated language descriptions of the action the robot is attempting to complete. Due to time and resource constraints from the human evaluations, we select only a small subset of the CALVIN dataset (30 samples) and three disparate subtasks ("rotate", "lift", "slide"). The trajectories are often complex and may not intuitively always resemble the action being completed, with many extraneous movements (see Fig.  1). Therefore, we also create a version of this dataset called "CALVIN-Cleaned" in which a human annotator extracts the parts of the trajectory that match with the specific action (e.g. only the lifting portion, see Fig.  1). Note, the CALVIN-Cleaned dataset retains the original "rotate" trajectories, as the task description is linked to the back-and-forth changes in the entirety of the motion, not any particular subsection. Finally, the dataset is normalized to the range $[0,300]\in\mathbb{Z}$ to increase the granularity of the trajectory but optimize for token conservation due to the long-range of many CALVIN trajectories (upwards of 50 - 100 points).

In order to demonstrate the efficacy of the Spatial Prefix-Prompting mechanism beyond just 3D trajectory classification, we also run experiments with Llama-2-7B on the entire test set of the SpartQA dataset [16] (of size 510 instances), a textual QA benchmark for deeper spatial reasoning questions of four types: find relation (FR), find blocks (FB), choose object (CO) (see Fig.  1).