Technical Summary: Heterogeneity-Aware Dataset Scheduling for Efficient Audio Large Language Model Training

Problem Statement

Training general-purpose Audio Large Language Models (ALLMs) on diverse datasets is critical for holistic audio understanding but is hindered by significant dataset heterogeneity. Current practices predominantly rely on uniform mixture (parallel training), where all datasets are sampled jointly. This approach often leads to conflicting gradients and slow convergence because datasets with disparate domains, annotation styles, and task complexities induce gradient inconsistencies. While multi-task learning and curriculum learning have been explored in traditional machine learning, existing approaches for ALLMs often rely on naive mixing or require architectural modifications and task-specific assumptions. There is a lack of systematic understanding regarding how dataset-level heterogeneity governs training dynamics and how it can be explicitly leveraged to schedule efficient training without architectural overheads.

Methodology

Theoretical Analysis: Convergence Bounds

The authors analyze multi-dataset AudioQA training through the lens of convergence analysis, treating each dataset as a distinct "client" in a federated-like setting. They establish theoretical bounds for two extreme regimes:

1. Parallel Training: Benefits from joint sampling which reduces effective gradient variance ($\sigma^2$) but is fully exposed to global dataset heterogeneity terms ($\beta^2, \zeta^2$).
2. Sequential Training: Suppresses heterogeneity effects by factors of $1/M$ and $1/M^{1/3}$, effectively decoupling interference, but suffers from the loss of joint variance reduction and risks catastrophic forgetting.

The analysis reveals a trade-off: parallel training is robust for mild heterogeneity, while sequential training becomes theoretically superior as dataset divergence increases. This motivates a strategy that interpolates between these extremes.

Proposed Solution: Grouped Sequential Training (GST)

The paper proposes Grouped Sequential Training (GST), a hybrid paradigm that partitions $M$ datasets into $K$ affinity-based groups $\{G_k\}$.

• Mechanism: Datasets are grouped to minimize intra-group heterogeneity. Training proceeds sequentially across groups, but within each group, datasets are trained in parallel (mixed).
• Theoretical Guarantee: By decomposing global gradient variance into intra-group and inter-group components, GST achieves tighter convergence bounds than both fully parallel and fully sequential training. It retains the variance reduction benefits of parallel mixing within groups while suppressing inter-group heterogeneity through sequential stages.
• Affinity Metrics: To implement grouping without prohibitive computational costs, the authors propose Gradient-Based Affinity, quantifying the pairwise distance between dataset-level gradients ($\|\nabla F_m - \nabla F_n\|^2$). This avoids the expensive empirical transferability estimation (which requires $M^2$ testing pairs).
• Progressive Implementation: To address practical hurdles like catastrophic forgetting and data loading overhead in strict sequential cycling, the authors introduce a Progressive Grouped Training schedule. This starts with the first group and incrementally expands the training pool to include subsequent groups ($G_1 \to G_1 \cup G_2 \to \dots$), approximating the optimal grouped sequential strategy while maintaining stability.

Key Contributions

1. Theoretical Framework: The paper establishes a convergence analysis for ALLMs that rigorously characterizes the impact of gradient variance and dataset heterogeneity, deriving bounds for parallel, sequential, and grouped sequential regimes.
2. Paradigm Shift: It proposes GST, which moves dataset affinity estimation from post-hoc analysis into the training loop design, interpolating between parallel and sequential optimization.
3. Scalable Implementation: A progressive training procedure is introduced that approximates optimal grouping strategies, enabling efficient large-scale training without architectural changes or model-specific assumptions.
4. Benchmark & Insights: The authors consolidate 14 diverse AudioQA datasets (covering speech, music, and environmental sounds) into a unified format, providing systematic insights and actionable strategies for heterogeneous multi-dataset training.

Experimental Results

The method was evaluated on 14 AudioQA datasets using the SALMONN-13B model (Whisper encoder + Vicuna-13B backbone).

• Full-Data Training:

  • Performance: GST variants consistently outperformed the standard parallel "Mix-all" baseline. GST-T2 achieved the highest unweighted average accuracy, while GST-G3 led in sample-size weighted average. All GST variants significantly surpassed the original SALMONN baseline and approached the performance of independent single-task training (the empirical upper bound).
  • Efficiency: GST variants achieved 30–40% faster convergence (reducing training time from ~4 days to ~2 days) compared to the parallel baseline.
  • Stability: Unlike pure sequential training, which suffered from catastrophic forgetting, the progressive GST implementation maintained high performance ceilings.

• Low-Resource Finetuning:

  • In scenarios with scarce data (250 samples per dataset), GST variants achieved performance comparable to the parallel baseline, demonstrating robustness against data sparsity.
  • The acceleration benefit was less pronounced here, attributed to a trade-off where the model may overfit to small initial groups before the data pool expands.

• Convergence Dynamics:

  • Analysis of training curves showed that strict sequential training exhibited volatile performance drops (catastrophic forgetting), whereas progressive GST showed steady, monotonic improvement.
  • "Stability-First" curriculum ordering (starting with high-affinity, low-variance groups) was found to be superior to reverse ordering or strict cycling, preventing the model from entering suboptimal regions of the loss landscape early in training.

Significance and Claims

The paper claims that affinity-aware scheduling is a critical, yet underexplored, factor in optimizing large-scale ALLMs. By explicitly managing dataset heterogeneity through GST, the authors provide a model-agnostic framework that:

• Accelerates convergence without requiring architectural modifications.
• Balances the stability of parallel training with the efficiency of sequential optimization.
• Offers a practical pathway to handle the inherent complexities of long audio sequences and cross-modal alignment in multi-dataset settings.

The authors conclude that their work provides both theoretical insights into the convergence behavior of heterogeneous multi-dataset training and a practical, scalable solution for developing robust Audio LLMs. They acknowledge limitations regarding the scalability to even larger models (e.g., 70B parameters) and the potential for dynamic, feedback-driven ordering mechanisms in future work.