Outlier-Friendly Feature Distillation

• The paper introduces FLEX loss, a method for robust multi-scale feature alignment between teacher and student models, effectively handling outlier activations.
• It employs cross-normalized transformer feature alignment with percentile gating to suppress noisy gradients and ensure stable knowledge transfer.
• Generative rectified-flow synthesis enables the student model to produce teacher-level features, enhancing restoration quality and computational efficiency.

Outlier-Friendly Feature Knowledge Distillation (OFF) is a framework for robust knowledge transfer between teacher and student models, designed to address the instability of conventional feature matching in transformer-based architectures. OFF enables compact student models to synthesize “teacher-quality” features, especially in the presence of outlier activations, by introducing mechanisms such as FLEX loss, cross-normalized transformer feature alignment, and a generative rectified-flow process in latent space. Employed within the RestoRect image restoration system, OFF underpins improvements in restoration quality, stability, and computational efficiency across diverse degradation scenarios and datasets (Verma et al., 27 Sep 2025).

1. Feature Layer EXtraction (FLEX) Loss

The cornerstone of OFF is the Feature Layer Extraction (FLEX) loss, which aligns multi-scale features from teacher and student networks with explicit outlier handling under heterogeneous transformer architectures. Each teacher feature map $$\mathbf f_{\mathrm{teach}}^l\in\mathbb R^{C_l\times H_l\times W_l}$$ and student feature map $$\mathbf f_{\mathrm{stud}}^l\in\mathbb R^{C_l\times H_l\times W_l}$$ undergo channel-wise cross-normalization using student statistics:

$$\mu_{\mathrm{stud}}^{l,c} = \frac{1}{H_l W_l} \sum_{h,w} \mathbf f_{\mathrm{stud}}^{l,c,h,w},\quad \sigma_{\mathrm{stud}}^{l,c} = \sqrt{\frac{1}{H_l W_l} \sum_{h,w} \big(\mathbf f_{\mathrm{stud}}^{l,c,h,w} - \mu_{\mathrm{stud}}^{l,c}\big)^2 + \epsilon}$$

Both teacher and student features are normalized via these statistics to yield directly comparable representation scales:

$$\mathbf f_{\mathrm{teach}}^{l,c,\mathrm{norm},h,w} = \frac{\mathbf f_{\mathrm{teach}}^{l,c,h,w} - \mu_{\mathrm{stud}}^{l,c}}{\sigma_{\mathrm{stud}}^{l,c}},\quad \mathbf f_{\mathrm{stud}}^{l,c,\mathrm{norm},h,w} = \frac{\mathbf f_{\mathrm{stud}}^{l,c,h,w} - \mu_{\mathrm{stud}}^{l,c}}{\sigma_{\mathrm{stud}}^{l,c}}$$

A percentile-based outlier masking strategy is then applied. The $p$-th percentile $\tau_p^{l,c}$ of the absolute normalized student activations determines the gating threshold:

$$\tau_p^{l,c} = \mathrm{Percentile}\left(\left| \mathbf f_{\mathrm{stud}}^{l,c,\mathrm{norm}} \right|,p\right)$$

$$M_{\mathrm{rel}}^{l,c,h,w} = \begin{cases} 1 & \left| \mathbf f_{\mathrm{stud}}^{l,c,\mathrm{norm},h,w} \right| \leq \tau_p^{l,c}\ 0 & \text{otherwise} \end{cases}$$

Layer-wise resolution-aware weighting prevents domination by high-resolution feature maps:

$$w^{\mathrm{res}}_l = \max\left( \left( \frac{H_{\mathrm{base}} W_{\mathrm{base}}}{H_l W_l} \right)^{0.25}, 0.1 \right),\quad (H_{\mathrm{base}}, W_{\mathrm{base}})=(64,64)$$

Aggregated over all layers, channels, and spatial positions, the FLEX loss is

$$\mathcal L_{\mathrm{FLEX}} = \sum_{l} w_l^{\mathrm{layer}}\,w_l^{\mathrm{res}}\, \frac{ \sum_{c,h,w} M_{\mathrm{rel}}^{l,c,h,w} \left\| \mathbf f_{\mathrm{teach}}^{l,c,\mathrm{norm},h,w} - \mathbf f_{\mathrm{stud}}^{l,c,\mathrm{norm},h,w} \right\|_2^2 }{ \sum_{c,h,w} M_{\mathrm{rel}}^{l,c,h,w}+\epsilon }$$

This procedure ensures outlier activations—common in transformer networks—are masked, promoting stable distillation and robust feature alignment (Verma et al., 27 Sep 2025).

2. Cross-Normalized Feature Alignment in Transformers

OFF utilizes cross-normalized transformer feature alignment to address activation scale discrepancies between student and teacher models, which may arise from architectural heterogeneity. Both features are normalized using the student’s per-layer, per-channel mean and standard deviation, followed by percentile-based gating of spatial/channel elements based on the $95$th percentile of student activations. This approach directly suppresses noisy, outlier-driven gradients that otherwise destabilize learning. Feature distillation is performed over multiple transformer blocks and scales, ensuring that only reliably matched representations contribute to the loss landscape.

The practical summary of these normalization and masking steps is as follows:

Statistic	Description	Equation / Operation
Student mean, std	Compute per channel over spatial dims	$\mu_{\mathrm{stud}}^{l,c}$, $\sigma_{\mathrm{stud}}^{l,c}$
Normalization	Normalize teacher/student features	$\mathbf f^{l,c,\mathrm{norm}}$
Percentile gating	Mask activations above percentile $\tau_p^{l,c}$	$M_{\mathrm{rel}}^{l,c,h,w}$

This cross-normalized alignment mechanism is integral to extracting transferable features even when student and teacher have disparate activation statistics (Verma et al., 27 Sep 2025).

3. Generative Rectified Flow Process for Feature Synthesis

While conventional knowledge distillation often employs static feature matching, OFF redefines feature generation as a learned stochastic process in latent space via rectified-flow ordinary differential equations (ODEs). Starting from Gaussian noise $\mathbf z \sim \mathcal N(0,I)$, a linear path interpolates toward the target teacher feature $\mathbf f_{\mathrm{teach}}$:

$$\mathbf x_t = (1-t)\, \mathbf z + t\, \mathbf f_{\mathrm{teach}}, \quad t \in [0,1]$$

The instantaneous velocity along this path is given by

$$\mathbf v(\mathbf x_t, t) = \frac{d\mathbf x_t}{dt} = \mathbf f_{\mathrm{teach}} - \mathbf z$$

A small neural network $\epsilon_\theta(\mathbf x_t, t, \mathbf c)$, where $\mathbf c$ denotes an image encoding, is trained to predict this velocity:

$$\mathcal L_{\mathrm{vel}} = \mathbb E_{t, \mathbf z, \mathbf f_{\mathrm{teach}}} \left\| \epsilon_\theta(\mathbf x_t, t, \mathbf c) - (\mathbf f_{\mathrm{teach}} - \mathbf z) \right\|_2^2$$

During inference, the ODE is discretized with $N_{\mathrm{steps}}=3\!\!-\!5$ Euler steps:

$$\mathbf x_{t+\Delta t} = \mathbf x_t + \Delta t\, \epsilon_\theta (\mathbf x_t, t, \mathbf c),\quad \Delta t=1/N_{\mathrm{steps}}$$

Regularization is applied to the synthesized feature trajectories for stability:

• $$L_{\mathrm{trans}} = \sum_{i=1}^{N-1} \| \mathbf f^{i+1}_{\mathrm{pred}} - \mathbf f^i_{\mathrm{pred}} \|_2^2$$
• $$L_{\mathrm{target}} = \| \mathbf f^{N}_{\mathrm{pred}} - \mathbf f_{\mathrm{teach}} \|_2^2$$
• $$L_{\mathrm{cons}} = \sum_{i=1}^N \mathrm{cos\_dist}(\mathbf f^i_{\mathrm{pred}}, \mathbf f_{\mathrm{teach}})$$

Combined, the trajectory stabilization loss is

$$\mathcal L_{\mathrm{traj}} = 0.1\,L_{\mathrm{trans}} + 0.5\,L_{\mathrm{target}} + 0.2\,L_{\mathrm{cons}}$$

This process enables the student to generatively synthesize robust features, as opposed to direct regression, and is central to OFF’s outlier resilience (Verma et al., 27 Sep 2025).

4. Integration with Physics-Based Priors: Retinex, Anisotropic Diffusion, and Polarized Color

OFF is applied within a broader setting where teacher and student networks encode reflectance $R$ and illumination $L$ streams via Retinex decomposition. The teacher produces two feature streams—one from image encoding, one from Retinex encoding—which the student is required to synthesize using rectified-flow dynamics.

• Retinex Decomposition: Segregates extreme intensity changes (illumination) from high-frequency details (reflectance). Outlier gating (via FLEX) ensures that spurious feature spikes in either stream are masked during knowledge transfer.
• Anisotropic Diffusion: Teacher features are regularized for texture consistency ($L_{\mathrm{tex}}$), resulting in smooth yet edge-preserving reflectance features. Student-induced overshoots are suppressed by outlier masking.
• Polarized HVI Color Loss: The teacher’s color encodings are uniformized in a trigonometric color space, and any trajectory artifacts in the student’s synthetic feature flows are subjected to the same percentile masking, enforcing robust color matching.

The combination of these physics-based priors with rectified flow and outlier gating ensures that only the reliable, physically consistent teaching signals are transferred, minimizing the effect of outliers and instabilities during student training (Verma et al., 27 Sep 2025).

5. Empirical Evaluation and Performance Attributes

RestoRect, embodying OFF, demonstrates enhanced robustness, efficiency, and restoration quality across diverse image degradation benchmarks:

• Outlier Robustness: Ablating the FLEX percentile masking leads to PSNR drops of $>3$ dB and FID instability during distillation.
• Training Stability: Teacher models employing SCLN + QK normalization converge $\sim30\%$ faster; student models attain target FID in $3$–$4$ rectified-flow steps compared to $10+$ DDIM steps.
• Restoration Quality: On LOL-v1, results include PSNR $27.84$ dB ($+2.5$ dB vs. RetiDiff), SSIM $0.945$, FID $38.67$. Gains persist across underwater (UIEB: $+1.76$ dB), backlit (BAID: $+4.48$ dB), and fundus (BIQI $6.03$ vs. $6.14$) benchmarks.
• Inference Efficiency: $3$–$5$ rectified-flow steps ($\sim160$ ms on H100) suffice, outperforming standard diffusion-based student distillation methods in both computation and quality (Verma et al., 27 Sep 2025).

These results confirm that the dual principles of cross-normalized FLEX masking and generative rectified-flow feature synthesis jointly enable outlier-tolerant, high-fidelity transfer of transformer features to compact student architectures.

6. Significance and Implications

OFF, as realized in RestoRect, represents a significant advancement in feature knowledge distillation for architectures susceptible to unstable or outlier activations, notably transformers. By replacing static alignment with a trajectory-guided, stochastic ODE process and rigorously masking outliers, OFF achieves stable and efficient student training without loss of restoration fidelity. A plausible implication is broader applicability to additional domains requiring cross-architecture distillation under real-world outlier conditions and nonstationary activation statistics (Verma et al., 27 Sep 2025).