Deep Learning · Data Imbalance in Deep LearningChapter 09 / 12

Two-Stage Decoupling: Representation Learning vs. Classifier Calibration (cRT and $\tau$-Normalization)

October 2026 · 6 min read

Separate backbone feature learning from linear classifier calibration, see why head classifier weights grow larger $\ell_2$ norms $\|w_k\|_2$, and compare cRT, NCM, LWS, and $\tau$-normalization.

When you train a deep network end-to-end with class-balanced oversampling ($q=0$) or heavy tail loss weights ($\alpha_k=1/n_k$), are you improving the backbone representation $f_\theta(x)$ or just fixing the final linear classifier $W$? Kang et al. (ICLR 2020) tested both parts separately and uncovered a surprising result: instance-balanced training ($q=1$) learns the highest-quality feature extractor $f_\theta(x)$, and almost all of the long-tailed accuracy drop sits inside the final linear layer $z_k=w_k^\top f_\theta(x)+b_k$.

1. Why classifier weight norms $\|w_k\|_2$ grow with class count $n_k$

Write the logit for class $k$ in terms of the classifier vector norm $\|w_k\|_2$, the feature norm $\|f_\theta(x)\|_2$, and the angle $\theta_k$ between them:

$$z_k=w_k^\top f_\theta(x)+b_k=\|w_k\|_2\,\|f_\theta(x)\|_2\cos\theta_k+b_k.$$

Now inspect the gradient of the cross-entropy loss with respect to class $k$’s classifier weight vector $w_k$ over a mini-batch $\mathcal B$:

$$\frac{\partial L}{\partial w_k}=\frac{1}{|\mathcal B|}\sum_{i\in\mathcal B}\bigl(p_{i,k}-\mathbf{1}[y_i=k]\bigr)f_\theta(x_i)=\underbrace{-\frac{1}{|\mathcal B|}\sum_{y_i=k}(1-p_{i,k})f_\theta(x_i)}_{\text{positive pull along class }k\text{ features}}+\underbrace{\frac{1}{|\mathcal B|}\sum_{y_i\ne k}p_{i,k}f_\theta(x_i)}_{\text{negative push from all other }K-1\text{ classes}}.$$

Head class ($n_1\gg 1$): Appears in almost every mini-batch, receiving frequent positive pulls along its feature cluster that grow $\|w_1\|_2$.
Tail class ($n_K\ll 1$): Rarely appears as $y_i=K$, so in $99\%$ of mini-batches $w_K$ only receives negative pushes $+p_{i,K}f_\theta(x_i)$ and weight decay—shrinking $\|w_K\|_2$ to a fraction of $\|w_1\|_2$.

Because $z_k$ multiplies $\cos\theta_k$ by $\|w_k\|_2$, a head class with $\|w_1\|_2=4.0$ beats a tail class with $\|w_2\|_2=1.0$ even when the test feature $f_\theta(x)$ points much closer to the tail direction ($\cos\theta_2=0.80$ vs. $\cos\theta_1=0.25$):

$$z_1=4.0\times 0.25=1.00\; > \;z_2=1.0\times 0.80=0.80\quad(\text{Head wins despite }\cos\theta_2\gg\cos\theta_1!).$$

Left: classifier norm skew ||w_1||=4.0 vs ||w_2||=1.0 before and after tau-normalization. Right: two-stage decoupled training pipeline freezing f_theta(x) after Stage 1 and recalibrating the classifier head in Stage 2. — Left: head classifier norms $\|w_k\|_2$ grow much larger than tail norms during standard training; $\tau$-normalization rescales $\tilde w_k=w_k/\|w_k\|_2^\tau$. Right: Kang et al.’s two-stage decoupling trains the backbone at $q=1$ and then calibrates only the classifier head.

2. Four ways to calibrate the classifier head with $f_\theta(x)$ frozen

Once Stage 1 finishes training the backbone $f_\theta(x)$ with standard instance-balanced cross-entropy ($q=1$), freeze $\theta$ and apply one of four Stage 2 head calibrations:

$$\begin{aligned}\textbf{1. }\tau\textbf{-Normalization (zero extra training):}\quad &\tilde w_k=\frac{w_k}{\|w_k\|_2^\tau},\qquad \tilde b_k=0,\qquad \hat y=\arg\max_k\,\tilde w_k^\top f_\theta(x),\\[6pt]\textbf{2. Learnable Weight Scaling (LWS):}\quad &\tilde w_k=c_k\cdot w_k\quad(\text{learn scalar }c_k\text{ per class with }q=0\text{ sampler, }w_k\text{ frozen}),\\[6pt]\textbf{3. Classifier Re-Training (cRT):}\quad &\text{Re-initialize }W,b\text{ and train only the linear layer for }10\text{ epochs with }q=0,\\[6pt]\textbf{4. Nearest Class Mean (NCM):}\quad &\mu_k=\operatorname{normalize}\!\left(\frac{1}{n_k}\sum_{y_i=k}\frac{f_\theta(x_i)}{\|f_\theta(x_i)\|_2}\right),\qquad \hat y=\arg\max_k\,\mu_k^\top\frac{f_\theta(x)}{\|f_\theta(x)\|_2}.\end{aligned}$$

Check how $\tau\in[0,1]$ in $\tau$-normalization fixes the worked example above ($\|w_1\|_2=4.0,\cos\theta_1=0.25$ vs. $\|w_2\|_2=1.0,\cos\theta_2=0.80$) at $\tau=0.5$:

$$\begin{aligned}\|\tilde w_1\|_2&=\frac{\|w_1\|_2}{\|w_1\|_2^{0.5}}=\sqrt{4.0}=2.0,\qquad \|\tilde w_2\|_2=\frac{\|w_2\|_2}{\|w_2\|_2^{0.5}}=\sqrt{1.0}=1.0,\\[4pt]\tilde z_1&=2.0\times 0.25=0.50\; < \;\tilde z_2=1.0\times 0.80=0.80\quad(\text{Tail class now wins as it should!}).\end{aligned}$$

3. Questions worth answering

What do $\tau=0$ and $\tau=1$ correspond to in $\tau$-normalization?

Why does Nearest Class Mean (NCM) work well on Stage-1 features without any Stage-2 gradient steps?

Because each class centroid $\mu_k$ is an average over class $k$’s normalized features and is normalized to unit length $\|\mu_k\|_2=1$, making its norm completely independent of whether $n_k=10{,}000$ or $n_k=10$.