Kiran Kumar
Explorations in Building a Cost Prediction Platform
These notes are adapted from a talk I gave on building a cost prediction platform for a two sided marketplace. This is part one of two. It covers the model side: the problem, the architecture, and how one base model ended up serving several different teams. Part one stops right where the talk shifted gears, from model architecture to the system around it, and part two picks up exactly there, covering the repository structure, training automation, testing, deployment, and alerting that made the models actually usable in production.
The problem: predicting cost in a two sided marketplace
In a two sided marketplace, the platform sits between supply and demand and has to price both sides. On the supply side, cost is set through an auction: a set of suppliers bid to fulfill a piece of work, and the platform needs a reasonable expectation of what that auction will clear at. On the demand side, cost drives bid optimization and willingness to pay, and it doubles as a reserve price that anchors the supply side auction. Get the demand side cost wrong and the reserve price is wrong too, which affects contribution profit directly and shapes how well the platform can set expectations with customers about what a job should cost.
In practice, account executives had been adjusting prices largely on intuition and experience, which works reasonably well until the market shifts under them. A distributional cost forecast, a model that predicts a range of likely costs rather than a single number, gives a data driven alternative to that intuition without pretending the problem is simpler than it is. Optimal pricing stays genuinely hard even with a good model, because a lot of what moves cost is exogenous to the marketplace itself: incumbency effects from an existing relationship, macroeconomic conditions, and something as mundane as weather patterns along a route can all move the number a model has to predict.
Starting from a foundation model idea
Rather than building a separate point prediction model for every segment of the business, the team leaned on the foundation model idea that Bommasani and colleagues described in 2021: train one base model to learn the fundamental patterns in the data, then adapt that base model to a range of downstream tasks instead of training each one from scratch. For a cost prediction problem with many related but distinct segments (different lanes, different customer types, different modes), a single well trained base model is a much better starting point than dozens of independently trained models that all have to relearn the same underlying structure.
Model development
TabTransformer for tabular embeddings
The core of the platform is TabTransformer, the architecture described by Huang and colleagues in 2020. Most of the input data here is tabular: categorical fields like origin, destination, and customer segment, alongside continuous fields like distance and historical volume. TabTransformer passes the categorical fields through transformer layers to produce contextual embeddings for each category, then combines those embeddings with the continuous fields before the final prediction layers.
Two things make this a good fit here rather than just an interesting architecture. First, it delivers performance comparable to gradient boosted decision trees on tabular data, which had been the default choice for cost modeling before this. Second, and more useful for a platform meant to serve several teams, it produces embeddings, a dense vector representation of a category learned from data, that a tree based model simply does not give you. Those embeddings are reusable in a way a decision tree's splits are not. TabTransformer also supports pretraining with a masked or replaced token input scheme, the same general idea BERT style pretraining uses for text, applied here to table columns instead of words.
Predicting a distribution, not a point
The final layer of the model does not output a single predicted cost. It outputs the two parameters of a log normal distribution, a location and a scale. Cost data like this tends to be strictly positive and right skewed, a small number of unusually expensive outcomes pull the tail out, which a log normal fits far better than a model that assumes a symmetric, unbounded error like plain squared error regression would.
The model is trained by minimizing the negative log likelihood of the observed cost under the predicted log normal distribution, rather than minimizing squared error against a single point estimate. Predicting a full distribution rather than a point matters a great deal for how the output gets used downstream: a reserve price or a risk band needs a sense of the spread of likely outcomes, not just a best guess, and a point estimate simply cannot answer that on its own.
Fine tuning and transfer learning
Once a base model exists, two related techniques let it cover new ground without retraining from scratch every time, and the difference between them comes down to how much of the model is allowed to change.
Fine tuning keeps training the base model itself, just on a different dataset, for example a specific lane or customer segment. Only part of the model is allowed to move: the layers holding the broad, reusable pattern, typically the embeddings, are frozen, and the remaining layers adjust to the new data. Because the frozen layers never change, the model cannot forget what it already learned there while it adapts to the new segment.
One lighter weight variant explored on this platform was prefix tuning: rather than unfreezing any of the base model's own layers, a small set of trainable vectors is prepended to the input, and only those vectors are learned while the entire base model stays untouched. The change per segment is much smaller, which matters when many segments each need their own adaptation and the cost of updating a full copy of the model for every one of them adds up quickly.
Transfer learning goes further and freezes the entire base model. Rather than updating an existing model, it builds a new one: the base model stays fixed as a foundation, and a new layer sits on top of it to learn from features the base model was never trained on. This is the pattern a team reached for whenever a model needed access to a feature the original base model never saw, without giving up what the base model already understood about the underlying cost structure.
What mattered most for the platform as a whole, more than either technique on its own, was making the base model usable by teams who did not build it. Requiring every team to understand the internals of TabTransformer before they could benefit from it would have limited adoption to whoever happened to be on the modeling team. The fix was an API for model merging: a way for another team to combine their own features and data with the shared base model without needing to understand its architecture at all. That single decision did more for the platform's reach across the organization than any architecture change did.
References
Bommasani, R. et al. 2021. On the opportunities and risks of foundation models. arXiv preprint arXiv:2108.07258.
Huang, X., Khetan, A., Cvitkovic, M. and Karnin, Z. 2020. TabTransformer: Tabular data modeling using contextual embeddings. arXiv preprint arXiv:2012.06678.