Fine-Tuning and Multi-Stage Model Optimization for Financial AI Agents

robinhood_markets · Finance · 2024 · source Agentic: yes · promotional bias: low · source-backed · extracted with prompt 0.5.6 (schema v2)

Generated from extraction.json — do not edit by hand. Every claim cites verbatim source evidence; stated means the quote directly asserts the claim, implied means it supports it indirectly.

Problem

Domain: financial-services AI agents

Robinhood needed production AI agents for financial workflows, especially the CX AI agent, to serve millions of users while balancing the generative-AI trilemma of cost, quality, and latency in regulated, user-facing financial products.

Constraints:

agents had to work for millions of users simultaneously
financial answers could not be vague, hallucinated, or below guardrail quality thresholds
multi-step agent flows amplified latency and quality problems across model calls
evaluation needed to combine product, engineering, data-science, automated, and human-feedback signals

Architecture

Robinhood built a multi-stage agent optimization system around Bedrock-backed foundation models, knowledge and tool access, a CX planner/communicator agent flow, trajectory tuning with dynamic few-shot examples retrieved from a vector database, LoRA fine-tuning through SageMaker/Bedrock Custom Model Imports, and a three-layer evaluation system. The architecture uses evaluation data to decide whether prompt tuning, trajectory tuning, or fine-tuning is the right intervention for a given model stage.

Diagram confidence: reconstructed (reconstructed from prose).

flowchart LR
  cx-ai-agent[CX AI agent]
  foundation-model-runtime[Foundation model runtime]
  knowledge-and-tool-access[Knowledge and tool access]
  planner-stage[Planner stage]
  communicator-stage[Communicator stage]
  three-layer-evaluation-system[Three-layer evaluation system]
  few-shot-vector-database[Few-shot vector database]
  lora-fine-tuning-platform[LoRA fine-tuning platform]
  cx-ai-agent --> foundation-model-runtime
  cx-ai-agent --> knowledge-and-tool-access
  cx-ai-agent --> planner-stage
  cx-ai-agent --> communicator-stage
  planner-stage --> communicator-stage
  planner-stage --> few-shot-vector-database
  three-layer-evaluation-system --> planner-stage
  three-layer-evaluation-system --> lora-fine-tuning-platform
  lora-fine-tuning-platform --> foundation-model-runtime

Drag to pan · scroll or use +/− to zoom · reconstructed from the source's described component topology

Components

Component	Role	Interacts with
`CX AI agent`	Customer-support agent flow that answers personalized financial support questions at production scale.	`Foundation model runtime`, `Knowledge and tool access`, `Planner stage`, `Communicator stage`
`Foundation model runtime`	Bedrock-backed inference layer that turns user questions into tool calls, plans, actions, and responses.	`CX AI agent`
`Knowledge and tool access`	Internal tools and account-history context used to personalize answers.	`CX AI agent`
`Planner stage`	Agent stage tuned through prompt, trajectory, and fine-tuning interventions.	`Communicator stage`, `Few-shot vector database`
`Communicator stage`	Final response generation stage in the agent workflow.	`CX AI agent`
`Three-layer evaluation system`	Evaluation control plane, LLM-as-judge, human feedback, and task-specific metrics used to diagnose quality and regressions.	`Planner stage`, `LoRA fine-tuning platform`
`Few-shot vector database`	Stores golden few-shot examples retrieved at inference time for trajectory tuning.	`Planner stage`
`LoRA fine-tuning platform`	SageMaker/Bedrock Custom Model Import path for parameter-efficient fine-tuning and deployment of smaller models.	`Foundation model runtime`

Data flow

A customer question enters the CX agent, where foundation models translate messy user questions into tool calls, plans, actions, and responses. The bot uses account history and internal tools for personalization. Evaluation data identifies weak stages; trajectory tuning injects retrieved high-quality few-shot examples for planner behavior, while LoRA fine-tuning is used when more durable quality/latency gains are needed.

“We needed state-of-the-art models that take user question and translate this mass user questions in tool calls, in plan and actions, in high quality responses.”transcript, stated

“we do a similarity check or the embedding similarity, and then we retrieve those 5 or 10 few-shot examples from that vector DB, and that uplifts the quality.”transcript, stated

Why this architecture

Robinhood layered prompt tuning, trajectory tuning, and LoRA fine-tuning because each intervention has different cost/latency/quality economics: prompt tuning is cheap but limited, trajectory tuning can improve quality with dynamic few-shot examples but grows context, and fine-tuning can move a smaller model to production quality without the same runtime token burden.

“we apply prompt tuning, we apply trajectory tuning, and we are selective when to apply fine-tuning to improve on the quality and getting the best of both worlds or the quality and latency and also the cost.”transcript, stated

“Now, you inject more dynamic future examples in your context. Your context grows up. And with that, your input token goes up, and that has an impact on your latency and also the cost.”transcript, stated

Models used: frontier foundation models · LoRA fine-tuned smaller models

Scale

The CX AI agent was designed to serve millions of customers.
“To serve millions of customers, we had to build a solution that is more advanced and goes through different stages of maturity.”transcript, stated
The fine-tuned CX stage produced more than 50% latency savings.
“with a lot of fine-tuned model, one of the stages in CX agent, we have received more than 50% latency savings.”transcript, stated
Latency dropped from three to six seconds to within one second.
“our previous model was giving us three to six seconds of latency. And with the lot of fine-tuned model, we cut down and brought it within one second.”transcript, stated
Robinhood retrieves 5 or 10 few-shot examples at inference time for trajectory tuning.
“we do a similarity check or the embedding similarity, and then we retrieve those 5 or 10 few-shot examples from that vector DB, and that uplifts the quality.”transcript, stated

Alternatives considered (roads not taken)

use a large frontier model for quality

Rejected because: frontier quality would blow up latency and cost budgets in production.

“If we go after quality and use a large frontier model, we will definitely get quality but it blows up our latency and cost budget.”transcript, stated

use a tiny model only for latency or cost

Rejected because: quality could fall below safety thresholds and cause guardrail blocks.

“But if we go with a tiny model to improve latency or cost, then the quality often dips and it goes below our safety threshold, and the answers get blocked by our guardrails.”transcript, stated

evaluate millions of chat logs manually

Rejected because: human-only evaluation did not scale, so Robinhood used LLM-as-judge backed by human feedback and curated eval data.

“it's not sustainable to eval millions of chat logs with, for instance, like by yourself or like human graded. So we heavily leverage LLM as a judge for end-to-end automated evaluation.”transcript, stated

Patterns

`tool-use-orchestration` (orchestration)

The CX agent translates user questions into tool calls, plans, actions, and responses, with access to internal tools and account history.

“We gave our bot knowledge of our internal tools, knowledge of the account history so that when customer ask for "what is my tax form" or "what is the status of my latest crypto transfer," our bot will have access to those tools”transcript, stated

`eval-driven-development` (evaluation)

Robinhood uses an eval-first loop with unified control plane, LLM judge, human feedback, curated datasets, and task-specific metrics to decide optimization steps and prevent regressions.

“we implemented a three-layer evaluation system. And this is especially valid for the end-to-end eval.”transcript, stated

“we heavily leverage LLM as a judge for end-to-end automated evaluation. But we don't stop at LLM as a judge. We also backstop it with human feedback and also hand-curated eval dataset”transcript, stated

`fine-tuning-adaptation` (model-operations)

Robinhood productionized LoRA fine-tuning for selected agent stages when prompt and trajectory tuning could not achieve the needed quality/latency tradeoff.

“We use LoRA heavily. This is what we have productionized, and this is what we believe that are easy to adopt and easy to roll out in production.”transcript, stated

`rag-retrieval` (context)

Trajectory tuning stores high-quality few-shot examples in a vector database and retrieves the most similar examples at inference time.

“The factuality, and then a vector DB, which is where you store your updated or high quality few-shot examples.”transcript, stated

“we do a similarity check or the embedding similarity, and then we retrieve those 5 or 10 few-shot examples from that vector DB, and that uplifts the quality.”transcript, stated

`latency-optimization` (serving)

Fine-tuning, model choice, hardware choice, prompt caching, and prompt compression are used to reduce latency and cost while preserving quality.

“We pick hardware, whether it's H100, A100, or some other hardware based on our needs, whether we want to optimize latency or cost.”transcript, stated

“move your static contents towards the beginning. So the model isn't building the attention KV cache on every user question. It reduces your cost and also reduces the latency.”transcript, stated

Tradeoffs

Optimize model stages through a ladder of tuning methods.

Chose prompt tuning, then trajectory tuning, then fine-tuning selectively over fine-tuning every problem immediately.

Rationale: The ladder uses cheaper interventions first and reserves fine-tuning for additional quality gains where the evals show it is needed.

“not using the fine-tuning hammer all over the place. Prompting and then the trajectory tuning for the dynamic future examples and then the fine-tuning for the additional quality gain.”transcript, stated

Use dynamic few-shot trajectory tuning.

Chose retrieve high-quality examples from a vector database at inference time over static planner prompting alone.

Rationale: Retrieved examples teach the planner domain-specific behavior with higher fidelity to the user question.

“You will not be injecting that example if somebody ask a question around credit cards. So we call it a dynamic few-shot example injection.”transcript, stated

Use LoRA for production fine-tuning.

Chose low-rank adaptation over full fine-tuning of all model weights.

Rationale: LoRA provides comparable fine-tuning performance with lower training cost, portability, and no inference overhead after merging.

“we talked about training way less parameters, and this also means very short training time, which means we can train for many, many use cases”transcript, stated

“we are getting the performance of a full fine-tuned model at a fraction of the cost and at no extra latency and at inference time.”transcript, stated

Gotchas (what bit them)

Agent quality, latency, and cost cannot be optimized independently.

Root cause: A multi-step agent pipeline amplifies the impact of any slow or low-quality model call.

Resolution: Robinhood uses staged prompt tuning, trajectory tuning, and fine-tuning so each intervention matches the failing stage and constraint.

“an agent isn't a single turn conversation. It's a multi-tier pipeline, where we are making N numbers of model calls. So if one of the calls in your entire agent flow is either slow or producing an inferior quality, it may jeopardize your end-to-end user-facing results”transcript, stated

Dynamic few-shot trajectory tuning improves quality but increases runtime context.

Root cause: Injecting examples into the prompt grows the input token count.

Resolution: Use trajectory tuning when it has quality uplift, then move to fine-tuning when durable improvements justify training.

Naive LLM-as-judge prompts were overwhelmed by account information.

Root cause: Throwing all account signals into the judge prompt overloaded the model.

Resolution: Robinhood built a two-tiered evaluation approach that first collects only the necessary signals for the user question.

“can we just throw all the account signals in a prompt and ask the model to do the evaluation? But it overwhelms, because there are just a lot of account information”transcript, stated

“We built a two-tiered approach, leveraging prompt tuning, where we first collect the necessary signals which are needed to answer a particular user question”transcript, stated

Evaluation

Robinhood evaluates at multiple layers: unified control plane for product/engineering/data-science success criteria, LLM-as-judge backed by human feedback and curated data, task-specific metrics for agent stages, and human annotations for training/eval data.

Metrics: LLM-as-judge end-to-end evaluations · human feedback and hand-curated eval datasets · task-specific categorical correctness · latency savings and P90/P95 tail latency · self-solve rate / successful bot answers

“This is important because evaluation is not just an engineering task. Most importantly, also a product requirement.”transcript, stated

Outcomes

CX agent stage achieved more than 50% latency savings
latency dropped from three to six seconds to within one second
quality parity was maintained with the trajectory-tuned frontier model
the approach is planned for extension to other Cortex agents

Robinhood describes the target as millions of users and reports production latency improvements in the CX agent, including long-tail latency reductions from P90/P95 above five seconds.

“with a lot of fine-tuned model, one of the stages in CX agent, we have received more than 50% latency savings.”transcript, stated

“We received or we were able to get or we were able to match the categorical correctness of the trajectory-tuned frontier model, and that was very essential for us to really productionize this.”transcript, stated