Skip to content

NusRAT-LiA/LRE

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

1 Commit
 
 
 
 
 
 
 
 
 
 

Repository files navigation

LRE — Lightweight Learned Relevance Eviction

A tiny relevance scorer (logistic regression on 10 causal features, trained once on logged data) evicts context under a token budget at near-zero cost: no LLM in the eviction loop, no neural encoder at inference, CPU-only, and extractive (kept text is verbatim, so exact identifiers survive).

LRE is evaluated in two settings that share one method core (lre/):

  • Conversational QA (experiments/conversational_qa/): relevance-bound retention. Can a cheap learned scorer predict which past units a future question will need? Measured by leave-one-group-out AUC + recall@budget, plus a downstream answer-quality check.
  • Agents (experiments/agents/): continuity-bound tool use on AppWorld, integrated into the ACON agent. LRE matches ACON's LLM history-compression on task success while removing the per-step compression calls.

This repository ships code only. The ACON agent and the AppWorld benchmark are third-party and are fetched at pinned versions by the steps below; every result is regenerated locally and saved to results/summaries/*.json.


Repo layout

path what
lre/ the methodfeatures.py (10 causal features), labels.py (training labels), scorer.py (LR/GBT), select.py (budgeted extractive selector).
experiments/agents/ the agent study — run_study.sh (driver), configs, train_scorer.py, replay_eval.py, report.py, acon_cost.py, selection_quality.py, task_ids/, patch/.
experiments/conversational_qa/ the conversational-QA harness.
data/ dataset downloader for the conversational-QA experiments.
results/ created at run time; all numbers land in results/summaries/*.json (gitignored).

0. Common setup (both settings)

python -m venv .venv && . .venv/bin/activate
pip install -e ".[experiments]"
cp .env.example .env        # put your OpenRouter key in OPENROUTER_API_KEY=

All model calls route through OpenRouter via environment variables (OPENAI_BASE_URL, OPENAI_API_KEY); the agent model is openai/gpt-4.1-mini. The two settings are independent — set up only the one you need.


1. Conversational QA — end to end

Download the datasets once, then run. Per-baseline results are written to results/summaries/conv_qa_auc_<dataset>.json (AUC, recall@budget, per-method cost tier) and results/summaries/conv_qa_downstream_<dataset>.json.

Full run — all baselines, all questions (--dense + --llmlingua add the heavy BGE / LLMLingua-2 baselines; run_downstream is paid — it calls the QA/judge models via OpenRouter):

python data/download.py                                              # once: LoCoMo + LongMemEval_S -> data/raw/
python -m experiments.conversational_qa.run_retrieval_auc --dataset locomo      --dense BAAI/bge-small-en-v1.5 --llmlingua
python -m experiments.conversational_qa.run_retrieval_auc --dataset longmemeval --dense BAAI/bge-small-en-v1.5 --llmlingua
python -m experiments.conversational_qa.run_downstream  --dataset locomo      --n 100000 --budget 0.20
python -m experiments.conversational_qa.run_downstream  --dataset longmemeval --n 100000 --budget 0.20

Flags — drop or change these to shrink the run or the baseline set:

flag applies to default effect
--dense BAAI/bge-small-en-v1.5 run_retrieval_auc off add the BGE dense baseline; omit it for the lightweight baselines only. Accepts comma-separated encoders.
--llmlingua run_retrieval_auc off add the LLMLingua-2 baseline (needs torch; slow on CPU)
--max_items N run_retrieval_auc all LongMemEval only — cap to the first N questions (ignored for LoCoMo, fixed at 10 conversations)
--n N run_downstream 30 number of QA questions evaluated (both datasets); the code uses min(N, total), so a large N runs all
--budget F run_downstream 0.20 retention budget fraction

Without --dense/--llmlingua, only the lightweight baselines run (recency, memorybank, content_salience, and the LRE variants). Quick smoke pass:

python -m experiments.conversational_qa.run_retrieval_auc --dataset longmemeval --max_items 20
python -m experiments.conversational_qa.run_downstream  --dataset locomo --n 20

2. Agents (AppWorld + ACON) — end to end

2a. One-time setup of the third-party agent + benchmark

These versions are pinned — using different ones will break the patch and/or the frozen task-id buckets.

# (i) ACON, at the exact commit the patch was generated against
git clone https://github.com/microsoft/acon experiments/agents/acon
git -C experiments/agents/acon checkout d63f9ae18959dc7215ff62899c94c5e8c56847ae
git -C experiments/agents/acon apply ../patch/01_memory_lre.diff   # the single additive LRE branch
pip install -e experiments/agents/acon


git lfs install                                  # the install bundles are git-LFS
git clone https://github.com/StonyBrookNLP/appworld experiments/agents/appworld-src
git -C experiments/agents/appworld-src checkout a072b7a86e7c1d5b1d7175659d750ebb9b79f10a
pip install -e experiments/agents/appworld-src


appworld install

# download the benchmark data into ACON's appworld experiment dir (its runner resolves ./data):
cd experiments/agents/acon/experiments/appworld && appworld download data && cd -

# (iii) (optional) regenerate the task-id buckets; the committed task_ids/*.txt already hold them
cd experiments/agents/acon/experiments/appworld && python ../../../make_task_ids.py && cd -

The patch and the config keys it consumes are documented in experiments/agents/patch/PATCH.md. ACON's upstream code is otherwise unmodified; the provider (OpenRouter) and the corrected TGC metric are handled outside it.

2b. Run

Full run — the single command. The defaults already are the full grid (ARMS="nocomp fifo acon_history lre", BUCKETS="med hard", SEEDS="1 2 3", all tasks), so no env vars are needed. This is paid (many gpt-4.1-mini rollouts):

bash experiments/agents/run_study.sh all

all runs, in order: train the LR scorer on the train split → run all arms on the eval buckets → score (replay each task's logged actions through evaluate() for true TGC) → report (per-baseline TGC mean±SE, token cost, and the LRE-vs-ACON sign test). The report is printed and saved to results/summaries/agent_report.json.

Env vars — set these to shrink the run; defaults are the full configuration:

env var default (= full) effect
ARMS nocomp fifo acon_history lre which baselines to run
BUCKETS med hard difficulty buckets (Easy 57 / Med 48 / Hard 63)
SEEDS 1 2 3 independent seeds (use ≥3 for the reported result)
MAXTASKS 0 (= all) cap to the first N task ids per bucket and the train split
MAXITER 50 agent steps per task, not a task-count cap (leave as-is)
MAXPAR 8 max parallel jobs, launch stage only (leave as-is)

Quick smoke pass (3 tasks, one bucket, one seed):

MAXTASKS=3 SEEDS=1 BUCKETS=med bash experiments/agents/run_study.sh all

Run stages individually if you prefer: preflight (free sanity), train, launch (parallel arms), status, score, report. Offline token/latency analysis (free, after a run):

AW=experiments/agents/acon/experiments/appworld/outputs
.venv/bin/python experiments/agents/acon_cost.py --base "$AW" \
    --arms nocomp,fifo,acon_history,lre --buckets med,hard --seeds 1,2,3 --save
.venv/bin/python experiments/agents/selection_quality.py \
    --trajectories "$AW/openai_gpt-4.1-mini_train/train" \
    --scorer results/lre_model_lr_s90.joblib --save

The metric

TGC (task goal completion = all hidden AppWorld unit tests pass) is the metric. ACON ships world.evaluate() commented out, so the run's self-reported success is not trusted; replay_eval.py replays the logged actions into a fresh AppWorld and calls evaluate() to recover true TGC offline. report.py reads tgc.json only.


Where results land

file produced by
results/summaries/conv_qa_auc_<ds>.json run_retrieval_auc
results/summaries/conv_qa_downstream_<ds>.json run_downstream
results/summaries/agent_report.json run_study.sh all (the per-baseline agent results)
results/summaries/acon_cost.json, selection_quality.json the offline analysis scripts
results/lre_model_lr_s90.joblib run_study.sh train (the trained scorer)

results/ and data/raw/ are gitignored; every artifact above is regenerated by the steps here.


Pinned versions

dependency pin comments
ACON commit d63f9ae18959dc7215ff62899c94c5e8c56847ae the patch's line anchors target this commit; a different checkout may not apply
AppWorld (reports 0.2.0.dev0; not on PyPI) the frozen task_ids/*.txt (Easy 57 / Med 48 / Hard 63) were sampled from this build's data
agent model openai/gpt-4.1-mini via OpenRouter the operating point all arms share

About

No description, website, or topics provided.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors