Distributional Reinforcement Learning with Information Bottleneck for Uncertainty-Aware DRAM Equalization

Muhammad Usama; Dong Eui Chang

cs.LGMar 5, 2026

Distributional Reinforcement Learning with Information Bottleneck for Uncertainty-Aware DRAM Equalization

Authors: Muhammad Usama, Dong Eui Chang

Organizations: Control Laboratory, School of Electrical Engineering, KAIST, Daejeon, 34141, Republic of Korea

Abstract

Equalizer parameter optimization is critical for signal integrity in high-speed memory systems operating at multi-gigabit data rates. However, existing methods suffer from computationally expensive eye diagram evaluation, optimization of expected rather than worst-case performance, and absence of uncertainty quantification for deployment decisions. In this paper, we propose a distributional risk-sensitive reinforcement learning framework integrating Information Bottleneck latent representations with Conditional Value-at-Risk optimization. We introduce rate-distortion optimal signal compression achieving 51 times speedup over eye diagrams while quantifying epistemic uncertainty through Monte Carlo dropout. Distributional reinforcement learning with quantile regression enables explicit worst-case optimization, while PAC-Bayesian regularization certifies generalization bounds. Experimental validation on 2.4 million waveforms from eight memory units demonstrated mean improvements of 37.1% and 41.5% for 4-tap and 8-tap equalizer configurations with worst-case guarantees of 33.8% and 38.2%, representing 80.7% and 89.1% improvements over Q-learning baselines. The framework achieved 62.5% high-reliability classification eliminating manual validation for most configurations. These results suggest the proposed framework provides a practical solution for production-scale equalizer optimization with certified worst-case guarantees.

Figures & tables

Panel of Figure 1 — Fig. 1 : The figure shows (a) the server memory system with double-sided DIMMs used to generate our dataset, and (b) a visualization of 1000 sample values for DRAM 1 from the dataset that plots the DRAM output waveform and the corresponding input waveform.

Panel of Figure 2 — Fig. 2 : Illustration of the signal validity labeling criteria. The rectangular window (80 mV $\times$ 35 ps) is shown in red. (a) An invalid signal where signal transitions intersect the window, and (b) a valid signal where no transitions occur within the window region.

Method	Mean	Worst-10%	Std Dev	Time	Rel.
	(%)	(%)	(%)	( $\mu$ s)	Time
DR-IB-A2C	37.1	33.8	1.42	167.4	1.0×
Std. Bayesian A2C	35.2	24.1	2.87	162.8	0.97×
Deterministic A2C [ 29 ]	36.8	30.9	1.89	158.5	0.95×
Q-learning [ 28 ]	26.1	18.7	3.45	1177.5	7.03×
DDPG [ 47 ]	15.5	11.2	2.98	676.9	4.04×

TABLE I : Performance Comparison for 4-Tap DFE Optimization. Worst-10% denotes CVaR at risk level

\alpha=0.1

Fig. 3 : Comparison of window area improvement distributions between Deterministic A2C (red) and Risk-Sensitive DR-IB-A2C (blue). The plot highlights the impact of CVaR optimization: while the mean performance is comparable, the risk-sensitive approach significantly shifts the lower tail (worst 10%) to the right, improving the 10th percentile performance from 29.8% to 38.1% (+8.4% absolute gain). This visualizes the trade-off between maximizing expected return and ensuring reliability for worst-case channels.

Method	Mean	Worst-10%	Std Dev	Time	Rel.
	(%)	(%)	(%)	( $\mu$ s)	Time
DR-IB-A2C	41.5	38.2	1.68	186.4	1.0×
Std. Bayesian A2C	39.4	28.9	3.12	181.2	0.97×
Deterministic A2C [ 29 ]	42.7	29.5	2.34	176.8	0.95×
Q-learning [ 28 ]	28.5	20.2	4.23	2429.5	13.03×
DDPG [ 47 ]	21.3	16.8	3.45	950.1	5.10×

TABLE II : Performance Comparison for Cascaded CTLE+DFE Optimization. Worst-10% denotes CVaR at risk level

\alpha=0.1

Latent Learning	Mean	Worst-10%	Silhouette	Recon.
Method	(%)	(%)	Score	MSE
IB Variational	41.5	38.2	0.72	0.0234
Standard AE	39.1	33.6	0.58	0.0189
AE + Classification	40.3	35.8	0.65	0.0212
Contractive AE	39.8	34.2	0.61	0.0198

TABLE III : Information Bottleneck vs Standard Autoencoder Comparison

Panel of Figure 7 — Fig. 4 : t-SNE visualization comparing latent spaces. (a) Information Bottleneck method exhibits clear cluster separation with anchor positioned near valid cluster centroid. (b) Standard autoencoder shows overlapping clusters with suboptimal anchor placement.

RL Variant	Mean	Worst-10%	KL Div.	Conv.
	(%)	(%)		Epochs
Distributional RL	41.5	38.2	0.142	178
Standard A2C	42.7	29.5	0.387	203
Categorical DQN	38.9	32.1	0.198	245
Implicit Quantile	40.2	35.8	0.165	192

TABLE IV : Return Distribution Modeling Comparison

Panel of Figure 9 — Fig. 5 : Return distribution visualization. (a) Histogram showing full distribution captured by DR-IB-A2C with mean = 24.8 mV $\cdot$ UI. (b) Value-based method with mean = 23.2 mV $\cdot$ UI demonstrating concentrated distribution that underestimates tail risks (KL divergence = 0.142).

Risk	Mean	Worst	Worst	Worst	Sharpe
Measure	(%)	10% (%)	5% (%)	1% (%)	Ratio
CVaR $\alpha$ =0.1	41.5	38.2	36.8	33.4	2.34
Mean-Var $\lambda$ =0.1	43.1	32.7	28.9	22.1	1.89
Mean-Var $\lambda$ =0.5	39.2	35.1	31.4	26.7	2.01
Expectile $\tau$ =0.1	40.8	36.5	34.2	30.1	2.18
VaR $\alpha$ =0.1	39.7	37.1	33.8	28.9	2.12

TABLE V : Risk-Sensitive Objective Comparison

Panel of Figure 11 — Fig. 6 : Risk-return trade-off analysis. (a) Pareto frontier showing DR-IB-A2C CVaR optimization achieves superior risk-adjusted performance. (b) Mean-variance methods demonstrate suboptimal tradeoff with conservative risk exposure across all operating points.

Distance	Mean	Worst-10%	Reward	Conv.
Metric	(%)	(%)	Variance	Rate
Sliced Wass.	41.5	38.2	0.187	1.0×
Euclidean L2	40.1	35.4	0.312	1.18×
Cosine	39.3	34.1	0.298	1.24×
Mahalanobis	40.6	36.2	0.245	1.09×
MMD	39.8	35.7	0.268	1.15×

TABLE VI : Latent Space Distance Metric Comparison

Panel of Figure 13 — Fig. 7 : Training stability comparison. (a) Reward curves showing Sliced Wasserstein distance provides smoother convergence with final reward = 24.8 mV $\cdot$ UI. (b) Euclidean L2 exhibits 67% higher variance (0.312 vs 0.187) demonstrating less stable policy learning dynamics (final reward = 23.5 mV $\cdot$ UI).

Method	Train	Test	Gen. Gap	KL(Q $\\|$ P)
	(%)	(%)	(%)
With PAC ( $\lambda$ =0.001)	40.8	38.9	1.9	2.34
Without PAC	43.2	35.7	7.5	8.97
L2 Reg. ( $\lambda$ =0.001)	41.5	37.2	4.3	4.12
Dropout (p=0.3)	40.9	37.8	3.1	3.45

TABLE VII : Generalization with PAC-Bayesian Bounds

Fig. 8 : PAC-Bayesian generalization analysis. (a) Learning curves showing 74.7% reduction in generalization gap with PAC regularization. (b) PAC bound tightness improving with training data according to

\sqrt{1/n}

scaling from Theorem III.4 .

Method	Clean	Noise	Noise	Adv.	Cert.
	(%)	$\sigma$ =0.01	$\sigma$ =0.05	$\epsilon$ =0.01	Rad.
Spectral (K=1)	41.5	40.8	38.6	39.2	0.0048
No Spectral	42.1	37.3	29.8	31.5	N/A
Batch Norm	41.8	38.9	33.2	34.7	N/A
Layer Norm	41.2	39.1	34.8	35.9	N/A

TABLE VIII : Robustness to Input Perturbations

Panel of Figure 17 — Fig. 9 : Robustness evaluation under perturbations. (a) Performance degradation with noise level $\sigma=0.01$ showing DR-IB-A2C maintaining 29.5% better performance than baseline. (b) Higher noise $\sigma=0.05$ demonstrating certified robustness region with empirical performance remaining within theoretical bounds.

Fig. 10 : Latent dimension analysis. (a) Performance versus dimension showing elbow at

l=11

with diminishing returns beyond. (b) Information Bottleneck lower bound saturation at

l=11

while training time increases linearly, confirming optimal selection.

Fig. 11 : Ablation study summary showing component contributions. Full DR-IB-A2C (highlighted) achieves best worst-case performance. Removing distributional RL or CVaR most severely degrades tail risk (

>12\%

reduction), validating core framework design.

Fig. 12 : Information plane analysis validating Theorem III.1 . Curves show rate-distortion trade-off for different $\beta$ values. Optimal operating point at $\beta=0.01$ (green star) dominates standard autoencoder (red X) in task relevance $I(\textbf{Z};Y)$ while using less compression $I(\textbf{Z};\textbf{D}_{o})$ .

Panel of Figure 21 — Fig. 13 : Wasserstein convergence validating Theorem III.2 . Scatter points show per-epoch measurements with inherent stochastic variance; solid line represents 10-epoch moving average. (a) Exponential convergence in Wasserstein distance with discount factor $\gamma=0.98$ showing linear trend on log scale despite fluctuations. (b) Bound tightness ratio fluctuating around 1.09, confirming theoretical guarantee holds with moderate variance typical of distributional RL.

Fig. 14 : Deployment classification based on CVaR performance and epistemic uncertainty. (a) Scatter plot showing three deployment categories with decision boundaries. High reliability configurations (green) satisfy

\text{CVaR}_{0.05}\geq 36\%

and

\sigma_{\text{unc}}<0.02

. (b) Distribution histogram showing 62.5% of CTLE+DFE configurations qualify for high reliability deployment.

Uncertainty	Mean	Expected	Calib.
Quartile	Perf. (%)	Perf. (%)	Error
Q1 (0–0.015)	42.8	42.3	1.2%
Q2 (0.015–0.020)	39.7	39.2	1.3%
Q3 (0.020–0.025)	35.4	35.9	1.4%
Q4 ( $>$ 0.025)	28.9	29.7	2.7%

TABLE IX : Epistemic Uncertainty Calibration

Panel of Figure 24 — Fig. 15 : Training dynamics over 300 epochs. (a) Actor loss convergence from 2.5 to 0.3 with smooth moving average. (b) Quantile regression critic loss stabilizing at 0.5. (c) Average reward evolution from 16 to 26 mV $\cdot$ UI demonstrating convergence with validation gap of 1.9%.

Risk-Sensitive RL

Risk-sensitive reinforcement learning optimizes policies for downside risk, tail outcomes, uncertainty, or safety constraints beyond expected reward.

See all 85 papers

Quantile Geometry Regularization for Distributional Reinforcement Learning
May 5, 2026Zhaofan Zhang, Minghao Yang, Rufeng Chen +2Distributional RL Quantile Regression
Information Hub, AI Thrust Hong Kong University of Science and Technology (Guangzhou) Guangzhou, Guangdong, China
Replay-buffer engineering for noise-aware quantum circuit optimization
Apr 23, 2026Akash Kundu, Sebastian FeldQuantum Architecture Search Reinforcement Learning
Delft University of Technology, Delft, The Netherlands · Quantum Computing Division, QuTech, The Netherlands
Risk and Anomaly Identification for Distribution Network Optimal Operation Based on Reinforcement Learning and Uncertainty Quantification
Sep 3, 2026Ziqi ZhangRisk-Sensitive RL OOD Detection
College of Automation Engineering, Nanjing University of Aeronautics and Astronautics, Nanjing 211100, China

May 5, 2026J/K move · Enter open · S save

Quantile Geometry Regularization for Distributional Reinforcement Learning

Zhaofan Zhang, Minghao Yang, Rufeng Chen, Sihong Xie+1

Information Hub, AI Thrust Hong Kong University of Science and Technology (Guangzhou) Guangzhou, Guangdong, China

Distributional RL Quantile Regression Risk-Sensitive RL Distributionally Robust Optimization Distributionally Robust RL

Open paper AskPDF arXiv

Quantile-based distributional reinforcement learning methods learn return distributions through sampled quantile regression, but their bootstrapped target quantiles may induce distorted or degenerate distribution estimates. We propose Robust Quantile-based Implicit Quantile Networks (RQIQN), a lightweight Wasserstein distributionally robust enhancement boosted from a quantile estimation perspective. We first reinterpret a snapshot of IQN loss as a collection of local empirical quantile estimation problems over sampled current fractions. We then robustify each local slot with a Wasserstein distributionally robust quantile estimation formulation, yielding a closed-form, fraction-dependent correction to the Bellman target. This correction directly addresses distributional degeneration: its median antisymmetry preserves the risk-neutral quantile average, while its monotonicity enlarges upper-lower quantile gaps and counteracts collapsed distributional spread. RQIQN thus regularizes quantile geometry without changing the underlying value objective or requiring additional sample set reconstruction. Finally, we empirically show that the proposed RQIQN outperforms other existing quantile-based distributional reinforcement learning algorithms in risk-sensitive navigation and Atari games.

Abstract

Risk-Sensitive RL

Explore similar work

Quantile Geometry Regularization for Distributional Reinforcement Learning

Replay-buffer engineering for noise-aware quantum circuit optimization

Risk and Anomaly Identification for Distribution Network Optimal Operation Based on Reinforcement Learning and Uncertainty Quantification