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
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:
- 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.
- 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-elsestatements—perfect for hardware conversion. - 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.
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.
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.
| Model | Avg. Latency (ps) | Area (μm²) |
|---|---|---|
| MLP (DNN) | 25,659 | 11,446 |
| LMT (Depth 1) | 23,354 | 45.9 |
| FIFO | 2,113,339 | 0.0 |
| Oracle | 21,492 | N/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.
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.
