PaperlayerView exact revision on arXiv

Source-linked AI summary

DARE: Aligning LLM Agents with the R Statistical Ecosystem via Distribution-Aware Retrieval

Maojun Sun, Yue Wu, Yifei Xie, Ruijian Han, Binyan Jiang, Defeng Sun, Yancheng Yuan, Jian Huang

arXiv:2603.04743v1cs.IRcs.AIcs.CL

TL;DR

LLM agents underuse rigorous R statistical methods because they have limited R capability and existing retrieval methods overlook data-distribution constraints. The paper introduces RPKB, DARE, and RCodingAgent; DARE achieves 93.47% NDCG@10 and improves downstream statistical analysis performance by up to 56.25%.

  • Problem

    LLM agents have limited capability with R, while existing retrieval approaches focus on function semantics and overlook data-distribution conditions relevant to statistical-method applicability.

  • Method

    The paper constructs RPKB from 8,191 high-quality CRAN packages, proposes DARE for distribution-aware function retrieval, and integrates it into RCodingAgent for R statistical workflows.

  • Results

    93.47% NDCG@10: DARE outperforms prior state-of-the-art embedding models by up to 17% using only 23M parameters, while integration into RCodingAgent improves downstream performance by up to 56.25%.

  • Takeaways & Limitations

    DARE provides a lightweight retrieval module for connecting LLM agents with the R statistical ecosystem and supporting automated statistical analysis workflows.

  • Takeaways & Limitations

    The framework’s retrieved tools are supplied through demonstration-style JSON augmentation, which may not fully capture hierarchical and compositional relationships among statistical tools.

Abstract

from arXiv · show

Large Language Model (LLM) agents can automate data-science workflows, but many rigorous statistical methods implemented in R remain underused because LLMs struggle with statistical knowledge and tool retrieval. Existing retrieval-augmented approaches focus on function-level semantics and ignore data distribution, producing suboptimal matches. We propose DARE (Distribution-Aware Retrieval Embedding), a lightweight, plug-and-play retrieval model that incorporates data distribution information into function representations for R package retrieval. Our main contributions are: (i) RPKB, a curated R Package Knowledge Base derived from 8,191 high-quality CRAN packages; (ii) DARE, an embedding model that fuses distributional features with function metadata to improve retrieval relevance; and (iii) RCodingAgent, an R-oriented LLM agent for reliable R code generation and a suite of statistical analysis tasks for systematically evaluating LLM agents in realistic analytical scenarios. Empirically, DARE achieves an NDCG at 10 of 93.47%, outperforming state-of-the-art open-source embedding models by up to 17% on package retrieval while using substantially fewer parameters. Integrating DARE into RCodingAgent yields significant gains on downstream analysis tasks. This work helps narrow the gap between LLM automation and the mature R statistical ecosystem.

1. Introduction

LLM agents remain limited in using R’s rigorous statistical ecosystem because retrieval based mainly on semantic similarity misses data-distribution constraints. The paper addresses this gap with RPKB, DARE, and RCodingAgent, reporting improved retrieval and downstream analysis performance.

  • R-based statistical methods remain underused because LLM agents predominantly operate within Python-centric ecosystems and have limited R capability.
  • Distribution-aware retrieval is needed because statistical-method applicability depends on sparsity, dimensionality, distributional assumptions, and modality-specific constraints beyond semantic intent.
  • 93.47% NDCG@10: DARE outperforms prior state-of-the-art embedding models by up to 17% while using only 23M parameters.
  • RCodingAgent combines iterative reasoning, tool retrieval, code generation, and execution-based validation across 16 R-based statistical tasks.
  • RPKB is a curated repository derived from 8,191 high-quality CRAN packages and functions spanning diverse statistical domains.
  • DARE is a lightweight, plug-and-play retrieval embedding model that incorporates data distribution into function representations.

2. Related Works

Prior data-science agents largely rely on Python, while existing retrieval and tool-learning approaches do not explicitly account for statistical data distributions. DARE adds a dedicated distribution-aware retrieval module for scalable access to R statistical tools.

  • Existing data-science agents predominantly rely on Python, while efforts to incorporate R remain limited and often produce lower-quality code.
  • Dense retrieval methods map queries and documents into a shared latent space, but DARE explicitly incorporates data distribution into representation learning.
  • DARE decouples statistical tool selection from generation through a dedicated, data-aware retrieval module.
  • When augmented with DARE, LLM agents can utilize a broader range of statistical packages and improve downstream performance.

3. Methodology

The methodology builds a curated R-function knowledge base, uses distribution-aware bi-encoder retrieval, and integrates it into an agent that generates and validates R analyses across diverse tasks.

  • RPKB construction: RPKB curates 8,191 high-quality R functions from CRAN after extracting documentation, filtering for statistical content, generating data profiles, and indexing the repository.The data profiles encode attributes such as modality, distribution assumptions, and dimensionality.
  • Problem formulation: DARE retrieves functions using both natural-language intent and a query-side data profile derived from dataset characteristics and user-query signals.Candidate functions combine documentation with structured profiles encoding modality, distribution, and dimensionality constraints.
  • DARE modeling: DARE uses a shared-weight bi-encoder to embed query-context and function-documentation pairs, scoring candidates with cosine similarity for efficient retrieval.Precomputed function embeddings support Maximum Inner Product Search.
  • DARE modeling: The retrieval encoder is fine-tuned with InfoNCE, treating each paired function as positive and other functions in the batch as negatives.The objective increases similarity for candidates satisfying both semantic and distributional constraints.
  • RCodingAgent and evaluation: RCodingAgent iteratively retrieves tools, generates R code, and validates execution, while retrieved metadata and examples guide tool invocation and code generation.The evaluation framework contains 16 execution-grounded R tasks spanning statistical and scientific domains.

4. Experiments

The experiments evaluate DARE against open-source embedding baselines, measure retrieval and efficiency, and test its effect on RCodingAgent across 16 statistical tasks and multiple LLMs.

  • Experimental setup: 245,730 synthetic queries were generated, split into 85% training and 15% testing, with one ground-truth item per query.Each function contributed 30 queries using five prompt templates spanning different complexity levels.
  • Experimental setup: DARE is compared with open-source embedding models including BGE-M3, E5-large-v2, Arctic-embed-l, Jina embeddings, MXBAI, UAE-Large, GTE-large, MPNet, and MiniLM.The baselines cover substantially larger state-of-the-art embedding models as well as DARE’s base model.
  • Downstream evaluation: RCodingAgent is evaluated with and without DARE across 16 statistical tasks and six LLMs to assess robustness across model capabilities.The evaluation compares the same agent settings under retrieval-enabled and retrieval-disabled conditions.
  • Evaluation metrics: Retrieval evaluation uses Recall@k, NDCG@k, and MRR@k to measure top-k coverage, rank-sensitive relevance, and the position of the first relevant function.Recall@1 corresponds to top-1 accuracy, while NDCG discounts later correct results logarithmically.
  • Evaluation metrics: Success Rate measures the proportion of analysis tasks whose agent output matches the ground-truth answer.The experiments also report latency and throughput to assess real-time deployment feasibility.

5. Experimental Results

DARE establishes strong retrieval accuracy and efficiency on RPKB, then improves RCodingAgent performance across statistical analysis tasks. Its distribution-aware representations address fine-grained statistical compatibility while maintaining low computational overhead.

  • 5.1. Performance on Retrieval: DARE achieves an NDCG@10 of 93.47%, outperforming the strongest baseline by a relative margin of 17.8%.It establishes a new state of the art across all reported retrieval metrics.
  • 5.1. Performance on Retrieval: DARE uses 23M parameters yet substantially outperforms competing models with up to 568M parameters.Compared with the base model, NDCG@10 rises from 61.27% to 93.47%.
  • 5.1. Performance on Retrieval: 87.39% Recall@1 and 98.63% Recall@10 show strong top-ranked accuracy and broad retrieval coverage.DARE also reaches an MRR@10 of 91.76%, reflecting consistent ranking quality across queries.
  • 5.1. Performance on Retrieval: DARE conditions retrieval on distributional characteristics to distinguish statistically similar but distributionally distinct functions.This design targets cases such as differentiating glm from glm.nb, where semantic similarity alone can be insufficient.
  • 5.1. Performance on Retrieval: DARE reaches 8,512 QPS at 3.7ms latency, occupying the favorable high-throughput, low-latency region.The reported speed advantage over large baselines is approximately three to four times.
  • 5.2. Inference Efficiency Analysis: Adding DARE substantially improves end-to-end success rates across RCodingAgent evaluations on 16 statistical analysis tasks.Grok-4.1-fast rises from 18.75% to 75.00%, while gpt-5.2 improves from 25.00% to 62.50%.

6. Future Directions

The paper identifies future directions spanning LLMs’ native R proficiency, structured tool learning, and expansion of the statistical knowledge base.

  • Enhancing native R proficiency of LLMs: Future work may build large-scale R-centric corpora to improve LLMs’ native statistical reasoning and programming proficiency.
  • Advancing structured tool learning and utilization: The current demonstration-style JSON augmentation may not fully capture hierarchical and compositional relationships among statistical tools.
  • Advancing structured tool learning and utilization: Future tool-learning strategies include dynamic tool abstraction, function-level reasoning graphs, and memory-based tool selection mechanisms.
  • Expanding and sustaining the statistical knowledge base: Although RPKB contains over 8,191 high-quality R functions, the broader R ecosystem includes substantially more domain-specific packages and specialized tools.
  • Expanding and sustaining the statistical knowledge base: The authors plan to open-source RPKB to encourage community-driven expansion toward a comprehensive, continuously evolving statistical knowledge infrastructure.

D. Training and Evaluation Details

The training and evaluation details describe DARE’s training-loss monitoring and RCodingAgent’s R-specialized operating instructions. The agent retrieves R documentation, prioritizes specified functions, and validates work through constrained execution guidance.

  • Training process: Figure 6 presents DARE training loss curves across epochs to demonstrate convergence behavior.
  • RCodingAgent system prompt: RCodingAgent is instructed to solve complex data-analysis tasks by writing robust, efficient, and tidy R code.
  • RCodingAgent system prompt: The agent’s guidelines restrict file access to designated paths, permit package installation, require concise numerical outputs, and prohibit harmful system actions.
  • Training process: Figure 7 presents evaluation loss curves for DARE during training.
  • Retrieved documentation: When retrieved R documentation is available, the agent must analyze it, prioritize specific functions, and follow documented usage and examples.

F. Construction of Statistical Analysis Evaluation Tasks

The evaluation framework constructs realistic R statistical-analysis queries from function metadata and tests whether retrieval follows statistical logic rather than package or function names.

  • Query construction: Queries are generated from R function metadata without naming the package or function, while embedding distribution assumptions, data modalities, and specific constraints.
  • Query construction: The four query styles range from vague beginner requests to parameter-focused technical logic, with rigorous researcher queries explicitly including statistical profiles.
  • Output constraints: Every query requires set.seed(123), a specified ground-truth value for verification, pre-loaded data, and no package-name or function-identifier leakage.
  • Output constraints: The example format asks for an R script using a data modality, distribution assumption, and task type, followed by printing a ground-truth metric.

G. Example of Statistical Analysis Task

The example task evaluates Independent Component Analysis on high-dimensional, non-Gaussian numerical data. Its generated queries vary in style but consistently require two components, reproducibility, and verifiable output.

  • Task specification: The example task is Independent Component Analysis for tabular, numerical, high-dimensional data with a non-Gaussian distribution.
  • Task specification: The ground-truth verification records df ica$IC1[1] = 0.5040647 and df ica$IC2[1] = -1.372758.
  • Generated query styles: The beginner-style query asks for two independent signals to be extracted, added to the original data, and reported from the first row.
  • Generated query styles: The practitioner-style query requests two-component feature extraction, fortification into the dataset, set.seed(123), and first-row output.
  • Generated query styles: The rigorous researcher query explicitly combines non-Gaussian assumptions, blind source separation, two components, data-frame fortification, and reproducible output.
  • Generated query styles: The technical logic query requests density-based ICA with fortification, seed 123, and display of the resulting ICA columns’ top row.

H. Details of 16 Statistical Analysis Evaluation Tasks

The evaluation suite covers 16 representative statistical analysis tasks and tests whether distribution-aware retrieval identifies appropriate R functions for realistic data contexts. One genomic example shows DARE retrieving the specialized sharpr2 primitive and enabling executable task completion.

  • 16 representative statistical analysis tasks form the RCodingAgent evaluation suite.
  • The tasks span domains including financial quality scoring, ancestry score computation, linear mixed modeling, Gaussian mixture overlap, and factor-analysis preprocessing.
  • DARE conditions retrieval on genomic data metadata, including modality, sparsity, high dimensionality, completeness requirements, and ATAC-STARR constraints.
  • The task profiles include package and function names, descriptions, usage information, arguments, values, examples, and distribution-aware metadata.
  • The top retrieved document identifies sharpr2, whose documentation matches the query’s fragment-length and DNA-count filters for regulatory-score estimation.

I.2. Correct Implementation by the RCodingAgent with DARE

With DARE-provided documentation, RCodingAgent generates an executable R script that successfully solves the statistical analysis problem.

  • With the correct R package documentation, RCodingAgent generates an executable R script and successfully solves the problem.

Code Log: Task ID 4406

The code log contrasts a DARE-guided implementation using sharpr2 with a no-DARE implementation based on a generic log2-ratio heuristic. The specialized implementation applies the requested filters and produces the target regulatory score, whereas the heuristic produces a different value.

  • With DARE: set.seed(123) initializes the DARE-guided workflow, which loads the hidra dataset and the specialized sharpr2 library.
  • With DARE: The sharpr2 call applies fragment-length bounds of 150–600 bp, f_dna = 5, f_rna = 0, and sig = FALSE.
  • With DARE: 1326 reads remain after filtering and 957 tiled regions are called in the specialized workflow.
  • With DARE: -0.2185472 is the first estimated regulatory score produced for the first region by the DARE-guided implementation.
  • Without DARE: Without DARE, the agent manually filters the data and estimates scores with a generic log2 ratio instead of identifying sharpr2.
  • Without DARE: -0.3479233 is the first estimated score returned by the no-DARE heuristic, differing from the specialized implementation’s value.
Loading 2603.04743v1…