[Paper Review] EBJR: Energy-Based Joint Reasoning for Adaptive Inference

Mohammad Kazem Akbari, Amin Banitalebi-Dehkordi+1 more|arXiv (Cornell University)|Oct 20, 2021

Advanced Neural Network Applications48 references4 citations

TL;DR

EBJR proposes an energy-based joint reasoning framework that adaptively routes inputs between a lightweight student model and a large pre-trained teacher model to achieve near-teacher accuracy with significantly reduced latency. By using a learnable energy function to route samples without retraining, EBJR achieves up to 1.8× speed-up on MS-COCO while maintaining 97% of the teacher's mAP, outperforming prior adaptive inference methods.

ABSTRACT

State-of-the-art deep learning models have achieved significant performance levels on various benchmarks. However, the excellent performance comes at a cost of inefficient computational cost. Light-weight architectures, on the other hand, achieve moderate accuracies, but at a much more desirable latency. This paper presents a new method of jointly using the large accurate models together with the small fast ones. To this end, we propose an Energy-Based Joint Reasoning (EBJR) framework that adaptively distributes the samples between shallow and deep models to achieve an accuracy close to the deep model, but latency close to the shallow one. Our method is applicable to out-of-the-box pre-trained models as it does not require an architecture change nor re-training. Moreover, it is easy to use and deploy, especially for cloud services. Through a comprehensive set of experiments on different down-stream tasks, we show that our method outperforms strong state-of-the-art approaches with a considerable margin. In addition, we propose specialized EBJR, an extension of our method where we create a smaller specialized side model that performs the target task only partially, but yields an even higher accuracy and faster inference. We verify the strengths of our methods with both theoretical and experimental evaluations.

Motivation & Objective

• To address the trade-off between inference latency and model accuracy in deep learning by enabling adaptive routing between small and large pre-trained models.
• To develop a method that does not require architectural changes or retraining of existing models, enabling plug-and-play deployment.
• To improve efficiency in real-world applications by dynamically selecting the optimal model path based on input complexity.
• To explore specialized student models that focus only on high-frequency classes to further enhance inference speed and accuracy.
• To provide a flexible, deployable solution for cloud and edge inference with minimal memory overhead.

Proposed method

• The method employs an energy-based routing module that computes a score for each input, determining whether to route it to the student (fast, small) or teacher (accurate, large) model.
• The energy function is learned during a one-time calibration phase using a small validation set, without fine-tuning the student or teacher models.
• The routing decision is based on minimizing an energy function that balances prediction confidence and computational cost.
• The framework is compatible with any pre-trained student-teacher pair, enabling plug-and-play integration into existing systems.
• A specialized variant of EBJR is introduced, where the student model is trained only on the most frequent classes to improve routing efficiency.
• The method supports both static and dynamic routing, with the router making decisions per-sample based on input complexity.

Experimental results

Research questions

• RQ1 Can a pre-trained student model be effectively combined with a large teacher model to achieve high accuracy with low latency without retraining?
• RQ2 How can an energy-based routing mechanism be designed to dynamically route inputs to the most appropriate model based on input complexity?
• RQ3 Does training a specialized student model on high-frequency classes improve inference efficiency and accuracy in real-world scenarios?
• RQ4 How does EBJR compare to state-of-the-art adaptive inference methods in terms of latency-accuracy trade-off across diverse downstream tasks?
• RQ5 What is the impact of data distribution (e.g., proportion of common vs. rare classes) on the performance of specialized EBJR?

Key findings

• EBJR achieves 97% of the teacher model’s mAP on MS-COCO with a 1.8× speed-up, outperforming the baseline EfficientDet-D4 and prior adaptive methods.
• On the same mAP level, EBJR provides 1.8× speed-up compared to the 1.3× speed-up reported by Zhou et al. (2017) for object detection.
• Specialized EBJR reduces FLOPs by approximately 1.5× compared to the teacher model while maintaining the same accuracy, under the assumption that 75% of inputs belong to the top-50 classes.
• When 50% or more of inputs belong to the top-50 classes, specialized EBJR with C̄=50 achieves optimal efficiency, outperforming both general EBJR and standard models.
• The method is orthogonal to existing dynamic inference techniques, and when applied on RS-Net, it further improves performance without architectural changes.
• Memory overhead is manageable by deploying the largest feasible student on edge and routing only hard samples to the cloud, preserving high accuracy while minimizing resource use.