Tao fine-tunes domain-specific models from prompts without requiring human-labeled response data.
2
The method uses a reward model to score multiple generated responses, then applies reinforcement learning to increase the likelihood of high-scoring outputs.
3
Tao can make smaller models competitive on narrow tasks, although repeated training can lead to reward hacking and performance depends on task difficulty and model diversity.
Summary
Prithviraj Ammanabrolu explains Tao, or test-time adaptive optimization, as a way to fine-tune models for private, domain-specific tasks using prompts instead of labeled answers. A policy model generates several responses, a reward model scores them, and reinforcement learning updates the policy toward higher-scoring outputs. The extra inference happens during training, so the deployed model keeps ordinary inference behavior and latency. Raj describes why response diversity matters, why simply raising temperature is insufficient, and how repeated optimization can overfit the reward model's errors. He says larger and more varied prompt sets improve results, although Tao can expand a small prompt sample internally. In his examples, 8B models perform well on narrow enterprise tasks, while harder tasks need larger models. He also discusses the cost-quality trade-off, the limits of small 1B and 3B models, and why models that have been heavily RLHF-trained may generate too little diversity for this process.
Tao customizes models without requiring labeled answers
Prithviraj Ammanabrolu describes Tao as a way for people to fine-tune models for their own domains and private data. The user supplies prompts that describe the tasks, rather than a large set of examples annotated by humans. Raj says the idea came from the cost and difficulty of collecting labels and from customers who knew what tasks they wanted but had little data showing how to perform them. Tao lowers that entry barrier by generating the training responses and feedback inside the process. The aim is a continuously improving, domain-specific model without asking every customer to build an expensive annotation pipeline.
The missing labels are replaced by synthetic responses and reinforcement learning
Tao generates its data under the hood. Given a prompt, the model produces one or more responses, receives feedback, and learns from its own mistakes. Raj compares this with personalized supervised learning: supervised learning copies human labels, while reinforcement learning lets the model generate candidate answers and learn which candidates receive better feedback. He argues that model mistakes can differ from human mistakes, so human-written data is not always the best source of learning. This approach still depends on data, but the system hides much of the data-generation and labeling work from the user.
A reward model makes verification cheaper than writing every answer
The Tao process uses a policy model to generate responses and a separate reward model to score them. The reward model assigns scalar values to outputs, then reinforcement learning adjusts the policy so high-scoring responses become more likely and low-scoring responses become less likely. Raj says the reward model is easier to train because judging whether an answer is correct is generally simpler than generating the correct answer. Databricks used a reward model called DBRM, trained on a broad range of enterprise tasks. Raj says the same reward model worked across tasks such as text-to-SQL and finance question answering.
Test-time compute is spent during training instead of deployment
The test-time part of Tao refers to extra inference used while training. The policy can generate multiple responses for a prompt, score them, and use those scores as learning signals. This consumes more compute before deployment, but the final model does not need to generate multiple answers for each user request. Raj contrasts this with systems that spend substantial reasoning compute on every live query. Once Tao training is complete, the customer receives a model with the expected deployment latency of that model, because the additional response generation was paid for ahead of time.
Response diversity matters more than generating many near-duplicates
Generating more responses is useful only when those responses contain different ideas or approaches. Raj says naive best-of-n sampling can waste compute because several outputs may be almost identical and receive the same reward. Turning the temperature up also failed as a general solution. Very high randomness may help creative writing, but enterprise tasks usually need more control. Tao therefore conditions later generations on earlier ones and tries to ensure that new responses are sufficiently different. The goal is to give the reward model fresh information rather than repeatedly scoring the same answer in different words.
Repeated optimization can turn reward-model errors into the training target
Raj says Tao training initially produced a graph that kept improving as the process ran for longer. The problem appeared when the policy extracted all the useful signal from the reward model and began learning its remaining noise. He describes this as an RL version of overfitting, or reward hacking. The optimized model may score well against the reward model while performing worse for the customer, because the reward is only a proxy for the desired behavior. Checkpoints allow the team to roll back. The ideal longer-term loop adds new customer feedback, updates the reward model, and then trains again with a better signal.
A deployed Tao model should not be retrained on the same prompt set indefinitely. Raj says the initial process already optimizes heavily over those prompts, so repeating them produces little new information after the useful signal has been extracted. Retraining makes sense when the customer has new tasks and can provide prompts that differ in substance, rather than the same requests rewritten. Tao also works to expand a small prompt sample into a broader set that covers different ways users might ask for the same task. This makes prompt diversity more useful without requiring every customer to submit thousands of examples.
Small models can match larger models on narrow enterprise work
Raj says Tao can make models in the 8B range perform impressively when their scope is narrowed to a specific type of work. He still believes scaling is useful for harder and more general tasks, but many customers do not need a model with the deployment burden of a very large system. Smaller models can occupy a different point on the cost-quality trade-off, with more training compute buying better task performance. Further distillation below 8B worked for some tasks, but Raj saw sharp drops on harder benchmarks with 1B and 3B models. He says the gap between 3B and 8B was larger than the cost difference suggested.
Model diversity before Tao affects how well the method works
Tao needs a policy model that can produce a diverse range of candidate responses. Raj says models that have been heavily RLHF-trained can have very spiky output distributions, where they are constrained to a small set of responses or repeat safety language. That leaves the reward model with less useful variation to score and gives reinforcement learning less signal to work with. Models with more generation diversity tend to work better with Tao. Raj says customers do not need to make this choice themselves because the team evaluates which models have suitable distributions and can sometimes reduce the spikiness before training.
"The reward model is easier to train to judge whether something is correct or not as opposed to actually trying to come up with the correct answer yourself."Prithviraj Ammanabrolu11:51
Who should watch
You have prompts for a private domain but lack a large set of labeled responses.
You are weighing a smaller fine-tuned model against the cost and latency of a much larger general model.
You are designing reinforcement-learning training and need to reason about reward hacking, sampling diversity, or when to stop.
