SSTP: Efficient Sample Selection for Trajectory Prediction

## Abstract

Abstract Trajectory prediction is a core task in autonomous driving.
However, training advanced trajectory prediction models on existing large-scale datasets is both time-consuming and computationally expensive.
More critically, these datasets are highly imbalanced in scenario density, with normal driving scenes (low-moderate traffic) overwhelmingly dominating the datasets, while high-density and safety-critical cases are underrepresented.
As a result, models tend to overfit low/moderate-density scenarios and perform poorly in high-density scenarios.
To address these challenges, we propose the SSTP framework, which constructs a compact yet density-balanced dataset tailored to trajectory prediction.
SSTP consists of two main stages: (1) Extraction , where a baseline model is pretrained for a few epochs to obtain stable gradient estimates, and the dataset is partitioned by scenario density.
(2) Selection , where gradient-based scores and a submodular objective select representative samples within each density category, while biased sampling emphasizes rare high-density interactions to avoid dominance by low-density cases.
This approach significantly reduces the dataset size and mitigates scenario imbalance, without sacrificing prediction accuracy.
Experiments on the Argoverse 1 and Argoverse 2 datasets with recent state-of-the-art models show that SSTP achieves comparable performance to full-dataset training using only half the data while delivering substantial improvements in high-density traffic scenes and significantly reducing training time. Robust trajectory prediction depends not only on data scale but also on balancing scene density to ensure reliable performance under complex multi agent interactions.

## 1 Introduction

Trajectory prediction aims to predict the future locations of agents conditioned on their past observations, which plays a key role in the domain of autonomous driving.
This task is essential yet challenging due to the complex uncertain driving situations.
With rapid developments in deep learning, various methods (Zhou et al., 2022c; Feng et al., 2023; Chai et al., 2019; Gu et al., 2021; Zhang et al., 2021; Zhao et al., 2021; Ngiam et al., 2021) have been proposed with promising trajectory prediction performance.
Meanwhile, more large-scale realistic datasets (Caesar et al., 2020; Zhan et al., 2019; Chang et al., 2019; Wilson et al., 2023; Sun et al., 2020) have been released by research institutes and self-driving companies, which further push the boundary of this task.

However, one common issue is that training these data-driven methods requires enormous computational resources and is time-consuming due to the large scale of the datasets.
For example, the recent state-of-the-art method MTR (Ngiam et al., 2021) has over 66 million model parameters, and the Waymo Open Motion Dataset (WOMD) (Sun et al., 2020) has over 2.2 million trajectory samples.
Training the complete model on this dataset requires a substantial amount of GPU hours, posing a significant computational burden.

Q1: To what extent does training on massive trajectory datasets improve performance, considering their high computational cost? Q2: How can we effectively reduce the training data volume without significantly compromising model accuracy?

Prompted by these questions, we take a deep dive into recent trajectory prediction benchmarks and reveal one finding: the imbalance in the dataset.
Specifically, the majority of scenarios involve only a limited number of interacting agents, whereas high-density interaction scenarios are significantly underrepresented.
This imbalance is evident in both Argoverse 1 (Chang et al., 2019) and Argoverse 2 (Wilson et al., 2023) (Figure [2](https://arxiv.org/html/2409.17385v3#S1.F2)), where most cases are low-density. Yet such high-density cases are precisely the most safety-critical for autonomous driving, since they involve complex multi-agent interactions where small prediction errors may directly compromise driving safety.
Similar long-tail phenomena have been observed in the literature (Makansi et al., 2021; Wang et al., 2023; Pourkeshavarz et al., 2023; Chen et al., 2024).
From a conventional machine learning perspective, standard tasks typically evaluate model performance using the average accuracy across all samples.
However, in autonomous driving trajectory prediction, this metric inherently biases models toward driving scenarios with abundant data, leading to suboptimal performance in data-scarce scenarios. For safe driving, an ideal trajectory predictor should exhibit robust performance across both data-rich and data-scarce scenarios.
However, due to the imbalanced distribution of driving scenarios, state-of-the-art baselines such as HiVT (Zhou et al., 2022c) and HPNet (Tang et al., 2024) struggle to maintain consistent performance across diverse densities, as illustrated in Figure [1](https://arxiv.org/html/2409.17385v3#S1.F1).

To address this issue, we introduce the Sample Selection for Trajectory Prediction (SSTP) framework, the first framework designed to construct a compact and balanced dataset for trajectory prediction.
SSTP consists of two main stages: extraction and selection.
In the extraction stage, a baseline model is pretrained on the full dataset for a few epochs to obtain stable gradient estimates, and the dataset is partitioned by scenario density, measured by the number of agents. This design mitigates the drawbacks of traditional trajectory prediction training:
in rare high-density scenarios, the scarcity of training data leads to systematic under-training.
For Transformer-based architectures, this issue is particularly pronounced. High-density scenarios require a model to accurately grasp the complex agent–agent dependencies, yet are significantly underrepresented in the training dataset. Consequently, the model does not have sufficient chances to learn such dependencies, leading to poor performance at inference.
Moreover, in standard Transformer self-attention, parameters are shared across samples and rely on broad coverage for generalization.
When high-density scenarios are scarce, updates are dominated by low-density scenarios, further reinforcing systematic under-training.
In the selection stage, for each density category, we calculate gradient-based influence scores for every sample, then apply a submodular objective that selects a subset that captures the most representative cases while reducing redundancy. Across partitions, we employ biased sampling to explicitly upweight rare but critical high-density scenarios, preventing dominance by the majority of low-density scenarios.
As shown in Figure [1](https://arxiv.org/html/2409.17385v3#S1.F1), SSTP
reduces data volume while maintaining performance, and achieves even better results in high-density scenarios.

Figure: Figure 1: Comparison of minADE across different scenario densities for HiVT and HPNet on Argoverse 1. Models are trained with the full dataset, 50% randomly selected data, 50% cluster-based selected data, and 50% SSTP (ours). SSTP consistently achieves lower errors in high-density scenarios.
Refer to caption: x1.png

We evaluate our proposed method on the Argoverse 1 (Chang et al., 2019) and Argoverse 2 (Wilson et al., 2023) datasets with multiple baseline methods.
Empirical results demonstrate that SSTP successfully constructs a compact and well-balanced dataset.
In high-density scenarios, training on the selected subset significantly outperforms using the full dataset, demonstrating the effectiveness of our method.
Moreover, our work offers a resource-efficient dataset that maintains balanced performance across various driving scenarios, making it well-suited for training state-of-the-art trajectory prediction models.
The main contributions are summarized as follows:

- •
We reveal that modern trajectory prediction datasets are heavily skewed, with scene density and other long-tailed factors biasing models toward majority regimes and limiting performance in safety-critical, rare scenarios.
- •
We introduce SSTP, a framework that partitions data by scene density and applies submodular selection with gradient-based influence scores, combined with biased sampling to emphasize underrepresented yet safety-critical cases.
- •
By selecting only 50% of the data, SSTP achieves performance comparable to or better than full-data training, substantially reduces computational cost, reduces error in high-density interaction scenes, and transfers robustly across diverse model architectures.

Figure: Figure 2: Visualization of different-density scenarios in trajectory prediction datasets. Left shows examples from Argoverse 1, and right shows examples from Argoverse 2.
Refer to caption: x2.png

## 2 Method

Denote the training set as $\mathcal{D}=\{S_{j}\}_{j=1}^{M}$. Each sample is represented as a triple $S=(X,Y,\mathcal{O})$, where $X$ and $Y$ denote the observed trajectories and future trajectories of all agents, and $\mathcal{O}$ the driving context information (e.g., maps).
The goal is to estimate $Y$ conditioned on $X$ and $\mathcal{O}$.
We aim to construct a compact subset $\mathcal{C}\subseteq\mathcal{D}$ of target size $B$.

Figure: Algorithm 1 Sample Selection for Trajectory Prediction

A model trained on $\mathcal{C}$ is expected to achieve trajectory prediction performance comparable to that obtained from training on the full dataset $\mathcal{D}$. Our method is formally described in Algorithm [1](https://arxiv.org/html/2409.17385v3#alg1). Next, we explain the key algorithmic components in our methods.

Data Partitioning.
As pointed out in Section [1](https://arxiv.org/html/2409.17385v3#S1), data imbalance in agent density is a common characteristic of datasets in the autonomous navigation domain. The impacts of such an imbalance are particularly pronounced when Transformer architectures or GNNs that rely on self-attention or graphs to capture dependencies among agents (Zhou et al., 2022c; 2023; Tang et al., 2024; Zhang et al., 2025).
Without density-based partitioning, gradient updates are dominated by abundant low-density scenarios, leading to overfitting on sparse interactions and systematic under-training of high-density scenarios. This imbalance prevents the model from learning robust attention patterns.
According to this observation, we first compute a density level $\rho(S_{j})$ for each sample $S_{j}\in\mathcal{D}$ based on the number of agents present. Then, we use a fixed interval $\tau$ to partition the dataset into $K$ disjoint subsets $\mathcal{D}_{k}$ for $k\in[K]$ based on these density values such that $S\in\mathcal{D}_{k}$ if $\rho(S)\in[\rho_{\min}+(k-1)\tau,\;\rho_{\min}+k\tau)$, where $\rho_{min}$ is the minimum density level in $\mathcal{D}$.

Gradient Extraction.
We calculate the total loss as follows:

$$ $\mathcal{L}=\mathcal{L}_{reg}+\mathcal{L}_{cls},$ (1) $$

where $\mathcal{L}_{reg}$ measures the L2 norm between the best mode prediction and the ground truth trajectory, and $\mathcal{L}_{cls}$ aligns the mode predicted probabilities $\pi$ with the best mode.
By backpropagating loss $\mathcal{L}$, we obtain the gradient with respect to $\hat{\mathbf{Y}}$ as follows:

$$ $\nabla_{\hat{\mathbf{Y}}}\,\mathcal{L}=\nabla_{\hat{\mathbf{Y}}}\,(\mathcal{L}_{reg}+\mathcal{L}_{cls}).$ (2) $$

Finally, we perform an element-wise multiplication of the calculated gradient with the corresponding decoder latent vector $\mathbf{E}$ as follows:

$$ $\mathbf{g}=\nabla_{\hat{\mathbf{Y}}}\,\mathcal{L}\odot\mathbf{E},$ (3) $$

where $\mathbf{g}$ captures the joint variations in the gradient and embedding spaces. In this manner, for each subset with $|\mathcal{D}_{k}|>n_{k}$, we construct a set of gradient feature vectors $\mathcal{G}_{k}=\{\mathbf{g}_{j}\}_{S_{j}\in\mathcal{D}_{k}}$ for all trajectory samples in that group.

Sample Selection.
Recall that $B$ is the target size of the constructed training subset.
To ensure fair representation across scene complexities, we adopt a dynamic allocation strategy.
For each partition $\mathcal{D}_{k}$, a local budget $n_{k}$ is assigned by the DynamicSelect function (see Algorithm [1](https://arxiv.org/html/2409.17385v3#alg1)).
This mechanism prioritizes high-density subsets that naturally contain fewer samples, preventing them from being underrepresented when the global budget is small. Formally, the allocation satisfies:

$$ $\sum_{k=1}^{K}n_{k}=B,\quad n_{k}\geq 0.$ (4) $$

Given total $K$ disjoint subsets and the corresponding gradient feature vectors $\mathcal{G}_{k}$, we present a submodular gain function $P(\cdot)$ to evaluate the contribution of each sample:

$$ $P(S_{j})=\sum_{S_{i}\in\mathcal{C}_{k}}\frac{\mathbf{g}_{i}\cdot\mathbf{g}_{j}}{\|\mathbf{g}_{i}\|\|\mathbf{g}_{j}\|}-\sum_{S_{i}\in\mathcal{D}_{k}\setminus\mathcal{C}_{k}}\frac{\mathbf{g}_{i}\cdot\mathbf{g}_{j}}{\|\mathbf{g}_{i}\|\|\mathbf{g}_{j}\|},$ (5) $$

where we use a cosine similarity kernel to measure the similarity between the sample $S_{j}$ and other samples.
We then apply a greedy optimization strategy, iteratively selecting the sample as follows:

$$ $S^{*}=\arg\min_{S_{j}\in\mathcal{D}_{k}\setminus\mathcal{C}_{k}}P(S_{j}).$ (6) $$

At each iteration, the selected sample $S^{*}$ is added to $\mathcal{C}_{k}$. This process continues
until the number of selected samples reaches the budget $n_{k}$ for subset $\mathcal{D}_{k}$.
Greedy optimization guarantees that the selected set $\mathcal{C}$ satisfies:

$$ $P(\mathcal{C})\;\;\geq\;\;\left(1-\tfrac{1}{e}\right)P(\mathcal{C}^{*}),$ (7) $$

where $\mathcal{C}^{*}$ denotes the optimal subset under target size $B$.
Notably, we process the subsets starting with those having higher density levels, as these subsets tend to be scarce and underrepresented. This prioritization is consistent with our dynamic allocation strategy, which ensures that complex, high-density scenarios are not discarded in early pruning and that long-tailed scenarios remain adequately represented. For subsets $\mathcal{D}_{k}$ where $|\mathcal{D}_{k}|\leq n_{k}$, we directly set $\mathcal{C}_{k}=\mathcal{D}_{k}$.
Finally, we can yield the target dataset $\mathcal{C}=\bigcup_{k=1}^{K}\mathcal{C}_{k}$.
By incorporating gradient-based similarity into submodular greedy selection, our method maximizes coverage of the gradient space while maintaining diversity, producing a smaller dataset that remains representative and informative.
In addition, we account for efficiency: the additional cost of sample selection is dominated by gradient computation and submodular updates. The overall computational complexity can be approximated as
$\mathcal{O}(\text{selection})=\mathcal{O}(|\mathcal{D}|\cdot d)+\mathcal{O}(B\cdot d)$,
where $d$ is the dimension of the gradient feature vectors.

Why Naive Strategies Fall Short.
To motivate the necessity of density-aware, gradient-guided selection, we examine several alternative strategies:
Re-weighting: assigns larger loss weights to high-density samples while retaining the full dataset;
Augmenting: duplicates high-density samples to artificially increase their proportion;
High-density+Random preserves all high-density samples first and fills the remainder via random sampling;
Epoch-wise: dynamically re-selects half the data at the start of each training epoch.
These alternatives highlight critical limitations. Both random down-sampling and epoch-wise re-selection significantly reduce training data, but the former discards informative scenarios (Table [1](https://arxiv.org/html/2409.17385v3#S2.T1), line 4) while the latter introduces instability despite added diversity (line 3).
Simple re-weighting or duplication naively equalizes sample presence but often introduces redundancy, showing that effective density balancing requires more than adjusting counts (line 1–2).

**Table 1: Comparison with other selection strategies on Argoverse 1. All methods are evaluated at 50% budget (95k samples) unless otherwise specified.**
| Method | #Samples | mADE$\downarrow$ | mFDE$\downarrow$ | MR$\downarrow$ |
| --- | --- | --- | --- | --- |
| Augmenting | 220k | 0.718 | 1.106 | 0.115 |
| Weighting | 190k | 0.715 | 1.108 | 0.114 |
| Epoch-wise | 95k | 0.752 | 1.189 | 0.130 |
| Random | 95k | 0.741 | 1.164 | 0.125 |
| High-density+Random | 95k | 0.724 | 1.111 | 0.117 |
| SSTP | 95k | 0.704 | 1.073 | 0.111 |

Preservation of high-density scenarios is intuitively beneficial, yet without a principled mechanism to regulate redundancy and ensure representativeness, gains remain limited (line 5). Importantly, in trajectory prediction, the scarcity of high-density scenarios is critical, as reliable performance requires learning complex multi-agent interactions that cannot be captured by oversampling or naive mixing. A central observation is that competence in complex, high-density interactions transfers naturally to simpler scenarios as shown in Figure [2](https://arxiv.org/html/2409.17385v3#S1.F2), whereas the reverse does not hold. This asymmetry motivates SSTP, which partitions by density and uses gradient-based selection to preserve scarce high-density scenarios while avoiding redundant low-density ones.

## 3 Experiments

### 3.1 Benchmarks and Setup

Datasets.
We evaluated the effectiveness of our proposed SSTP method on Argoverse Motion Forecasting Dataset 1.1  (Chang et al., 2019) and Argoverse 2 (Wilson et al., 2023). The Argoverse 1 dataset contains 323,557 real-world driving scenarios. All the training and validation scenarios are 5-second sequences sampled at 10 Hz. The length of the historical trajectory for each scenario is 2 seconds, and the length of the predicted future trajectory is 3 seconds. The Argoverse 2 dataset contains 250,000 scenarios, with the same sampling frequency of 10 Hz. Each trajectory has a larger observation window with 5 seconds and a longer prediction horizon with 6 seconds.

Baselines.
For Argoverse 1, we validate our SSTP method on two SOTA models HiVT (Zhou et al., 2022c) and HPNet (Tang et al., 2024) for evaluation.
For Argoverse 2, we evaluate our SSTP method using two SOTA models QCNet (Zhou et al., 2023) and DeMo (Zhang et al., 2025). For a more comprehensive comparison, we also include three following data selection approaches:

(1) Random Selection (Rebuffi et al., 2017): randomly selects a certain proportion of training samples from the original training set.

(2) K-Means Clustering (Likas et al., 2003): clusters trajectories within the observation window based on their features, and then selects the trajectory sample closest to the cluster center in each cluster as a representative.

(3) Herding Selection (Castro et al., 2018): a greedy strategy that first computes the mean feature of all trajectories within the observation window and then iteratively selects trajectory samples that bring the mean of the selected subset as close as possible to the overall mean.

**Table 2: Performance comparison results on Argoverse 1 with data retention ratios of 60%, 50%, and 40%. The compared methods include Random Selection, K-Means Clustering, and Herding Selection. The model used for data selection is HiVT-64, while the evaluation is conducted on HiVT-64, HiVT-128, and HPNet. $\alpha$ (%) represents the proportion of retained data relative to the full training set.**
| Methods | $\alpha$(%) | HiVT-64 | HiVT-128 | HPNet |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| mADE$\downarrow$ | mFDE$\downarrow$ | xMR$\downarrow$x | mADE$\downarrow$ | mFDE$\downarrow$ | xMR$\downarrow$x | mADE$\downarrow$ | mFDE$\downarrow$ | xMR$\downarrow$x |  |  |
| Argoverse 1 | 100 | 0.695 | 1.037 | 0.109 | 0.666 | 0.978 | 0.091 | 0.647 | 0.871 | 0.070 |
| Random | 60 | 0.745 | 1.163 | 0.132 | 0.719 | 1.078 | 0.129 | 0.680 | 0.951 | 0.091 |
| Cluster | 0.716 | 1.097 | 0.121 | 0.697 | 1.025 | 0.108 | 0.673 | 0.930 | 0.081 |  |
| Herding | 0.723 | 1.101 | 0.125 | 0.685 | 1.018 | 0.106 | 0.666 | 0.922 | 0.085 |  |
| SSTP | 0.702 | 1.064 | 0.110 | 0.674 | 0.994 | 0.093 | 0.653 | 0.901 | 0.071 |  |
| Random | 50 | 0.750 | 1.175 | 0.137 | 0.728 | 1.098 | 0.126 | 0.687 | 0.967 | 0.091 |
| Cluster | 0.725 | 1.117 | 0.124 | 0.692 | 1.033 | 0.118 | 0.676 | 0.952 | 0.085 |  |
| Herding | 0.728 | 1.107 | 0.126 | 0.698 | 1.036 | 0.119 | 0.674 | 0.938 | 0.089 |  |
| SSTP | 0.706 | 1.074 | 0.110 | 0.684 | 1.022 | 0.101 | 0.661 | 0.913 | 0.074 |  |
| Random | 40 | 0.752 | 1.183 | 0.139 | 0.727 | 1.109 | 0.126 | 0.696 | 0.987 | 0.099 |
| Cluster | 0.732 | 1.141 | 0.127 | 0.703 | 1.058 | 0.121 | 0.681 | 0.962 | 0.089 |  |
| Herding | 0.722 | 1.123 | 0.128 | 0.704 | 1.056 | 0.119 | 0.684 | 0.956 | 0.093 |  |
| SSTP | 0.711 | 1.088 | 0.114 | 0.696 | 1.048 | 0.106 | 0.671 | 0.931 | 0.076 |  |

**Table 3: Comparison of model performance across scenario densities when trained on the full dataset, a 50% random subset, and a 50% SSTP subset, where SSTP consistently achieves superior results.**
|  | Agent$<$40 | Agent$>=$40 | Agent$>=$60 | Agent$>=$80 |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | mADE$\downarrow$ | mFDE$\downarrow$ | MR$\downarrow$ | mADE$\downarrow$ | mFDE$\downarrow$ | MR$\downarrow$ | mADE$\downarrow$ | mFDE$\downarrow$ | MR$\downarrow$ | mADE$\downarrow$ | mFDE$\downarrow$ | MR$\downarrow$ |
| Full | 0.700 | 1.071 | 0.108 | 0.950 | 1.456 | 0.171 | 1.248 | 1.898 | 0.283 | 1.450 | 2.059 | 0.361 |
| Random | 0.734 | 1.127 | 0.119 | 0.997 | 1.552 | 0.193 | 1.287 | 1.989 | 0.315 | 1.638 | 2.359 | 0.389 |
| Cluster | 0.719 | 1.121 | 0.115 | 0.971 | 1.501 | 0.188 | 1.251 | 1.898 | 0.304 | 1.568 | 2.257 | 0.376 |
| Herding | 0.726 | 1.122 | 0.116 | 0.980 | 1.528 | 0.190 | 1.258 | 1.922 | 0.310 | 1.581 | 2.381 | 0.372 |
| SSTP | 0.714 | 1.098 | 0.113 | 0.962 | 1.497 | 0.183 | 1.219 | 1.835 | 0.280 | 1.373 | 1.762 | 0.277 |

Metrics.
Following the baselines, we also generate a total $6$ future trajectories and use the metrics minimum Average Displacement Error (minADE), minimum Final Displacement Error (minFDE), and Missing Rate (MR) to evaluate the prediction performance.

Implementation Details.
We primarily utilized the pre-trained HiVT-64 and QCNet as backbone models to perform sample selection on the Argoverse 1 and Argoverse 2 datasets, respectively. To evaluate the performance of the selected subset, we follow their official training and validation protocols. We experimented with different selection ratios, different intervals, and assessed the prediction accuracy of the trajectory models after training on the corresponding subsets. Further implementation details are provided in Appendix [A.1](https://arxiv.org/html/2409.17385v3#A1.SS1).

### 3.2 Main Results

We present the primary experimental results in this section, while a comprehensive ablation study, including analyses of alternative strategies and design choices, is provided in Appendix [A.3](https://arxiv.org/html/2409.17385v3#A1.SS3). In addition, qualitative results and visual analyses are presented in Appendix [A.4](https://arxiv.org/html/2409.17385v3#A1.SS4).

Table [2](https://arxiv.org/html/2409.17385v3#S3.T2) showcases the strong performance of our selected subset on the Argoverse 1 dataset across all compression rates. Following the same experimental setup as the baseline models, we trained HiVT and HPNet from scratch on the subset. Our method significantly reduces data volume while maintaining nearly lossless performance. Even with only half of the data, models trained on our selected subset still perform comparably to those trained on the full dataset and consistently outperform random selection, clustering, and herding. Furthermore, our subset also demonstrates strong results on larger models such as HiVT-128 and HPNet. Additional results for other data retention ratios are provided in Appendix [A.2.1](https://arxiv.org/html/2409.17385v3#A1.SS2.SSS1), with detailed numbers reported in Table [6](https://arxiv.org/html/2409.17385v3#A1.T6).

**Table 4: Performance comparison results on Argoverse 2 with data retention ratios of 60%, 50%, and 40%. The compared methods include Random Selection and K-means clustering. The model used for data selection is QCNet, while the evaluation is conducted on QCNet and DeMo.**
| Method | $\alpha$(%) | QCNet | DeMo |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| mADE$\downarrow$ | mFDE$\downarrow$ | xMR$\downarrow$x | mADE$\downarrow$ | mFDE$\downarrow$ | xMR$\downarrow$x |  |  |
| Argoverse 2 | 100 | 0.724 | 1.258 | 0.162 | 0.657 | 1.254 | 0.163 |
| Random | 60 | 0.787 | 1.419 | 0.208 | 0.755 | 1.433 | 0.198 |
| Cluster | 0.773 | 1.406 | 0.192 | 0.693 | 1.386 | 0.187 |  |
| Herding | 0.778 | 1.402 | 0.191 | 0.701 | 1.494 | 0.189 |  |
| SSTP | 0.740 | 1.316 | 0.163 | 0.682 | 1.344 | 0.164 |  |
| Random | 50 | 0.805 | 1.447 | 0.219 | 0.756 | 1.448 | 0.203 |
| Cluster | 0.798 | 1.435 | 0.193 | 0.732 | 1.437 | 0.191 |  |
| Herding | 0.782 | 1.407 | 0.190 | 0.731 | 1.434 | 0.190 |  |
| SSTP | 0.754 | 1.352 | 0.172 | 0.704 | 1.414 | 0.173 |  |
| Random | 40 | 0.811 | 1.471 | 0.226 | 0.763 | 1.475 | 0.202 |
| Cluster | 0.813 | 1.495 | 0.214 | 0.732 | 1.456 | 0.195 |  |
| Herding | 0.807 | 1.454 | 0.198 | 0.739 | 1.460 | 0.197 |  |
| SSTP | 0.778 | 1.410 | 0.183 | 0.723 | 1.450 | 0.191 |  |

Performance Enhancement.
Scene density in autonomous driving varies substantially, yet most existing trajectory prediction datasets are dominated by low-density scenarios. From a safety perspective, however, an ideal trajectory predictor must perform reliably across the full spectrum of scene complexities. To this end, we evaluate our proposed method on multiple models and across different density levels.
As shown in Table [3](https://arxiv.org/html/2409.17385v3#S3.T3), our method consistently outperforms models trained on the full dataset, particularly in high-density scenarios, on both the Argoverse 1 and Argoverse 2 datasets. In low-density settings (fewer than 40 agents), our selected subset achieves performance nearly identical to the full dataset, with only marginal increases of minADE and minFDE, while MR remains almost unchanged, which is negligible given the limited interactions in such scenes.

When the agent density increases, our method brings the most substantial gains. Compared to all other baselines, SSTP achieves lower displacement errors and notably reduces MR. For example, when the number of agents exceeds 80, SSTP cuts the missing rate by more than 8% relative to random selection, and also outperforms clustering and herding by clear margins.
These results highlight that while clustering and herding provide partial improvements by ensuring representativeness, they are still insufficient for handling highly complex traffic scenes. In contrast, SSTP effectively balances scene density while selecting informative samples, leading to consistently superior performance across all density levels. Similarly, the SOTA model HPNet, when trained on our selected subset, even outperforms its counterpart trained on the full dataset. Detailed results are provided in Appendix [A.2.2](https://arxiv.org/html/2409.17385v3#A1.SS2.SSS2), with comprehensive numbers reported in Table [8](https://arxiv.org/html/2409.17385v3#A1.T8).

Figure: Figure 3: Performance of HiVT and HPNet trained on subsets selected by different methods at varying retention ratios, evaluated with minADE.
Refer to caption: x3.png

Generalization Across Datasets.
We further evaluated our proposed method on the Argoverse 2 dataset, which presents greater challenges due to its more diverse driving scenarios and longer prediction horizons. As shown in Table [4](https://arxiv.org/html/2409.17385v3#S3.T4), our method consistently outperforms other data selection strategies across all data retention rates, achieving lower minADE and minFDE while maintaining a lower MR.
These results further validate the robustness of our approach, as it maintains strong performance across different datasets. This demonstrates that our method is not only effective within a specific dataset but also generalizes well to more complex and diverse trajectory scenarios, such as those found in Argoverse 2. Additional results for other data retention ratios are provided in Appendix [A.2.1](https://arxiv.org/html/2409.17385v3#A1.SS2.SSS1), with detailed numbers reported in Table [7](https://arxiv.org/html/2409.17385v3#A1.T7).

**Table 5: Performance comparison of data selection strategies on HiVT trained with Argoverse 1. This table shows the impact of partitioning and selection with Submodular Gain strategies on data distribution and model performance. The whole dataset and random selection serve as baselines, while different variations of SSTP are evaluated. Our method (SSTP), integrates both strategies, achieves the best results by maintaining a balanced data distribution and reducing minADE, minFDE, and MR.**
| Variants | Data Distribution(%) | Model Performance |  |  |  |
| --- | --- | --- | --- | --- | --- |
| Agent$<$40 | Agent$>=$40 | mADE$\downarrow$ | mFDE$\downarrow$ | xMR$\downarrow$x |  |
| Full dataset | 93.88 | 6.12 | 0.692 | 1.047 | 0.104 |
| Random | 85.16 | 14.84 | 0.741 | 1.164 | 0.125 |
| SSTP w/ Submodular | 93.88 | 6.12 | 0.724 | 1.115 | 0.116 |
| SSTP w/ Partition | 70.35 | 29.65 | 0.729 | 1.116 | 0.118 |
| SSTP | 70.35 | 29.65 | 0.704 | 1.073 | 0.111 |

Data Retention Ratio.
To examine the impact of different data retention ratios on model performance, we conducted experiments with retention rates of $\alpha$ = 60, 50, 40, 30, 20, 10% as shown in Figure [4](https://arxiv.org/html/2409.17385v3#S3.F4). When higher model performance is required, retaining 50% of the data already achieves results comparable to training with the full dataset. This demonstrates the effectiveness of our SSTP method, as the selected subset is of higher quality compared to equally sized subsets chosen by other methods. Furthermore, under limited computational resources, even retaining only 20% of the data still yields reasonably good results.
We further analyze the behavior under extremely low retention ratios, with detailed results provided in Appendix [A.2.3](https://arxiv.org/html/2409.17385v3#A1.SS2.SSS3) and Table [9](https://arxiv.org/html/2409.17385v3#A1.T9).

Density Balancing. Our method explicitly controls scene density distribution during selection, ensuring a more balanced dataset, as illustrated in Figure [4](https://arxiv.org/html/2409.17385v3#S3.F4). In contrast, random selection fails to maintain this balance, leading to uneven representation of scenarios with varying complexity and weaker generalization. As shown in Table [5](https://arxiv.org/html/2409.17385v3#S3.T5) line 4, applying scene balancing alone already improves performance compared to random selection, demonstrating that controlling scene density enhances the effectiveness of trajectory prediction models. However, balancing alone remains insufficient. By further integrating submodular selection to account for sample informativeness, our method achieves the best performance across all metrics. These findings indicate that while scene balancing is beneficial, it is insufficient to achieve optimal performance without also considering sample informativeness.

Figure: Figure 5: Training time comparison of HiVT and HPNet. Our method achieves over 50% reduction in time while maintaining comparable accuracy.
Refer to caption: x5.png

Efficiency. Our method significantly reduces computational time while maintaining strong performance, as shown in Figure [5](https://arxiv.org/html/2409.17385v3#S3.F5). Training the HiVT model on the full dataset requires 7.78 hours. In contrast, utilizing our SSTP method to select a 60% subset requires only 0.6 hour for pre-training and 2.30 hours for selection, reducing the overall training time. When training on the selected subset, the total training time decreases to 7.25 hours (0.6 + 2.30 + 4.35), with only a minor increase of 0.007 in minADE.

Beyond efficiency on the original backbone, a key advantage of SSTP is that the selected subset is model-agnostic and can be directly reused to train other trajectory predictors. For instance, when applied to HPNet, training on the full dataset takes 307.2 hours, whereas training with the selected subset reduces the time by over 100 hours, cutting the training cost by nearly 45%, while achieving comparable or even better accuracy.
This decisively confirms the superior efficiency of our approach in balancing training cost and model performance.

Generalizability across Backbones. To further assess the generalizability of our method, we evaluated its performance with different backbone models. Specifically, we used pre-trained HiVT-64 and HPNet as feature extractors on the Argoverse 1 dataset and conducted experiments under varying data retention ratios. Detailed results are provided in Appendix [A.2.4](https://arxiv.org/html/2409.17385v3#A1.SS2.SSS4), with numerical values reported in Table [10](https://arxiv.org/html/2409.17385v3#A1.T10). As shown, regardless of the backbone used for subset selection, the final trajectory prediction performance remains nearly identical. This demonstrates that our subset selection strategy is largely backbone-agnostic, underscoring its robustness and broad applicability for optimizing diverse trajectory prediction models.

## 4 Related Work

Trajectory Prediction.
Trajectory prediction infers an agent’s future motion from its historical observations. In recent years, research has increasingly concentrated on capturing complex multi-agent interactions, driving advances in various predictive methods (Wang et al., 2024; Giuliari et al., 2021; Xu et al., 2020; Liu et al., 2021; Ngiam et al., 2021; Zhou et al., 2022c; 2023; Tang et al., 2024; Zhang et al., 2025; Liang et al., 2020). Furthermore, novel approaches including pretraining (Chen et al., 2023; Cheng et al., 2023; Lan et al., 2023), historical prediction structure design (Park et al., 2024; Tang et al., 2024), GPT-style next-token prediction (Philion et al., 2023; Seff et al., 2023), and post-processing optimization (Zhou et al., 2024; Choi et al., 2023) have significantly enhanced model performance, demonstrating strong results across various datasets. However, most state-of-the-art methods rely on training with large-scale datasets (Caesar et al., 2020; Chang et al., 2019; Sun et al., 2020; Wilson et al., 2023; Ettinger et al., 2021), leading to significant computational costs. In contrast, we propose a pioneering data selection strategy in the trajectory prediction domain that constructs a compact, balanced, yet highly representative dataset, significantly reducing training time while preserving model performance.

Long-Tail in Trajectory Prediction Dataset
The performance of the trajectory prediction models is evaluated on the overall average.
While they excel on benchmarks, these models often struggle with challenging scenarios (Makansi et al., 2021; Wang et al., 2023; Pourkeshavarz et al., 2023) due to the long-tail data distribution, where common cases dominate and complex or rare situations are underrepresented (Chen et al., 2024). Recent studies (Wang et al., 2023; Zhang et al., 2024; Zhou et al., 2022a; Lan et al., 2024) have started addressing the long-tail problem in trajectory prediction, primarily by leveraging contrastive learning to enhance feature representations.
However, these methods mainly focus on optimizing feature-level learning while overlooking scenario level distribution and the importance of individual samples. In contrast, our proposed method assesses the contribution of each sample and applies a refined selection strategy to build a balanced and compact dataset. Experimental results show that it significantly boosts performance in complex scenarios.

Training Sample Selection
Deep neural networks, especially Transformers, depend on large datasets and incur high computational costs. To reduce these costs and shorten training time, various methods to improve data efficiency have been proposed, including frequent parameter updates (Robbins & Monro, 1951), fewer iterations (Sutskever et al., 2013), and dynamic learning rate adjustments (Kingma, 2014; Duchi et al., 2011). To directly reduce data volume, dataset condensation compresses raw data into compact synthetic samples (Wang et al., 2018; Zhao et al., 2020; Zhao & Bilen, 2021; Kim et al., 2022; Wang et al., 2022; Zhao & Bilen, 2023; Cazenavette et al., 2022). Another widely studied approach is coreset selection, which constructs a weighted subset that closely approximates the statistical distribution of the original dataset (Har-Peled & Mazumdar, 2004; Coleman et al., 2019; Margatina et al., 2021; Mirzasoleiman et al., 2020). However, those approaches on dataset condensation (Nguyen et al., 2020; Loo et al., 2022; Zhou et al., 2022b; Sajedi et al., 2023; Zhao et al., 2023) and coreset selection (Killamsetty et al., 2021; Paul et al., 2021) has been predominantly focused on image classification. In contrast, we propose a sample selection strategy based on submodular functions for the trajectory prediction domain. As a pioneering method, our approach significantly reduces the training data required while maintaining model performance.

## 5 Conclusion

In this paper, we presented the Sample Selection for Trajectory Prediction (SSTP) framework, a novel, data-centric approach that constructs a compact yet balanced dataset for trajectory prediction.
SSTP effectively tackles the challenges posed by data imbalance and the high training costs inherent in large-scale trajectory datasets.
By significantly reducing the training data volume while maintaining, and even enhancing the model performance in high-density scenarios.
SSTP not only accelerates training but also delivers results comparable to, or better than, those achieved with full-dataset training.
Extensive evaluations on the Argoverse 1 and Argoverse 2 benchmarks across a wide range of state-of-the-art models underscore the practical value of our approach in improving both efficiency and robustness in trajectory prediction for autonomous driving.

Limitations.
While effective, our method still incurs non-negligible computational overhead during the gradient extraction and submodular optimization stages of sample selection. Although this is a one-time cost, further optimization to reduce its complexity would be desirable. Moreover, model performance degrades under extremely low data retention ratios (e.g., 10%), highlighting the challenge of preserving robustness when only very limited data are available. Finally, the present study is restricted to trajectory prediction tasks on the Argoverse benchmarks; extending SSTP to other autonomous driving tasks and additional datasets remains an important avenue for future research.

## Ethics Statement

This work uses only publicly available autonomous driving datasets (Argoverse 1 and 2) containing vehicle trajectories and maps, with no human subjects or private data involved. We adhere to the conference Code of Ethics and confirm full compliance with its standards.

## Reproducibility Statement

Implementation details, including dataset preprocessing, training protocols, and hyperparameter settings, are provided in Section 4 and Appendix A.1. We also release our code and checkpoints (via the anonymous link) to facilitate full reproducibility.

## References

- Caesar et al. (2020)
Holger Caesar, Varun Bankiti, Alex H Lang, Sourabh Vora, Venice Erin Liong, Qiang Xu, Anush Krishnan, Yu Pan, Giancarlo Baldan, and Oscar Beijbom.
nuscenes: A multimodal dataset for autonomous driving.
In *Proceedings of the IEEE/CVF conference on computer vision and pattern recognition*, pp. 11621–11631, 2020.
- Castro et al. (2018)
Francisco M Castro, Manuel J Marín-Jiménez, Nicolás Guil, Cordelia Schmid, and Karteek Alahari.
End-to-end incremental learning.
In *Proceedings of the European conference on computer vision (ECCV)*, pp. 233–248, 2018.
- Cazenavette et al. (2022)
George Cazenavette, Tongzhou Wang, Antonio Torralba, Alexei A Efros, and Jun-Yan Zhu.
Dataset distillation by matching training trajectories.
In *Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition*, pp. 4750–4759, 2022.
- Chai et al. (2019)
Yuning Chai, Benjamin Sapp, Mayank Bansal, and Dragomir Anguelov.
Multipath: Multiple probabilistic anchor trajectory hypotheses for behavior prediction.
*arXiv preprint arXiv:1910.05449*, 2019.
- Chang et al. (2019)
Ming-Fang Chang, John Lambert, Patsorn Sangkloy, Jagjeet Singh, Slawomir Bak, Andrew Hartnett, De Wang, Peter Carr, Simon Lucey, Deva Ramanan, et al.
Argoverse: 3d tracking and forecasting with rich maps.
In *Proceedings of the IEEE/CVF conference on computer vision and pattern recognition*, pp. 8748–8757, 2019.
- Chen et al. (2024)
Changhe Chen, Mozhgan Pourkeshavarz, and Amir Rasouli.
Criteria: a new benchmarking paradigm for evaluating trajectory prediction models for autonomous driving.
In *2024 IEEE International Conference on Robotics and Automation (ICRA)*, pp. 8265–8271. IEEE, 2024.
- Chen et al. (2023)
Hao Chen, Jiaze Wang, Kun Shao, Furui Liu, Jianye Hao, Chenyong Guan, Guangyong Chen, and Pheng-Ann Heng.
Traj-mae: Masked autoencoders for trajectory prediction.
In *Proceedings of the IEEE/CVF International Conference on Computer Vision*, pp. 8351–8362, 2023.
- Cheng et al. (2023)
Jie Cheng, Xiaodong Mei, and Ming Liu.
Forecast-mae: Self-supervised pre-training for motion forecasting with masked autoencoders.
In *Proceedings of the IEEE/CVF International Conference on Computer Vision*, pp. 8679–8689, 2023.
- Choi et al. (2023)
Sehwan Choi, Jungho Kim, Junyong Yun, and Jun Won Choi.
R-pred: Two-stage motion prediction via tube-query attention-based trajectory refinement.
In *Proceedings of the IEEE/CVF International Conference on Computer Vision*, pp. 8525–8535, 2023.
- Coleman et al. (2019)
Cody Coleman, Christopher Yeh, Stephen Mussmann, Baharan Mirzasoleiman, Peter Bailis, Percy Liang, Jure Leskovec, and Matei Zaharia.
Selection via proxy: Efficient data selection for deep learning.
*arXiv preprint arXiv:1906.11829*, 2019.
- Duchi et al. (2011)
John Duchi, Elad Hazan, and Yoram Singer.
Adaptive subgradient methods for online learning and stochastic optimization.
*Journal of machine learning research*, 12(7), 2011.
- Ettinger et al. (2021)
Scott Ettinger, Shuyang Cheng, Benjamin Caine, Chenxi Liu, Hang Zhao, Sabeek Pradhan, Yuning Chai, Ben Sapp, Charles R Qi, Yin Zhou, et al.
Large scale interactive motion forecasting for autonomous driving: The waymo open motion dataset.
In *Proceedings of the IEEE/CVF International Conference on Computer Vision*, pp. 9710–9719, 2021.
- Feng et al. (2023)
Chen Feng, Hangning Zhou, Huadong Lin, Zhigang Zhang, Ziyao Xu, Chi Zhang, Boyu Zhou, and Shaojie Shen.
Macformer: Map-agent coupled transformer for real-time and robust trajectory prediction.
*IEEE Robotics and Automation Letters*, 2023.
- Giuliari et al. (2021)
Francesco Giuliari, Irtiza Hasan, Marco Cristani, and Fabio Galasso.
Transformer networks for trajectory forecasting.
In *2020 25th international conference on pattern recognition (ICPR)*, pp. 10335–10342. IEEE, 2021.
- Gu et al. (2021)
Junru Gu, Chen Sun, and Hang Zhao.
Densetnt: End-to-end trajectory prediction from dense goal sets.
In *Proceedings of the IEEE/CVF International Conference on Computer Vision*, pp. 15303–15312, 2021.
- Har-Peled & Mazumdar (2004)
Sariel Har-Peled and Soham Mazumdar.
On coresets for k-means and k-median clustering.
In *Proceedings of the thirty-sixth annual ACM symposium on Theory of computing*, pp. 291–300, 2004.
- Killamsetty et al. (2021)
Krishnateja Killamsetty, Sivasubramanian Durga, Ganesh Ramakrishnan, Abir De, and Rishabh Iyer.
Grad-match: Gradient matching based data subset selection for efficient deep model training.
In *International Conference on Machine Learning*, pp. 5464–5474. PMLR, 2021.
- Kim et al. (2022)
Jang-Hyun Kim, Jinuk Kim, Seong Joon Oh, Sangdoo Yun, Hwanjun Song, Joonhyun Jeong, Jung-Woo Ha, and Hyun Oh Song.
Dataset condensation via efficient synthetic-data parameterization.
In *International Conference on Machine Learning*, pp. 11102–11118. PMLR, 2022.
- Kingma (2014)
DP Kingma.
Adam: a method for stochastic optimization.
*arXiv preprint arXiv:1412.6980*, 2014.
- Lan et al. (2024)
Zhengxing Lan, Yilong Ren, Haiyang Yu, Lingshan Liu, Zhenning Li, Yinhai Wang, and Zhiyong Cui.
Hi-scl: Fighting long-tailed challenges in trajectory prediction with hierarchical wave-semantic contrastive learning.
*Transportation Research Part C: Emerging Technologies*, 165:104735, 2024.
- Lan et al. (2023)
Zhiqian Lan, Yuxuan Jiang, Yao Mu, Chen Chen, and Shengbo Eben Li.
Sept: Towards efficient scene representation learning for motion prediction.
*arXiv preprint arXiv:2309.15289*, 2023.
- Liang et al. (2020)
Junwei Liang, Lu Jiang, Kevin Murphy, Ting Yu, and Alexander Hauptmann.
The garden of forking paths: Towards multi-future trajectory prediction.
In *Proceedings of the IEEE/CVF conference on computer vision and pattern recognition*, pp. 10508–10518, 2020.
- Likas et al. (2003)
Aristidis Likas, Nikos Vlassis, and Jakob J Verbeek.
The global k-means clustering algorithm.
*Pattern recognition*, 36(2):451–461, 2003.
- Liu et al. (2021)
Yicheng Liu, Jinghuai Zhang, Liangji Fang, Qinhong Jiang, and Bolei Zhou.
Multimodal motion prediction with stacked transformers.
In *Proceedings of the IEEE/CVF conference on computer vision and pattern recognition*, pp. 7577–7586, 2021.
- Loo et al. (2022)
Noel Loo, Ramin Hasani, Alexander Amini, and Daniela Rus.
Efficient dataset distillation using random feature approximation.
*Advances in Neural Information Processing Systems*, 35:13877–13891, 2022.
- Makansi et al. (2021)
Osama Makansi, Özgün Çiçek, Yassine Marrakchi, and Thomas Brox.
On exposing the challenging long tail in future prediction of traffic actors.
In *Proceedings of the IEEE/CVF International Conference on Computer Vision*, pp. 13147–13157, 2021.
- Margatina et al. (2021)
Katerina Margatina, Giorgos Vernikos, Loïc Barrault, and Nikolaos Aletras.
Active learning by acquiring contrastive examples.
*arXiv preprint arXiv:2109.03764*, 2021.
- Mirzasoleiman et al. (2020)
Baharan Mirzasoleiman, Jeff Bilmes, and Jure Leskovec.
Coresets for data-efficient training of machine learning models.
In *International Conference on Machine Learning*, pp. 6950–6960. PMLR, 2020.
- Ngiam et al. (2021)
Jiquan Ngiam, Benjamin Caine, Vijay Vasudevan, Zhengdong Zhang, Hao-Tien Lewis Chiang, Jeffrey Ling, Rebecca Roelofs, Alex Bewley, Chenxi Liu, Ashish Venugopal, et al.
Scene transformer: A unified architecture for predicting multiple agent trajectories.
*arXiv preprint arXiv:2106.08417*, 2021.
- Nguyen et al. (2020)
Timothy Nguyen, Zhourong Chen, and Jaehoon Lee.
Dataset meta-learning from kernel ridge-regression.
*arXiv preprint arXiv:2011.00050*, 2020.
- Park et al. (2024)
Daehee Park, Jaeseok Jeong, Sung-Hoon Yoon, Jaewoo Jeong, and Kuk-Jin Yoon.
T4p: Test-time training of trajectory prediction via masked autoencoder and actor-specific token memory.
In *Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition*, pp. 15065–15076, 2024.
- Paul et al. (2021)
Mansheej Paul, Surya Ganguli, and Gintare Karolina Dziugaite.
Deep learning on a data diet: Finding important examples early in training.
*Advances in neural information processing systems*, 34:20596–20607, 2021.
- Philion et al. (2023)
Jonah Philion, Xue Bin Peng, and Sanja Fidler.
Trajeglish: Learning the language of driving scenarios.
*arXiv preprint arXiv:2312.04535*, 2(7):12, 2023.
- Pourkeshavarz et al. (2023)
Mozhgan Pourkeshavarz, Changhe Chen, and Amir Rasouli.
Learn tarot with mentor: A meta-learned self-supervised approach for trajectory prediction.
In *Proceedings of the IEEE/CVF International Conference on Computer Vision*, pp. 8384–8393, 2023.
- Rebuffi et al. (2017)
Sylvestre-Alvise Rebuffi, Alexander Kolesnikov, Georg Sperl, and Christoph H Lampert.
icarl: Incremental classifier and representation learning.
In *Proceedings of the IEEE conference on Computer Vision and Pattern Recognition*, pp. 2001–2010, 2017.
- Robbins & Monro (1951)
Herbert Robbins and Sutton Monro.
A stochastic approximation method.
*The annals of mathematical statistics*, pp. 400–407, 1951.
- Sajedi et al. (2023)
Ahmad Sajedi, Samir Khaki, Ehsan Amjadian, Lucy Z Liu, Yuri A Lawryshyn, and Konstantinos N Plataniotis.
Datadam: Efficient dataset distillation with attention matching.
In *Proceedings of the IEEE/CVF International Conference on Computer Vision*, pp. 17097–17107, 2023.
- Seff et al. (2023)
Ari Seff, Brian Cera, Dian Chen, Mason Ng, Aurick Zhou, Nigamaa Nayakanti, Khaled S Refaat, Rami Al-Rfou, and Benjamin Sapp.
Motionlm: Multi-agent motion forecasting as language modeling.
In *Proceedings of the IEEE/CVF International Conference on Computer Vision*, pp. 8579–8590, 2023.
- Sun et al. (2020)
Pei Sun, Henrik Kretzschmar, Xerxes Dotiwalla, Aurelien Chouard, Vijaysai Patnaik, Paul Tsui, James Guo, Yin Zhou, Yuning Chai, Benjamin Caine, et al.
Scalability in perception for autonomous driving: Waymo open dataset.
In *Proceedings of the IEEE/CVF conference on computer vision and pattern recognition*, pp. 2446–2454, 2020.
- Sutskever et al. (2013)
Ilya Sutskever, James Martens, George Dahl, and Geoffrey Hinton.
On the importance of initialization and momentum in deep learning.
In *International conference on machine learning*, pp. 1139–1147. PMLR, 2013.
- Tang et al. (2024)
Xiaolong Tang, Meina Kan, Shiguang Shan, Zhilong Ji, Jinfeng Bai, and Xilin Chen.
Hpnet: Dynamic trajectory forecasting with historical prediction attention.
In *Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition*, pp. 15261–15270, 2024.
- Wang et al. (2024)
Jiangwei Wang, Lili Su, Songyang Han, Dongjin Song, and Fei Miao.
Towards safe autonomy in hybrid traffic: Detecting unpredictable abnormal behaviors of human drivers via information sharing.
*ACM Transactions on Cyber-Physical Systems*, 8(2):1–25, 2024.
- Wang et al. (2022)
Kai Wang, Bo Zhao, Xiangyu Peng, Zheng Zhu, Shuo Yang, Shuo Wang, Guan Huang, Hakan Bilen, Xinchao Wang, and Yang You.
Cafe: Learning to condense dataset by aligning features.
In *Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition*, pp. 12196–12205, 2022.
- Wang et al. (2018)
Tongzhou Wang, Jun-Yan Zhu, Antonio Torralba, and Alexei A Efros.
Dataset distillation.
*arXiv preprint arXiv:1811.10959*, 2018.
- Wang et al. (2023)
Yuning Wang, Pu Zhang, Lei Bai, and Jianru Xue.
Fend: A future enhanced distribution-aware contrastive learning framework for long-tail trajectory prediction.
In *Proceedings of the IEEE/CVF conference on computer vision and pattern recognition*, pp. 1400–1409, 2023.
- Wilson et al. (2023)
Benjamin Wilson, William Qi, Tanmay Agarwal, John Lambert, Jagjeet Singh, Siddhesh Khandelwal, Bowen Pan, Ratnesh Kumar, Andrew Hartnett, Jhony Kaesemodel Pontes, et al.
Argoverse 2: Next generation datasets for self-driving perception and forecasting.
*arXiv preprint arXiv:2301.00493*, 2023.
- Xu et al. (2020)
Mingxing Xu, Wenrui Dai, Chunmiao Liu, Xing Gao, Weiyao Lin, Guo-Jun Qi, and Hongkai Xiong.
Spatial-temporal transformer networks for traffic flow forecasting.
*arXiv preprint arXiv:2001.02908*, 2020.
- Zhan et al. (2019)
Wei Zhan, Liting Sun, Di Wang, Haojie Shi, Aubrey Clausse, Maximilian Naumann, Julius Kummerle, Hendrik Konigshof, Christoph Stiller, Arnaud de La Fortelle, et al.
Interaction dataset: An international, adversarial and cooperative motion dataset in interactive driving scenarios with semantic maps.
*arXiv preprint arXiv:1910.03088*, 2019.
- Zhang et al. (2025)
Bozhou Zhang, Nan Song, and Li Zhang.
Decoupling motion forecasting into directional intentions and dynamic states.
*Advances in Neural Information Processing Systems*, 37:106582–106606, 2025.
- Zhang et al. (2024)
Junrui Zhang, Mozhgan Pourkeshavarz, and Amir Rasouli.
Tract: A training dynamics aware contrastive learning framework for long-tail trajectory prediction.
In *2024 IEEE Intelligent Vehicles Symposium (IV)*, pp. 3282–3288. IEEE, 2024.
- Zhang et al. (2021)
Lingyao Zhang, Po-Hsun Su, Jerrick Hoang, Galen Clark Haynes, and Micol Marchetti-Bowick.
Map-adaptive goal-based trajectory prediction.
In *Conference on Robot Learning*, pp. 1371–1383. PMLR, 2021.
- Zhao & Bilen (2021)
Bo Zhao and Hakan Bilen.
Dataset condensation with differentiable siamese augmentation.
In *International Conference on Machine Learning*, pp. 12674–12685. PMLR, 2021.
- Zhao & Bilen (2023)
Bo Zhao and Hakan Bilen.
Dataset condensation with distribution matching.
In *Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision*, pp. 6514–6523, 2023.
- Zhao et al. (2020)
Bo Zhao, Konda Reddy Mopuri, and Hakan Bilen.
Dataset condensation with gradient matching.
*arXiv preprint arXiv:2006.05929*, 2020.
- Zhao et al. (2023)
Ganlong Zhao, Guanbin Li, Yipeng Qin, and Yizhou Yu.
Improved distribution matching for dataset condensation.
In *Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition*, pp. 7856–7865, 2023.
- Zhao et al. (2021)
Hang Zhao, Jiyang Gao, Tian Lan, Chen Sun, Ben Sapp, Balakrishnan Varadarajan, Yue Shen, Yi Shen, Yuning Chai, Cordelia Schmid, et al.
Tnt: Target-driven trajectory prediction.
In *Conference on Robot Learning*, pp. 895–904. PMLR, 2021.
- Zhou et al. (2022a)
Weitao Zhou, Zhong Cao, Yunkang Xu, Nanshan Deng, Xiaoyu Liu, Kun Jiang, and Diange Yang.
Long-tail prediction uncertainty aware trajectory planning for self-driving vehicles.
In *2022 IEEE 25th International Conference on Intelligent Transportation Systems (ITSC)*, pp. 1275–1282. IEEE, 2022a.
- Zhou et al. (2024)
Yang Zhou, Hao Shao, Letian Wang, Steven L Waslander, Hongsheng Li, and Yu Liu.
Smartrefine: A scenario-adaptive refinement framework for efficient motion prediction.
In *Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition*, pp. 15281–15290, 2024.
- Zhou et al. (2022b)
Yongchao Zhou, Ehsan Nezhadarya, and Jimmy Ba.
Dataset distillation using neural feature regression.
*Advances in Neural Information Processing Systems*, 35:9813–9827, 2022b.
- Zhou et al. (2022c)
Zikang Zhou, Luyao Ye, Jianping Wang, Kui Wu, and Kejie Lu.
Hivt: Hierarchical vector transformer for multi-agent motion prediction.
In *Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition*, pp. 8823–8833, 2022c.
- Zhou et al. (2023)
Zikang Zhou, Jianping Wang, Yung-Hui Li, and Yu-Kai Huang.
Query-centric trajectory prediction.
In *Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition*, pp. 17863–17873, 2023.

## Appendix A Appendix

### A.1 Implementation Details

During the dataset selection stage, we employ HiVT (Zhou et al., 2022c), HPNet (Tang et al., 2024), and QCNet (Zhou et al., 2023) as backbones with batch sizes of 16, 4, and 4, respectively, selected using the NaiveGreedy algorithm. For model evaluation, one mainstream model and one SOTA model are chosen per dataset: HiVT and HPNet for Argoverse 1 (Chang et al., 2019), and QCNet and DeMo (Zhang et al., 2025) for Argoverse 2 (Wilson et al., 2023). All models follow their official experimental settings and are trained on both selected subsets and complete datasets.

HiVT: Batch size 32, LR $3\times 10^{-4}$, weight decay $1\times 10^{-4}$, dropout 0.1. Its architecture includes 1 interaction module, 4 temporal and 3 global modules, 8-head attention (50 m radius), 6 prediction modes, and hidden sizes of 64/128.

HPNet: Batch size 16, LR $5\times 10^{-4}$ (same weight decay and dropout). It has 1 spatiotemporal and 2 tri-factor attention layers (50 m radius, 20-frame windows) with data augmentation: horizontal flipping (0.5), agent (0.05) and lane (0.2) occlusion.

QCNet: Trained for 64 epochs with AdamW (batch size 32, LR $5\times 10^{-4}$, weight decay $1\times 10^{-4}$, dropout 0.1). It uses a 128-dim hidden space, gated 8-head attention, layer normalization on MLP and attention layers, 3 iterations for the trajectory proposal, and 2 multi-context attention modules in both encoder and decoder.

DeMo: Trained for 60 epochs with AdamW (batch size 16 per GPU, LR 0.003, weight decay 0.01, cosine annealing with 10 warm-up epochs). No data augmentation is applied.

### A.2 Additional Experiments

#### A.2.1 Additional Results on Argoverse 1 and Argoverse 2.

We present comprehensive results for different models and data retention ratios on SSTP-selected datasets for both Argoverse 1 and Argoverse 2. The results, shown in Table [6](https://arxiv.org/html/2409.17385v3#A1.T6) and Table [7](https://arxiv.org/html/2409.17385v3#A1.T7), follow trends consistent with those observed in the 60%, 50%, and 40% retention settings, further validating the effectiveness of our approach.

**Table 6: Performance comparison results on Argoverse 1 with data retention ratios of 30%, 20%, and 10%. Pretrained HiVT-64 is used for sample selection. Evaluation conducted on HiVT-64, HiVT-128, and HPNet.**
| Methods | $\alpha$(%) | HiVT-64 | HiVT-128 | HPNet |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| mADE$\downarrow$ | mFDE$\downarrow$ | xMR$\downarrow$x | mADE$\downarrow$ | mFDE$\downarrow$ | xMR$\downarrow$x | mADE$\downarrow$ | mFDE$\downarrow$ | xMR$\downarrow$x |  |  |
| Random | 30 | 0.76 | 1.21 | 0.13 | 0.73 | 1.13 | 0.12 | 0.69 | 0.98 | 0.09 |
| Cluster | 0.74 | 1.14 | 0.12 | 0.71 | 1.07 | 0.11 | 0.68 | 0.97 | 0.08 |  |
| Herding | 0.73 | 1.14 | 0.12 | 0.71 | 1.07 | 0.11 | 0.69 | 0.96 | 0.09 |  |
| SSTP | 0.72 | 1.12 | 0.12 | 0.70 | 1.06 | 0.11 | 0.68 | 0.95 | 0.07 |  |
| Random | 20 | 0.78 | 1.25 | 0.14 | 0.75 | 1.18 | 0.13 | 0.71 | 1.01 | 0.10 |
| Cluster | 0.76 | 1.21 | 0.14 | 0.72 | 1.11 | 0.12 | 0.69 | 0.98 | 0.09 |  |
| Herding | 0.76 | 1.19 | 0.13 | 0.74 | 1.14 | 0.13 | 0.70 | 0.99 | 0.09 |  |
| SSTP | 0.74 | 1.16 | 0.12 | 0.72 | 1.11 | 0.11 | 0.69 | 0.97 | 0.08 |  |
| Random | 10 | 0.84 | 1.40 | 0.16 | 0.80 | 1.29 | 0.15 | 0.88 | 1.40 | 0.18 |
| Cluster | 0.82 | 1.37 | 0.16 | 0.76 | 1.20 | 0.13 | 0.87 | 1.39 | 0.17 |  |
| Herding | 0.80 | 1.29 | 0.14 | 0.80 | 1.29 | 0.14 | 0.81 | 1.26 | 0.13 |  |
| SSTP | 0.78 | 1.25 | 0.13 | 0.75 | 1.17 | 0.12 | 0.78 | 1.17 | 0.12 |  |

#### A.2.2 Performance Enhancement on HPNet.

To further validate the generalization capability of the SSTP method, we conducted the same experimental setup on HPNet. The test set was partitioned based on different scene densities, and the model was evaluated on these subdivided datasets. The results, presented in Table [8](https://arxiv.org/html/2409.17385v3#A1.T8), demonstrate the effectiveness of our approach. When the agent density is below 40, our method achieves performance comparable to models trained on the full dataset, with only slight increases of 0.019 in minADE and 0.028 in minFDE, while MR remains nearly unchanged. However, in high-density scenarios where the number of agents exceeds 60, models trained on our selected subset exhibit notable improvements, with minADE and minFDE reduced by 0.001 and 0.002, respectively. This advantage becomes even more pronounced in scenarios with more than 80 agents, where minADE is reduced by approximately 0.022, minFDE by nearly 0.073, and MR by nearly 3.5%. These findings confirm that SSTP effectively maintains dataset diversity while preserving model performance, regardless of the model used.

**Table 7: Performance comparison results on Argoverse 2 with data retention ratios of 30%, 20%, and 10%. Pretrained QCNet is used for sample selection. Evaluation conducted on QCNet and DeMo.**
| Methods | $\alpha$(%) | QCNet | DeMo |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| mADE$\downarrow$ | mFDE$\downarrow$ | xMR$\downarrow$x | mADE$\downarrow$ | mFDE$\downarrow$ | xMR$\downarrow$x |  |  |
| Random | 30 | 0.827 | 1.486 | 0.224 | 0.779 | 1.547 | 0.216 |
| Cluster | 0.821 | 1.523 | 0.222 | 0.765 | 1.498 | 0.207 |  |
| Herding | 0.819 | 1.505 | 0.213 | 0.762 | 1.506 | 0.204 |  |
| SSTP |  | 0.794 | 1.453 | 0.192 | 0.743 | 1.501 | 0.192 |
| Random | 20 | 0.878 | 1.586 | 0.242 | 0.827 | 1.645 | 0.238 |
| Cluster | 0.843 | 1.561 | 0.222 | 0.794 | 1.618 | 0.229 |  |
| Herding | 0.844 | 1.545 | 0.222 | 0.801 | 1.625 | 0.227 |  |
| SSTP |  | 0.832 | 1.530 | 0.211 | 0.783 | 1.591 | 0.212 |
| Random | 10 | 0.944 | 1.797 | 0.272 | 0.918 | 1.817 | 0.263 |
| Cluster | 0.911 | 1.772 | 0.251 | 0.878 | 1.779 | 0.247 |  |
| Herding | 0.908 | 1.743 | 0.246 | 0.882 | 1.774 | 0.243 |  |
| SSTP |  | 0.891 | 1.684 | 0.241 | 0.871 | 1.770 | 0.241 |

**Table 8: Performance Comparison of HPNet across different scene densities when trained on different sample set with $\alpha=50\%$.**
| HPNet | Agent$<$40 | Agent$>=$40 | Agent$>=$60 | Agent$>=$80 |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| mADE$\downarrow$ | mFDE$\downarrow$ | MR$\downarrow$ | mADE$\downarrow$ | mFDE$\downarrow$ | MR$\downarrow$ | mADE$\downarrow$ | mFDE$\downarrow$ | MR$\downarrow$ | mADE$\downarrow$ | mFDE$\downarrow$ | MR$\downarrow$ |  |
| Full | 0.611 | 0.833 | 0.062 | 0.875 | 1.202 | 0.127 | 1.111 | 1.463 | 0.201 | 1.596 | 1.771 | 0.276 |
| Random | 0.652 | 0.903 | 0.071 | 0.932 | 1.336 | 0.147 | 1.267 | 1.850 | 0.274 | 1.649 | 2.086 | 0.375 |
| Cluster | 0.642 | 0.889 | 0.069 | 0.916 | 1.291 | 0.143 | 1.193 | 1.669 | 0.250 | 1.636 | 1.977 | 0.325 |
| Herding | 0.650 | 0.901 | 0.070 | 0.921 | 1.303 | 0.145 | 1.215 | 1.673 | 0.252 | 1.637 | 2.001 | 0.342 |
| SSTP | 0.630 | 0.861 | 0.064 | 0.878 | 1.210 | 0.132 | 1.110 | 1.461 | 0.198 | 1.574 | 1.698 | 0.241 |

**Table 9: Performance comparison under extremely low data retention settings ($\alpha\leq 10\%$). For each retention ratio, SSTP consistently outperforms random selection, though performance degrades sharply when the data retention ratio drops below 10%.**
| $\alpha$(%) | Method | mADE$\downarrow$ | mFDE$\downarrow$ | xMR$\downarrow$x |
| --- | --- | --- | --- | --- |
| 10 | Random | 0.818 | 1.329 | 0.167 |
| SSTP | 0.781 | 1.253 | 0.132 |  |
| 5 | Random | 0.869 | 1.484 | 0.183 |
| SSTP | 0.843 | 1.416 | 0.161 |  |
| 2 | Random | 1.055 | 1.902 | 0.245 |
| SSTP | 1.015 | 1.871 | 0.220 |  |
| 1 | Random | 2.033 | 4.613 | 0.452 |
| SSTP | 1.589 | 3.416 | 0.386 |  |

**Table 10: Performance comparison of different backbone models, HiVT-64 and HPNet, trained on subsets selected by different strategies at varying data retention ratios.**
| Backbone | $\alpha$(%) | HiVT-64 | HPNet |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| mADE$\downarrow$ | mFDE$\downarrow$ | xMR$\downarrow$x | mADE$\downarrow$ | mFDE$\downarrow$ | xMR$\downarrow$x |  |  |
| HiVT-64 | 10 | 0.786 | 1.254 | 0.138 | 0.782 | 1.172 | 0.121 |
| 20 | 0.749 | 1.167 | 0.125 | 0.691 | 0.970 | 0.081 |  |
| 30 | 0.727 | 1.120 | 0.120 | 0.678 | 0.951 | 0.079 |  |
| 40 | 0.712 | 1.089 | 0.114 | 0.671 | 0.931 | 0.076 |  |
| 50 | 0.704 | 1.073 | 0.111 | 0.661 | 0.913 | 0.074 |  |
| 60 | 0.703 | 1.065 | 0.111 | 0.654 | 0.901 | 0.072 |  |
| HPNet | 10 | 0.796 | 1.277 | 0.143 | 0.739 | 1.050 | 0.095 |
| 20 | 0.753 | 1.193 | 0.125 | 0.699 | 0.975 | 0.082 |  |
| 30 | 0.732 | 1.143 | 0.119 | 0.681 | 0.943 | 0.078 |  |
| 40 | 0.717 | 1.087 | 0.112 | 0.670 | 0.924 | 0.075 |  |
| 50 | 0.708 | 1.079 | 0.111 | 0.664 | 0.917 | 0.074 |  |
| 60 | 0.703 | 1.067 | 0.108 | 0.657 | 0.910 | 0.075 |  |

#### A.2.3 Low Data Retention Setting.

We further studied the extremely low-budget regime with $\alpha\in\{10,5,2,1\}\%$.
As shown in Table [9](https://arxiv.org/html/2409.17385v3#A1.T9), the performance of all methods drops sharply below 10%.
Nevertheless, our SSTP method consistently outperforms random selection at the same budget levels.
In particular, at 1% retention, SSTP reduces minADE from 2.033 to 1.589 and MR from 0.452 to 0.386, showing that even with very limited data, informative subset selection remains beneficial.
This advantage comes from prioritizing scarce high-density scenes during budget allocation, ensuring long-tailed scenarios remain represented.

#### A.2.4 Generalizability and Robustness.

To further assess the generalizability and robustness of our proposed method, we evaluate the impact of different data retention ratios on SSTP-selected datasets using various backbone networks. The results, presented in Table [10](https://arxiv.org/html/2409.17385v3#A1.T10), demonstrate that SSTP is not only adaptable across a diverse range of trajectory prediction models but also effectively reduces dataset size while maintaining model performance across different feature extractors. These findings underscore SSTP’s robustness and broad applicability in real-world autonomous driving scenarios.

### A.3 Ablation Study

Effectiveness of Submodular Gain.
To isolate the effect of submodular gain, we conducted an experiment where data selection was based solely on submodular importance scores, without considering scene density balancing, as shown in Table [5](https://arxiv.org/html/2409.17385v3#S3.T5) line 3. The results indicate that using only submodular selection achieves a minADE of 0.724, lower than the 0.741 obtained through random selection, demonstrating that submodular-based sample selection improves data quality. However, it still underperforms compared to our full method.
This is because prioritizing sample informativeness without adjusting for scene density leads to a dataset biased toward certain complexity levels, ultimately hindering the model’s generalization ability.
In contrast, our full method, which integrates scene balancing with submodular gain, achieves the best performance across all metrics. These findings highlight the necessity of jointly considering both scene distribution balance and sample informativeness to construct an effective training dataset.

**Table 11: Performance comparison of models pretrained on the full dataset for different numbers of epochs. The pretrained weights are then used to initialize the model for SSTP, selecting 50% of the data.**
| Epoch | 0 | 5 | 8 | 10 | 15 | 64 |
| --- | --- | --- | --- | --- | --- | --- |
| mADE$\downarrow$ | 0.713 | 0.704 | 0.707 | 0.708 | 0.710 | 0.712 |
| mFDE$\downarrow$ | 1.083 | 1.073 | 1.076 | 1.074 | 1.083 | 1.080 |
| MR$\downarrow$ | 0.112 | 0.111 | 0.111 | 0.111 | 0.111 | 0.111 |

Impact of Pretrained Backbone Epochs.
To examine the influence of the pretrained backbone on subset selection, we conducted a series of experiments using models initialized identically but trained with different numbers of pretraining epochs. Taking HiVT-64 as an example, the official training setup involves training the model for 64 epochs using the full dataset. In our experiments, we varied the number of pretraining epochs as {0, 5, 8, 10, 15, 64} and analyzed its impact on subset selection, as shown in Table [12](https://arxiv.org/html/2409.17385v3#A1.T12). The results indicate that moderate pretraining is crucial for effective subset selection. When the number of pretraining epochs is set to 5, the subset selection achieves optimal performance, consistently outperforming other configurations across all data retention ratios. As the pretraining epochs increase, subset selection continues to provide significant advantages over other data selection methods but does not surpass the performance observed at epoch 5. For models without pretraining, minADE degrades noticeably compared to models pretrained for 5 epochs. In contrast, at 64 pretraining epochs, as the model has already converged, the impact of sample selection on gradient updates diminishes. Although subset selection performance remains competitive, it does not yield further improvements over moderate pretraining.

Figure: Figure 6: Visualization of different agent density scenarios in Argoverse 1.
Refer to caption: x6.png

Partition Interval. Given the number of agents in different trajectory prediction scenarios varies significantly, we examine the impact of scene density partition intervals on the selected subset to validate the generalizability of our method. Specifically, we divide the scenes in the Argoverse 1 dataset based on agent counts with partition intervals $\tau$ of {5, 10, 20}, forming multiple scene density categories. Within each category, we perform data selection. As shown in Table [12](https://arxiv.org/html/2409.17385v3#A1.T12), the subsets selected using different partition intervals result in comparable model performance, with minimal variations in minADE and minFDE. This consistency across different settings highlights the robustness of our method. This suggests that our approach generalizes well to other datasets. When applying this method, the partition interval can be adjusted based on the dataset characteristics: if the dataset has relatively few agents per scene, a smaller interval is preferable; whereas for datasets with a high variance in agent count, a larger interval may be more suitable to better balance scene density.

Figure: Figure 7: Visualization of different agent density scenarios in Argoverse 2.
Refer to caption: x7.png

### A.4 Qualitative Results

In Figure [6](https://arxiv.org/html/2409.17385v3#A1.F6) and Figure [7](https://arxiv.org/html/2409.17385v3#A1.F7), we present trajectory samples from three different scene densities in Argoverse 1 and Argoverse 2. In the top row, the scenes feature only a few agents, making them relatively straightforward for the model to predict. Because low-density scenarios dominate the dataset, the model tends to be biased toward these simpler cases, and the lower interaction complexity naturally reduces motion-forecasting uncertainty. In the middle rows, moderate-density scenarios display an increased number of agents. Here, interactions among vehicles become more frequent, creating added complexity for trajectory prediction. In the bottom row, high-density scenes pose significant challenges for trajectory prediction models. These crowded urban intersections and multi-agent interactions are underrepresented in standard datasets. The large number of dynamic agents in these scenarios compounds uncertainty, making it harder for the model to produce accurate predictions. Yet, these are precisely the scenarios that are most critical for ensuring safe autonomous driving. Our SSTP method addresses this imbalance head-on by emphasizing high-density samples. Through a more balanced yet compact selection of training data, the model becomes well-prepared for both common and complex scenarios. As a result, SSTP bolsters model robustness in high-density environments, leading to more reliable trajectory predictions in real-world urban traffic conditions.

## LLM Usage

We used ChatGPT solely for grammar correction and LaTeX formatting. They were not involved in research ideation, experiment design, or data analysis. All scientific contributions, methodology, and results are entirely the work of the authors.

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 2
presentation: 2
questions: In Table 5, the ablation study compares the effects of using partitioning submodular selection, and the full SSTP framework. However, the paper does not clearly explain how each component contributes to performance across different scene densities. Could the authors provide model performance results broken down by the number of agents for each ablation variant and interpret how partition and submodular selection influence each density scenarios?
soundness: 2
strengths: 1. The paper clears out a key limitation in current imbalanced trajectory prediction dataset.
2. The method is systemically validated across multiple models and two benchmark datasets.
3. The paper is clearly written and easy to follow.
summary: This paper proposes SSTP, a framework designed to improve data efficiency and scene-density balance in trajectory prediction. The authors observe that existing large-scale trajectory prediction datasets are heavily imbalanced, with low-density scenarios dominating and high-density cases underrepresented. SSTP tackles this issue through two-stage process: density-based partitioning of the dataset and gradient-based submodular selection to identify representative samples within each partition. Experiments on Argoverse 1 and Argoverse 2 show that SSTP achieves comparable performance to full-dataset training while reducing training cost and improving performance in high-density scenarios.
weaknesses: 1. The paper raises an important problem regarding scene-density imbalance in trajectory prediction, but the proposed approach offers only a partial solution. SSTP mainly combines known concepts such as density-based partitioning and submodular selection rather than introducing a fundamentally new modeling idea. As a result, the contribution lacks novelty.

2. The comparison in Table 3 seems unfair and potentially misleading. Since high-density scenarios are few and mostly retained while low-density ones are pruned, performance improvement in dense scenes is somewhat expected. However, the noticeable drop in low-density performance raises doubts about the true benefit of SSTP, as improving performance in complex cases while degrading simpler ones limits its overall practical value.

3. The definition of scene density based solely on the number of agents may oversimplify scene complexity. A high agent count does not always indicate complex interactions, and low-density scenes can still involve diverse or challenging trajectories. Considering additional measures such as interaction diversity, trajectory variance, or spatial proximity could lead to a more accurate and interpretable characterization of scenario complexity, thereby improving the effectiveness of the density-aware selection process.

4. The paper lacks qualitative examples or visualizations of the prediction results. Although the authors claim that SSTP improves performance in scenarios with many agents, it is unclear how the predicted trajectories differ in these complex, high-density scenes. Visual comparisons or representative examples would help illustrate the practical impact of the proposed selection strategy.

5. Evaluation on a more diverse dataset such as Waymo Open Motion Dataset would strengthen the generality of the method.

---

confidence: 5
contribution: 2
initial_rating: 4
presentation: 3
questions: See the last section.
soundness: 3
strengths: The proposed SSTP framework aims to construct compact and balanced training subsets via density-based partitioning and gradient-informed sample selection.
summary: The paper aims to address an important problem of reducing dependency on large-scale datasets in trajectory prediction, particularly under imbalanced data distributions.
weaknesses: While the motivation is relevant and the experimental results show improved efficiency, I have several conceptual and technical concerns that prevent me from supporting acceptance in its current form. 

My main concerns are as follows:  
The method relies heavily on partitioning data based solely on the number of agents per scene. This definition of "scene density" appears oversimplified, as it does not account for other critical factors such as interaction complexity, road topology, agent behaviors, or environmental conditions. In many real-world scenarios, a high agent count (e.g., in a parking lot) may not correspond to high risk, while sparse scenes can still be highly complex. A more nuanced and context-aware measure of scene difficulty is needed to justify the selection strategy.  

The notion of "long-tail" in trajectory prediction is not rigorously defined or statistically established. The authors should provide a clearer formulation of the long-tail phenomenon specific to autonomous driving, supported by dataset statistics (e.g., distribution plots, tail indices) rather than qualitative descriptions. This would help align the problem statement with the proposed solution.  

Although the paper motivates the work with the long-tail challenge, it does not sufficiently demonstrate that SSTP effectively addresses tail scenarios. Evaluation on dedicated corner-case benchmarks or tail-oriented splits is missing. A comparison with existing long-tail learning methods adapted to trajectory prediction would also strengthen the validity of the approach.  

The technical components of SSTP, like density partitioning and gradient-based submodular selection, appear to be combinations of existing ideas rather than fundamentally novel contributions. The authors should more clearly differentiate their method from related work in coreset selection or gradient-based sample pruning. Additionally, there is no sensitivity analysis on key hyperparameters, such as the similarity threshold in the submodular function, which raises concerns about robustness and generalizability.  

The requirement of a pretrained model for gradient extraction introduces non-negligible computational overhead, partly counteracting the efficiency gains of data selection. Moreover, submodular optimization in high-dimensional gradient spaces may not scale efficiently. These practical limitations should be acknowledged and quantitatively analyzed.

The experimental evaluation is limited to the Argoverse 1 and 2 datasets, which share similar characteristics. To establish generalizability, validation on more diverse benchmarks such as nuScenes or Waymo Open Motion Dataset is essential. Additionally, comparisons are made only against simple baselines (random, clustering, herding). Including recent advanced selection methods like GradMatch or CRAIG would offer a more convincing and competitive assessment.  

The writing structure could be improved by moving the related work section earlier to better contextualize the contribution. 

Furthermore, the analysis would benefit from including failure cases or scenarios where SSTP does not perform well, offering deeper insight into its limitations.  

In summary, while the idea of data-efficient training for trajectory prediction is valuable and timely, the current manuscript lacks sufficient methodological novelty, rigorous evaluation, and generalizability to meet the bar for ICLR. I encourage the authors to address these points in a future revision, particularly by refining the notion of scene complexity, expanding experiments across datasets and baselines, and providing more thorough analytical validation.

---

confidence: 4
contribution: 2
initial_rating: 4
presentation: 3
questions: Have you tried other possible feature formulations to select data within bins?
soundness: 3
strengths: 1. SSTP achieves consistent empirical gains at 40-60% budget, and beats augmenting, weighting, epoch-wise reselection, random, and high-density+random in terms of minADE, minFDE and MR. 
2. The paper has comprehensive ablation studies, such as Table 5-12.
summary: This paper proposes SSTP, a two-stage sample selection framework that constructs a compact yet density-balanced dataset for trajectory prediction. It consists of two stages: (i) first partition the data by scene density; and (ii) select a compact and density-balanced subset via gradient-based scores and a submodular objective. The goal is to reduce training time and mitigate long-tail imbalance. On Argoverse 1 and 2 Datasets and several backbones (HiVT, HPNet, QCNet, DeMo), SSTP claims comparable average metrics to full-data training, and improving high-density performance with around 50% budget.
weaknesses: 1. The paper uses a simple metric of agent count as density to divide the data into bins, but uses a relatively complicated construction for gradient features to select data within bins. One might wonder (1) if there are other possible definitions of density, and (2) if there are other possible constructions of data selection criterion within bins.
2. The density imbalance in autonomous driving dataset is a well known problem and it is known that re-balancing the data during training might hurt the overall performance on the entire dataset, as shown in Table 3 and 4.  
3. In practice, for example in an autonomous driving company, practitioners would mostly likely first train on the full dataset and only then apply SSTP to potentially improve the model. In this sense, SSTP is usable primarily a posteriori. It offers a possible quality boost rather than a budget saver.

---

confidence: 5
contribution: 2
initial_rating: 4
presentation: 3
questions: 1. Could you elaborate on the choice of using the raw number of agents as the metric for scene density? Have you considered alternative, potentially more robust heuristics, such as filtering out non-moving agents or agents outside a certain proximity, to get a more accurate measure of interaction complexity?

2. The results indicate a performance drop in low-density scenarios in exchange for gains in high-density ones. Could you comment on this trade-off? Is this an inherent limitation of the approach, or could the framework be adjusted to mitigate this degradation while preserving the benefits for rare scenarios?

3. The paper's contribution to improving long-tail performance is clear. However, why were alternative methods designed specifically for long-tail trajectory prediction not included as baselines for comparison, particularly regarding performance in high-density scenarios?

4. Given that the training time is only halved, and the selection process itself adds computational overhead, could you further justify the motivation from a training efficiency perspective, especially when sample selection literature often targets much larger (10x-100x) speedups?

5. I would be very interested to see results of applying this methodology to fields other than trajectory prediction, where training models are indeed very expensive and data has a lower quality, allowing for higher gains in selecting interesting samples.
soundness: 2
strengths: * The paper tackles a relevant and practical problem in trajectory prediction: the long-tail distribution of driving scenarios. Improving performance on rare but safety-critical, high-density scenes is a valuable contribution.

* The proposed framework, SSTP, is novel for this domain. Applying a principled, gradient-based submodular selection approach to create a compact and balanced dataset for trajectory forecasting is a new and interesting direction.

* The experimental results are promising, showing consistent improvements in high-density scenarios across multiple datasets and models (HiVT, HPNet). This demonstrates the effectiveness of the proposed density-balancing strategy.

* The work shows strong generalizability. The subset selected by SSTP using one model backbone transfers well when used to train other, more powerful models, saving computation time on the more expensive model.
summary: This paper introduces SSTP, a sample selection framework for trajectory prediction. The primary motivation is to address two challenges in existing large-scale datasets: the high computational cost of training and the imbalance where common, low-density scenarios dominate over rare, safety-critical high-density ones. The proposed method consists of two stages. First, it partitions the dataset based on scene density (number of agents) and pre-trains a model to extract gradient information. Second, it uses a submodular selection objective with these gradient-based scores to select a compact and representative subset, while explicitly up-sampling high-density scenarios. Experiments on the Argoverse 1 and 2 datasets show that training on a 50% subset selected by SSTP can achieve comparable performance to training on the full dataset, while significantly improving performance in high-density scenes.
weaknesses: * Motivation for sample selection is not entirely convincing. The paper highlights two goals: reducing training time and improving long-tail performance.

    * Regarding training time, a 2x speedup by using 50% of the data is a relatively modest gain, especially when considering the non-negligible computational overhead of the selection process itself. Training times for trajectory prediction models (e.g., ~8 hours for HiVT , ~300 hours for HPNet ) are not as prohibitive as in other domains like large language models, or diffusion models for image and video generation, which makes the efficiency argument less compelling. Also, sample selection methods typically aim for much larger data reduction (e.g., 90-99%) to justify the method and implementation complexity by a larger gain in training time decreasing. However the results for this higher filtering rates in table 9 shows the performance drop is very high and reduces the contribution of this work.

    * If the main goal is improving long-tail performance, the paper lacks comparisons to the significant body of work that directly addresses this issue in trajectory prediction (e.g., using contrastive learning or other re-balancing techniques ), some of which are cited in the related work section.

* The definition of "agent density" is overly simplistic. The paper uses the total number of agents in a scene as a proxy for interaction density. This heuristic could be misleading, as high agent counts might result from parked cars or agents that are far away and not interacting with the vehicle of interest. The fact that the "low-density" bucket includes scenes with up to 40 agents seems to support this, as it's unlikely a driver would interact with so many agents in a short 5-second interval. A more refined metric, such as counting only moving agents within a certain proximity, could provide a more accurate measure of scene complexity.

* Performance degradation on common scenarios. The results show that while SSTP improves performance on high-density cases, it leads to a noticeable drop in performance on more common low-density scenarios. For example, in Table 8, the HPNet model's minADE on low-density scenes (Agent<40) increases from 0.611 to 0.630 when trained on the 50% SSTP subset. For many practical applications, this trade-off might not be acceptable, and one might prefer to train for twice as long to achieve better performance on the majority of cases.

* Potential redundancy of sample selection. Publicly available datasets like Argoverse are typically already curated to contain a variety of interesting and non-trivial driving scenarios. This pre-selection might limit the potential gains of an additional sample selection layer, making it harder to find and discard truly redundant samples.