Closing the Gap: Java Test Assertion Generation via Knowledge Distillation with Trident Loss

Abstract

Software testing is crucial in the software development process to ensure quality. However, automating test assertion generation remains a significant challenge in software engineering due to the need for both precise syntactic structure and semantic correctness. While large language models (LLMs) have shown impressive capabilities in generating test assertions, their high computational demands make them less practical for developers working in resource- constrained environments where cloud services are not a viable option. We present a knowledge distillation approach that trains a smaller student model (220M parameters) to mimic the behavior of a larger teacher model (770M parameters) through a novel Trident multi-component loss. Trident combines (1) a focal loss to focus training on hard-to-predict tokens, (2) a Jensen-Shannon Divergence (JSD) term to align the student with the teacher’s output distribution, and (3) a semantic similarity loss to preserve meaning, along with dynamic weight scheduling to balance these objectives. While knowledge distillation is established, its application to the nuanced task of generating test code assertions is underexplored. Our experimental evaluation on 7,000 Java unit tests demonstrates that the distilled student model achieves 90% of the teacher’s Code Quality Score while requiring 71% less memory. This significant reduction in resource requirements makes powerful LLM capabilities more accessible, particularly for developers in resource-constrained environments where cloud-based inference is not viable.