Distilling Arbitration Logic: Bridging the Gap Between Deep Learning and ASIC Reality

Distilling Arbitration Logic from Traces using Machine Learning: A Case Study on NoC

2021-11-08
Yuan Zhou, Hanyu Wang, Jieming Yin, Zhiru Zhang
Summary
Problem
Method
Results
Takeaways
Abstract

This paper introduces an automated framework to "distill" high-performance arbitration logic from simulation traces using tree-based machine learning models as an intermediary. Applied to Network-on-Chip (NoC) port arbitration, the methodology converts complex Deep Reinforcement Learning (DRL) agents into compact, hardware-friendly linear model trees (LMT), achieving up to 282x area reduction while maintaining SOTA performance.

TL;DR

Researchers have developed a method to automatically transform complex Deep Learning models into ultra-compact hardware circuits. By using Linear Model Trees (LMT) as a "distillation" bridge, they reduced the area of a Network-on-Chip (NoC) arbiter by 282x compared to a neural network implementation, while significantly outperforming traditional human-designed FIFO and Round-Robin heuristics.

Background: The Price of Intelligence

In modern computer architecture, arbiters are the "traffic cops" of hardware, deciding which data packet gets access to shared resources. Traditionally, these are simple heuristics like Round-Robin. While Deep Reinforcement Learning (DRL) has shown it can "learn" much smarter traffic management strategies, the resulting models are too "heavy." Implementing a Deep Neural Network (DNN) directly in a router's port would consume massive amounts of silicon area and introduce unacceptable latency.

The challenge: How do we get the brains of a Deep Learning model into the body of a tiny, nanosecond-speed circuit?

Methodology: The Distillation Pipeline

The authors propose a three-step automated flow to move from high-level simulation to low-level logic:

  1. Policy Learning: A Deep Q-Network (DQN) is trained in a NoC simulator (gem5/Garnet). It learns to assign "priority scores" to packets based on features like local age, hop count, and payload size.
  2. Tree-Based Distillation: Instead of directly synthesizing the DNN, the authors train a Linear Model Tree (LMT) to mimic the DNN. Why? Because trees are essentially nested if-else statements—perfect for hardware conversion.
  3. Logic Synthesis: The LMT is converted into combinational logic. To keep it lean, weights are log-quantized (shifting instead of multiplying), and LASSO regression is used to prune unnecessary features.

Overall Flow and Model Conversion Figure 1: The proposed flow from simulation traces to hardware logic via tree models.

Why Linear Model Trees (LMT)?

Standard Decision Trees create "step functions" (abrupt changes). However, neural networks with ReLU activations are essentially piecewise linear functions. LMTs naturally fit this mathematical structure. By using LMTs, the authors could use shallower trees (lower latency) to achieve better accuracy than deep standard decision trees.

Converting Trees to Logic Figure 2: Physical implementation of LMTs using comparators, MUXs, and shift-add logic.

Experimental Battlefront: SOTA Comparison

The results are striking. When compared to the "Oracle" (Global-Age-Based arbitration, which is theoretically perfect but physically impossible to build at scale), the distilled logic (LMT) almost matches its performance.

ModelAvg. Latency (ps)Area (μm²)
MLP (DNN)25,65911,446
LMT (Depth 1)23,35445.9
FIFO2,113,3390.0
Oracle21,492N/A

The LMT reaching 23,354 ps latency is nearly 100x faster than the standard FIFO policy (~2,113,339 ps) while occupying a negligible 45.9 μm² of area—a fraction of the original MLP's footprint.

Performance Under Traffic Figure 3: Performance of distilled LMT logic under Uniform Random traffic compared to baselines.

Critical Insights & Future Work

The distillation process doesn't just reduce area; it acts as a form of regularization. Interestingly, the LMT models sometimes outperformed the original DNN agent, likely because the structural constraints of the tree prevented the "overfitting" that occurred during DRL training.

Limitations: While the model generalized well to Transpose traffic, it isn't a "universal" solution. The authors suggest that future work should focus on reconfigurable RL agents—hardware that can update its tree logic on-the-fly as data center traffic patterns shift.

Conclusion

This research provides a scalable bridge between AI research and hardware engineering. It proves that we don't need massive AI accelerators inside every router; we just need to use AI to "distill" better, simpler logic that can be baked directly into the silicon.

Find Similar Papers

Try Our Examples

  • Examine recent literature on automated hardware synthesis of reinforcement learning policies for resource management in many-core systems.
  • Which seminal papers first utilized "Knowledge Distillation" to convert neural networks into decision trees, and how does this paper's application to combinational logic circuits differ?
  • Investigate the feasibility of applying linear model tree distillation to other timing-critical hardware components such as branch predictors or cache replacement units.
Contents
Distilling Arbitration Logic: Bridging the Gap Between Deep Learning and ASIC Reality
1. TL;DR
2. Background: The Price of Intelligence
3. Methodology: The Distillation Pipeline
4. Why Linear Model Trees (LMT)?
5. Experimental Battlefront: SOTA Comparison
6. Critical Insights & Future Work
7. Conclusion