Deep Reading Notes: Uncertainty-Aware KD

1. Key assessment

This paper questions a common knowledge-distillation assumption: teachers are not equally reliable at every 2D keypoint. Occlusion, weak local texture, or degenerate views may make certain predictions vary substantially with initialization; weighting them like stable points can reinforce incorrect knowledge in the student.

The authors couple two mechanisms:

• UAKD converts teacher keypoint variance into confidence mass in OT;
• PFKD aligns corresponding receptive fields along UAKD’s keypoint transport plan, rather than independently matching whole feature maps.

The clearest contribution is using the same correspondences to constrain outputs and features. All three LINEMOD students outperform undistilled students and ADLP+FKD, yet the lighter SPNv2 student still trails its teacher by 6.2 points, so distillation does not eliminate the capacity gap.

2. Base pose pipelines and task boundaries

The paper compresses existing two-stage RGB pose estimators, rather than proposing a new PnP solver or pose head:

1. A network predicts object 2D keypoints from one RGB image;
2. Known 3D keypoints / bounding-box corners correspond to the 2D predictions;
3. PnP recovers object rotation and translation relative to the camera.

Two base networks:

• WDRNet: DarkNet + FPN outputs local 2D keypoint clusters at multiple scales, with keypoint-existence scores;
• SPNv2: an EfficientDet + BiFPN multitask network; only its 2D keypoint-heatmap head is distilled here.

Student training retains GT keypoint losses. KD supplies additional supervision rather than label-free compression; known-object 3D keypoints, camera intrinsics, and other PnP priors remain required.

3. Estimating teacher uncertainty

Only epistemic uncertainty is estimated: train $E=4$ or 6 teachers with identical architectures and different initializations. For an input image and keypoint $i$, collect $E$ 2D predictions:

$$\{\hat k^{T_j}_i\}_{j=1}^E.$$

Compute the means and variances of $x,y$ separately, use the ensemble mean as the teacher target, and define:

$$\sigma_i^2=\sigma_{x,i}^2+\sigma_{y,i}^2,\qquad u_i=\tanh(\sigma_i^2).$$

Teacher uncertainty confidence is:

$$\alpha^{T,c}=\mathbf 1-u.$$

If WDRNet also provides existence probability $\alpha^{T,e}$:

$$\alpha^T=\lambda\alpha^{T,c}+(1-\lambda)\alpha^{T,e},$$

The default is $\lambda=0.5$. Multiplication was also tried but performed worse empirically; specific numbers are not given.

Three boundaries matter:

• Variance depends on initialization, training randomness, and predicted-coordinate scale;
• tanh lacks temperature or scale calibration, so input resolution and coordinate normalization change saturation behavior;
• This does not model image-noise aleatoric uncertainty or pose covariance.

If ensemble members predict different numbers of keypoints, majority voting retains alignable points, with uncertainty set to 1 for the rest. This branch is not evaluated separately.

4. UAKD: unbalanced optimal transport

With $N$ teacher points and $M$ student points, OT allows $N\neq M$. Its cost is 2D Euclidean distance:

$$L_{\mathrm{pred}} =\min_\pi\sum_{i=1}^{M}\sum_{j=1}^{N} \pi_{ij}\lVert \hat k_i^S-\hat k_j^T\rVert_2.$$

Student / teacher confidence controls the row / column marginal masses of $\pi$, solved with unbalanced Sinkhorn. Uncertain teacher points have smaller $\alpha_j^T$, reducing mass transferred to the student.

The unbalanced formulation matters: strict mass conservation would force unreliable teacher points to distribute all their mass. Sensitivity to Sinkhorn regularization, the unbalanced penalty, and iteration count is not systematically reported, posing a reproduction risk.

5. PFKD: from keypoints to receptive fields

Ordinary feature KD often aligns entire teacher / student maps independently of prediction-level correspondence. PFKD uses the keypoint head’s convolutional structure to map predicted coordinates back to features:

• Multiply output coordinates by spatial scale $\delta$ and convert to integers for region centers;
• Derive region size from the receptive field determined by each head layer’s kernel $\kappa_l$ and stride $s_l$;
• Align teacher channels to student channels with a $1\times1$ convolution;
• Align spatial dimensions with average pooling.

For teacher region $R_i^T$ and student region $R_j^S$, UAKD’s $\pi_{ij}$ weights feature MSE. Ensemble teacher feature regions are first averaged across models.

Finally:

$$L_{\mathrm{distill}} =\gamma_pL_{\mathrm{pred}}+\gamma_fL_{\mathrm{feat}}, \quad \gamma_p=5,\ \gamma_f=0.1.$$

PFKD shares output and feature correspondences, at the cost of assuming a clear spatial mapping in the keypoint head. Backtracking convolutional receptive fields may not apply directly to transformer tokens, global attention, or deformable sampling.

6. Model size: efficiency evidence and gaps

Model	Backbone	Parameters	FLOPs	Input
WDRNet teacher	DarkNet-53	52.1M	36.51B	640×480
WDRNet student	DarkNet-Tiny	8.5M	18.84B	640×480
WDRNet tiny-H	DarkNet-Tiny-H	2.3M	4.75B	640×480
SPNv2 teacher	EfficientDet-D6	57.8M	288.27B	768×512
SPNv2 student	EfficientDet-D0	3.8M	12.1B	768×512

Parameter / FLOP reductions are clear, but real-time / efficient-inference claims lack latency, FPS, peak-memory, or energy evidence. Distillation needs 4–6 full teachers and averaged feature regions, increasing training resources substantially; deployment retains only the student.

7. Main LINEMOD results

Mean ADD(-S)-0.1d:

Student	Teacher	Student	ADLP+FKD	UAKD	PFKD	UAKD+PFKD
DarkNet-Tiny-H	92.9	81.9	84.8	88.1	88.2	89.0
DarkNet-Tiny	92.9	88.7	90.4	92.2	92.0	92.3
SPNv2 φ=0	92.8	84.8	85.9	86.2	86.0	86.6

Observations:

• Weaker Tiny-H gains more in absolute terms: +7.1, yet still trails its teacher by 3.9;
• DarkNet-Tiny nearly matches its teacher: 92.3 versus 92.9;
• SPNv2 has the largest remaining student–teacher gap, 6.2 after distillation;
• UAKD / PFKD individually exceed the ADLP+FKD mean, supporting both designs;
• The combination exceeds the best individual component by only 0.8 / 0.1 / 0.4: complementary, but not linearly additive.

Students exceeding teachers on individual classes may reflect distillation regularization and class variation. No multiseed variance is given, so differences of 0.1–0.5 require caution.

8. SPEED+ cross-domain results

SPNv2 φ=0:

Method	Synthetic pose	Lightbox HIL pose	Sunlamp HIL pose
Teacher	0.024	0.202	0.263
Student	0.033	0.368	0.401
ADLP+FKD	0.027	0.336	0.388
UAKD	0.024	0.288	0.362
PFKD	0.028	0.257	0.364
UAKD+PFKD	0.024	0.248	0.360

The combination matches the teacher on synthetic data and improves over the student by 0.120 / 0.041 on lightbox / sunlamp, yet trails the teacher by 0.046 / 0.097. PFKD alone gives 0.257 on lightbox, near the combined 0.248, suggesting feature alignment is particularly helpful for synthetic-to-real lighting transfer.

9. Additional analysis

The authors separate ensemble-averaging gains from uncertainty weighting:

• Ensemble alone without uncertainty ($\lambda=0$): Tiny-H +0.2, Tiny +0.5;
• Full UAKD: +6.2 / +3.5, respectively.

This supports weighted distillation as the main source of improvement rather than simply more accurate teacher ensembling.

The $\lambda$ curve peaks at 0.5, equally weighting uncertainty confidence and WDRNet existence scores. A single curve without calibration metrics, keypoint-level correlations, or pose-error stratification still does not establish $u_i$ as a well-calibrated probability.

10. Limitations and reproduction risks

Uncertainty-guided PnP and other compression techniques are future work. Current boundaries:

• Uncertainty does not reach PnP; deployment outputs no pose uncertainty;
• The ensemble estimates only epistemic variance, without accounting for training costs of 4–6 teachers;
• tanh(variance), coordinate scale, ensemble seeds, and member count directly affect weights;
• No public code; OT / receptive-field details must be reconstructed from formulas;
• Only LINEMOD and single-spacecraft SPEED+ are tested, without harder-occlusion LM-O, YCB-V, or T-LESS;
• No detection / crop frontend explanation or end-to-end latency;
• LINEMOD uses only ADD(-S)-0.1d, without stricter thresholds, AUC, or BOP AR;
• Decimal-scale differences lack multiseed means / variances.

11. Reproduction checklist

• Fix every teacher’s data, augmentation, and training configuration, varying only initialization;
• State whether keypoint coordinates are pixels or normalized values;
• Record the choice of ensemble size 4 / 6 and keypoint majority voting;
• Check variance→tanh and existence-score weighting with $\lambda=0.5$;
• Reproduce unbalanced Sinkhorn entropy regularization, mass penalties, and iteration counts;
• Derive feature receptive fields correctly from head kernels / strides;
• Distinguish teacher-feature averaging from ensemble-mean keypoints;
• Compare all KD baselines with identical PnP and detectors;
• Report parameters, FLOPs, GPU / CPU / edge latency, peak memory, and training costs together.

12. Overall assessment

The method trades training cost for compact deployment: teacher disagreement filters keypoint-level knowledge, while prediction and feature distillation share one OT plan. Three student results suggest some cross-architecture applicability. However, uncertainty-aware pose is easily overinterpreted: uncertainty is currently a training-time teacher epistemic signal, while final PnP and deployment remain deterministic.