How Amazon Trains a 15M-Parameter Model to Replace LLM Calls for Product Upgrade Recommendations

E-commerce platforms need to recommend better products to customers, not just similar ones. A shopper buying a basic moisturizer should see a premium alternative that preserves their intent while offering improved formulation or certifications. This is trade-up recommendation: identifying higher-quality alternatives that keep the customer's shopping mission intact. Doing it well requires understanding subtle differences between products, which is exactly the kind of reasoning large language models excel at. But scoring hundreds of millions of product pairs through an LLM is operationally impractical. It costs millions of dollars and takes weeks. The team at Amazon Everyday Essentials Technologies proposes a two-level framework that distills LLM reasoning into a small, fast classifier and adapts its decision boundary per product category.

Why LLMs Cannot Scale to Catalog-Level Scoring

A large e-commerce store might contain hundreds of millions of product pairs. Running an LLM on each pair, even with batching, would cost millions in API fees and delay recommendation availability by weeks or months. The LLM also exposes only black-box outputs: a relation label and a free-text rationale, with no access to logits or hidden states. This rules out conventional distillation from teacher logits or intermediate features. The challenge is transferring the teacher's nuanced understanding of trade-up relations to a student model that consumes only two precomputed 768-dimensional product embeddings, with no LLM calls or text generation at inference.

The authors formulate trade-up as a directional, product-type-dependent relation classification. Given a base product b and candidate c, the model predicts whether c is a valid trade-up of b. The candidate must preserve the same shopping intent while providing evidence of an added benefit beyond price, quantity, or superficial variation. Differences in flavor, packaging, or pack size alone do not count. The taxonomy has four classes: similar/same tier, b is a trade-up of c, c is a trade-up of b, and incompatible/not meaningfully comparable. For binary evaluation, class 3 (candidate is a trade-up of the base) is positive.

Level 1: Reasoning Distillation from a Black-Box Teacher

The first level uses a retrieval-augmented few-shot LLM teacher. For each product pair, the teacher receives the product descriptions, product-type-specific trade-up criteria, and five dynamically retrieved expert demonstrations (selected by embedding similarity within the same product type). It generates a four-class relation label and a concise rationale of at most two sentences. These outputs form silver supervision: 1,019,241 teacher-annotated pairs across 29 product types, split 90/10 for training and validation.

The student is a lightweight embedding-pair classifier. It takes two precomputed embeddings, processes them through separate base and candidate branches, plus an ordered base-to-candidate branch that captures directionality. The three branch outputs are concatenated into a 2,296-dimensional representation, and an MLP maps this to binary or four-class logits. Total parameters: 15.5 million for the shallow model, 65.9 million for the deep variant. At inference, the model needs only the two embeddings. No text, no rationale, no LLM.

Three training objectives shape the student. Task supervision uses weighted cross-entropy for four-class labels or weighted focal binary cross-entropy for binary labels. Rationale alignment minimizes the L2 distance between a training-only projection of the pair representation and the encoded teacher rationale, forcing the student's internal representation to capture the semantic content of the teacher's reasoning. Contrastive distillation combines InfoNCE (contrasting each student projection against in-batch rationale embeddings) and a relational KL term that transfers pairwise similarity structure. The rationale projection and auxiliary objectives are discarded at inference.

The Interaction Between Rationale and Label Granularity

The results on a held-out human-annotated benchmark of 8,352 pairs reveal an important interaction. Within the same shallow architecture, adding rationale supervision to the binary model leaves AUC unchanged at 0.911. The rationale provides no benefit when non-trade-up cases are collapsed into a single negative class. But combining rationale supervision with four-class labels improves AUC from 0.912 to 0.924 (95% CI [0.918, 0.929]), a paired difference of +0.013. The reasoning benefit depends entirely on retaining the teacher's fine-grained relation structure.

This is not just an architectural curiosity. The rationale supervision encodes how the teacher distinguishes between similar/same-tier products and genuine trade-ups. When the label space collapses these distinctions, the rationale has nothing to anchor to. When the label space preserves them, the rationale guides the student's representation toward the same discriminative features the teacher uses. The deeper 65.9M model actually performs worse than the shallow one, reaching only 0.911 AUC with rationale and four-class labels. The extra capacity does not help because the bottleneck is supervision quality, not model expressiveness.

The student also recovers recall that the teacher misses. The retrieval-augmented teacher with five demonstrations achieves F1 0.749 with precision 0.967 and recall 0.610. The distilled student reaches F1 0.843 with precision 0.829 and recall 0.858 on the same benchmark. This comparison is specific to the evaluated teacher configuration, not evidence that the student outperforms frontier LLMs in general.

Level 2: Product-Type Test-Time Training

Trade-up criteria differ by category. What makes a battery a trade-up (higher capacity, better brand) differs from what makes a moisturizer a trade-up (improved formulation, relevant certifications) or pet food (ingredient sourcing, nutritional profile). The globally distilled student uses a shared scoring function across all product types, which cannot capture these differences.

Product-Type Test-Time Training (PT-TTT) addresses this by fitting lightweight LoRA adapters over the frozen student before scoring. For each product type, the same expert-labeled support examples that served as few-shot demonstrations for the teacher are reused as gradient-based supervision. The adaptation objective combines task loss on the support labels with a rationale-alignment term (lambda_reason = 0.1). The adapter is rank 8 with alpha 16, trained for 50 steps with AdamW at learning rate 1e-3. Each product type starts from the same global model and adapter initialization; adaptation is performed once per product type and amortized over all query pairs in that category.

At K=32 support examples per product type, PT-TTT improves AUC from 0.924 to 0.941 and average precision from 0.920 to 0.940. Performance plateaus between K=16 and K=32. The improvement is not explained by direct expert supervision alone: a pooled LoRA adapter trained on the same 928 expert-labeled examples across all 29 product types reaches only 0.929 AUC. Product-type-specific adaptation reaches 0.940. The macro within-product-type AUC increases from 0.910 to 0.925 for PT-TTT versus 0.914 for pooled LoRA, confirming improved within-type discrimination rather than cross-category score rescaling.

Rationale reuse during adaptation adds little. At K=32, reasoning-guided PT-TTT reaches 0.941 AUC versus 0.940 for label-only. The Level-2 gain comes primarily from product-type-specific supervised adaptation, not from reusing the rationale signal that shaped the global student.

Scale: 5,000x Faster, 10,000x Cheaper

On a proxy catalog of 100K randomly sampled product pairs, the distilled student on a single eight-GPU machine runs approximately 5,000x faster and 10,000x lower in estimated cost than direct LLM inference on the same workload. This is the practical payoff of the two-level design: the rationale encoder, projection head, and auxiliary distillation objectives are used only during training. Inference requires only the compact pair-classification model and precomputed embeddings. No generative inference, no token-by-token decoding, no prompt construction at serving time.

For a team operating at Amazon's scale, this means the recommendation pipeline can score the full catalog without the latency or cost bottlenecks of LLM inference. The precomputed embeddings are served once per product update cycle, and the 15.5M-parameter student classifies pairs in milliseconds.

Limitations and What Is Missing

The evaluation covers 29 product types and does not establish generalization to unseen categories. PT-TTT requires expert-labeled support examples and an explicit optimization step per product type, which adds operational complexity. Performance may depend on support-set composition and optimization randomness. The pooled-adapter and macro PT-AUC controls address two alternative explanations for the observed gain, but the authors do not fully separate product-type adaptation from simpler category-specific calibration. The support/golden split is pair-disjoint but not explicitly product- or brand-disjoint, leaving open the possibility of entity-level information leakage. The golden benchmark is designed for controlled model comparison and may have substantially different class prevalence from production candidate streams, so precision and F1 should not be interpreted directly as production positive predictive value.

The LLM comparison reflects the fixed retrieval-augmented teacher configuration evaluated in this study, not an exhaustive optimization over prompting, retrieval, or demonstration count. A stronger teacher could close the gap with the distilled student, though the 5,000x speed advantage would remain.

What This Means for E-Commerce ML Teams

The framework demonstrates that LLM reasoning can be transferred to a small, fast model without serving-time generative inference. The key insight is the division of labor: rationale-guided distillation shapes the global representation, while product-type-specific adaptation specializes the decision boundary. This separation means the expensive LLM is used offline to generate training signal, and the deployed model is a standard embedding-pair classifier that scales to hundreds of millions of pairs.

For teams building similar systems, the practical steps are: (1) use a retrieval-augmented LLM teacher to generate fine-grained labels and rationales, (2) distill into a compact classifier with alignment and contrastive objectives, (3) preserve the teacher's label taxonomy rather than collapsing to binary, and (4) apply lightweight per-category adaptation at serving time using the same expert demonstrations that trained the teacher. The 15.5M-parameter model with 50-step LoRA adaptation achieves AUC 0.941 on human-annotated data, which is competitive with much larger models and approaches the ceiling of the teacher's own accuracy.

Read the paper on arXiv