Expedia Group Technology — Data
How Keras 3 Helped Modernise Expedia Group’s Lodging Ranking Stack
What happened when we treated a framework migration as an architecture modernisation — and cut P99 inference latency by two-thirds
Expedia Group™ has always been a market leader in providing personalised search experiences for travellers. As our ranking models evolved, we saw an opportunity not just to migrate to Keras 3, but to modernise the broader stack around it so we can better serve travellers. This led us to rewrite key parts of our pipelines that made model training 30% faster and cut P99 inference latency by two-thirds.
Our main focus in this blog is to discuss our improvements to lodging ranking — our service which handles users’ search requests and returns a personalised property ranking on top of a lightweight candidate generator. This service is frequently retrained to improve customer experience as we get new signals over time.
As such, our velocity to make changes needs to be fast, but more importantly we need to be able to build models which are state-of-the-art, enabling customers to easily find and book the most relevant property for their trip.
The problem with Keras 2, eager training and combined serving
The challenge was not that the old stack was broken. It was that it had gradually become a limiting factor.
On the training side, we were carrying dependencies on older components and missing out on modern Keras APIs, new optimisers, new layers and cleaner support for accelerator-oriented training. On the serving side, we had a path that was simple and operationally familiar, but not especially well suited to aggressively optimised inference for heavier ranking architectures.
This blockage on utilising new layers and optimisers has become more pronounced as LLMs see blistering amounts of innovation. Naturally, these new features are only available in Keras 3 or as separate backends from our default TensorFlow-only serving. As Machine Learning Scientists, this significantly impacted our innovation velocity, as we either had to recreate layers — such as mixture of experts implementations in TensorFlow — or we could not simply utilise the functionalities (such as new optimisers like Muon).
That mattered more as we started thinking about models whose compute is dominated by dense matrix operations. Mixer- and transformer-style ranking models benefit heavily from GPU execution, graph compilation, kernel fusion and mixed precision. If the stack cannot take advantage of those capabilities cleanly, latency becomes the main constraint on what you can train and serve.
The prior state: a coupled preprocessing and scoring graph
In our ranking stack, all the feature transformations happen in preprocessing layers generated by Kamae. Kamae is our open source preprocessing repository, covered in our 2025 RecSys paper, that maps TensorFlow operations to Spark operations. This enables fast, distributed, and one-time preprocessing before training, as well as produces a single preprocessing graph for serving. Kamae handles operations such as categorical lookups, one-hot encoding and other transformations over raw request features, including string-heavy inputs such as country codes.
That design led naturally to a coupled serving graph: as TensorFlow supported the string transformations we needed, the simplest production architecture was to bundle preprocessing and the neural network together into one fused TensorFlow graph. That made testing and deployment straightforward, and for earlier model generations, it was a very reasonable trade-off.
But the same architecture also created a ceiling
Because preprocessing and scoring were tightly coupled, we could not easily move the neural network into more aggressively optimised runtimes without first separating those graphs. The fused graph was stable, but it limited how far we could push compiled serving and hardware-aware optimisation.
Our string-based operations in preprocessing could only be run with eager training in contrast to JIT (Just-In-Time compilation).
JIT support is not new to Keras 3 but is a major part of our migration. JIT means Keras traces the model into a static computation graph and hands it to XLA (accelerated linear algebra, an ML compiler), which can fuse operations together enabling more efficient hardware utilisation. The main catch is that compiled execution is built around numeric tensor operations, not string processing, so preprocessing stages that depend on things like one-hot encoding from raw strings or other vocabulary-based transformations cannot cleanly be compiled into the same optimised graph.
What Keras 3 brought — multiple backends and new operations
Keras 3 was a fairly substantial shift for the ecosystem because it moved Keras away from being treated primarily as tf.keras and back toward being its own higher-level framework with a cleaner import keras API, a stronger backend-agnostic model definition style and support for writing more logic in framework-neutral primitives such as keras.ops rather than directly in TensorFlow ops.
In practical terms, that meant cleaner imports, more emphasis on portability, a newer native .keras serialisation format and some breaking API cleanup around how models, variables, reductions and layer resolution are handled, compared with older tensorflow.keras patterns.
Alongside that, the broader ranking stack around Keras also changed. Keras 3 made it much more natural to target newer execution paths and backends, while older TensorFlow-coupled utilities became less attractive to depend on long term.
One important example is TensorFlow Ranking: its deprecation was not a feature of Keras 3 itself, but it landed in the same general modernisation wave, and many teams migrating to Keras 3 have used that moment to drop the dependency entirely and move ranking losses, metrics, and supporting logic into their own libraries or directly onto standard Keras APIs.
This has also opened new opportunities for us to move towards more backend-agnostic code which could enable the usage of PyTorch or JAX. This means, should we wish to use highly optimised kernels which do not have native bindings in our chosen backend, our layers can easily migrate to match. We still predominantly use TensorFlow, as our preprocessing for most models still requires strings, but we have added flexibility.
Beyond Keras 3 migration — speeding up training and inference
Rather than treating Keras 3 as just an API update, we used the migration to modernise the core ranking stack. That included rewriting layers, losses, metrics, and training logic to be Keras 3 native and XLA-safe, while preserving parity with existing production behavior.
This allowed us to run testing within our existing training and ensure performance parity while benefitting from the speed-ups and the cleaner, more backend-agnostic APIs.
Kamae operations were rewritten to use backend-agnostic operations, with the exception of string operations that are maintained for legacy support in TensorFlow only. For our current models which rely on string inputs, nothing changes and we can still produce parity. For future models we can now choose XLA-safe preprocessing layers and as a result, support different backends should we want the entire model graph in JAX for instance.
Get Conor Worthington’s stories in your inbox
Join Medium for free to get updates from this writer.
Once the neural-network portion of the stack was cleanly XLA-safe and Keras 3 native, a more optimised serving design became possible. The key architectural move was to split the old fused graph into two stages to retain compatibility with string-containing preprocessing. Our new separate graphs consist of:
- Preprocessing in TensorFlow, where the string-heavy transformations still belong
- Neural network in Keras, comprising purely numeric operations and enabling the graph to be exported to an optimised runtime
We currently make use of ONNX and TensorRT to provide our faster runtime for the trained graph.
Speeding up the graph
ONNX is an open-standard intermediate representation for machine learning models. ONNX Runtime takes an ONNX graph and applies hardware-specific optimisations at load time (e.g. collapsing a MatMul + BiasAdd + ReLU) into a single fused kernel. Because ONNX Runtime is a purpose-built inference engine rather than a full training framework, it carries none of the overhead that TensorFlow or PyTorch bring to serving.
TensorRT is NVIDIA’s GPU inference optimiser and runtime, designed to take a neural network graph and compile it into a form that runs more efficiently on specific NVIDIA hardware. It improves latency and throughput by doing things like kernel fusion (see example illustrated below), hardware-specific optimisation, and lower-precision execution such as FP16 where accuracy allows. It further refines ONNX graphs in our case and enables us to optimise for the GPU we specifically are serving on.
Serving multiple graphs
With these more optimised runtimes, we can ensemble the preprocessing and model graph together at serving time with Triton Inference Server. Triton Inference Server is a serving platform which enables users to utilise multiple graphs with low cost memory passes and supports TensorRT natively. Triton configurations are able to be autogenerated as we already know how the graphs relate from preprocessing to training.
This separation matters because ONNX and its runtimes are much better suited to aggressively optimised numeric-inference than to vocabulary based string-transformation. ONNX Runtime can apply graph optimisations such as operator fusion and hardware-specific execution-provider selection, TensorRT can then go further on NVIDIA GPUs through device-specific compilation, kernel selection, and mixed-precision execution.
So the migration ultimately did two things at once: it modernised the training stack and it created a cleaner boundary that made compiled serving practical but still able to combine both graphs should we need singular graph TF-Serving for legacy support.
What we gained in training
Some of the biggest wins from the migration showed up before inference even began.
Our earlier staging pipeline was slow and heavily constrained by serial work. In the Keras 3 rewrite, our data handling was redesigned around more parallel Spark-based handling and improved shard management. This aligned with our design philosophy behind Kamae of having these heavy preprocessing ops parallelized by Spark. On representative production workloads, that represented a 75% speed up in writing our datasets from preprocessing.
Training also became substantially faster. On a representative production ranking model, training has become 30% faster since enabling JIT and rewriting our graph to be JIT friendly without any cost to our metrics.
As part of the rewrite, some of the legacy code which handled callbacks such as variable batch size was refactored to enable multi-GPU training. This has allowed us to throw much more compute at problems — for some of our models this means training is down from a working week to a day.
Those changes are important not just because they save time, but because they change the cost of experimentation. Faster staging and shorter training cycles mean teams can evaluate more ideas, iterate on larger models quicker, and spend less of their time waiting on experiments to complete.
What we gained in serving
The serving gains were just as important.
In isolated neural-network benchmarks, each step toward a more compiled runtime improved latency. On our old CPU-oriented serving model we saw ONNX providing major latency reductions both when we tested the model graph in isolation and tested it end to end.
We can see that while our preprocessing cost has not changed at all, we observe a reduction of 35% in total compute time due to the model graph speeding up. This is encouraging and with some of our model architectures being far more complex within the model graph than the preprocessing graph, which should enable us to better utilise our resources.
Utilising TensorRT means we can make use of hardware optimisation for mixed precision. Mixed precision means running most inference in lower precision where it is safe to do so, and in our Triton setup we use TensorRT with FP16 so the neural network can make better use of its resources. One question that naturally comes up with compiled GPU serving is whether FP16 changes ranking quality in ways that matter.
In practice, mixed precision looked like a safe trade-off for this class of model. Across our simulations, FP16 inference achieved an RBO (Rank-Biased Overlap, a measure of similarity of two ranked lists where disagreements at the top of a list matter much more than disagreements at the bottom) greater than 99% relative to the higher-precision baseline, indicating that ranking order stayed effectively unchanged while still benefiting from faster TensorRT execution.
We broadly see a substantial improvement across the board as we move to better optimised runtimes. Our P99 is higher than expected on our Triton ONNX serving pattern but still substantially lower than TF-Serving.
We then deployed our new service and compared it end to end, including preprocessing, hitting it with real traffic from Expedia®.
That kind of improvement matters because it reduces tail latency and creates room for more expressive architectures. In this case our prior P50 is now our P99. Further improvements can be seen on TensorRT FP16 where our P50 latency is reduced to only 10% of our original P99 latency on TF-Serving.
Why this mattered beyond one migration
The most important outcome of this work is more than any one benchmark. It is that the stack now has a clearer shape while retaining total parity with our old stack.
Preprocessing and neural-network execution are separated to keep strings isolated. Training is built around Keras 3 native, XLA-safe components. Serving still supports our old bundled TF-Serving approach but now has practical faster compiled paths through Triton. The same pattern can be reused for future ranking models that need better accelerator utilisation and lower latency than the previous architecture could comfortably provide.
Once a stack can train efficiently, export separate numeric model graphs cleanly and serve them through a compiled inference path, it becomes much easier to operationalize models that would otherwise be too latency-sensitive to deploy.
Further, we have been unblocked on utilisation of code from other backends now and are much more capable to run experiments with newer features such as optimisers like Muon. This has massively increased the velocity of experimentation.
Modernisation is often described as maintenance work. Sometimes, it is. But sometimes, modernisation is what makes the next generation of systems possible. For us, Keras 3 was exactly that kind of migration.
Acknowledgements
This would not have been possible without the fantastic work of the Expedia Group Search Ranking AI team as a whole. However, a special shoutout must go to George Barrowclough, James Shinner and Daniele Donghi who actively assisted with the Keras migration alongside maintenance of our core ranking repositories. A further thanks to Rajat Malhotra who assisted with productionalisation and testing of the new serving pattern. Finally thanks to Anne Morvan, Hafiz Tiomoko Ali, Youssef Achenchabe and Zoe Yang for their careful review of the article and valuable feedback.
[1]: Deploying Deep Neural Networks with NVIDIA TensorRT (Gray et al., 2017)
