FiGKD Distills Only the Teacher’s High Frequency Logits

A teacher’s answer is not one thing

Knowledge distillation has a simple pitch. Train a big accurate model, then train a small cheap model to copy its answers rather than only the ground truth labels. The small model ends up smarter than it would from labels alone, because the big model’s answer carries extra information about which wrong classes are almost right. Geoffrey Hinton called this dark knowledge back in 2015, and the idea has held up well enough that it now sits inside a huge share of production model compression pipelines.

The trouble shows up on tasks where the classes look almost identical. A restaurant and a dining room. A Laysan albatross and a black footed albatross. A German shepherd and a Belgian malinois. These are the tasks researchers call fine grained recognition, and they are exactly where compact student models tend to fall apart even after distillation. The paper’s Figure 1 makes the point with a single comparison. A ResNet teacher tells a restaurant from a dining room without trouble. A student trained with plain knowledge distillation gets the two backward. The same student trained with FiGKD gets both right.

Why would that gap exist. Every method already discussed treats the teacher’s logit vector as one undivided signal. Every number in it, from the dominant predicted class down to the faintest hint of similarity toward an unrelated class, gets matched with equal weight through a KL divergence or a mean squared error term. Decoupled Knowledge Distillation, a strong 2022 method known as DKD, tried to fix this by splitting the logit vector into the target class score and the non target class scores and weighting each piece separately. That helps, and DKD remains a common baseline. But it still draws its line along class identity rather than along the kind of information a class score carries.

Splitting a prediction the way you would split an image

The frequency domain framing is not new to deep learning generally. Researchers have long known that convolutional networks favor low frequency structure in images and struggle relatively more with high frequency detail, and that robustness to noise tracks how well a network handles those high frequency signals. Wavelet based image translation methods already exploit this by supervising the high frequency bands of a generated image separately from its coarse structure, since faithfully reproducing sharp edges and fine texture is the harder part of the job. What is new here is applying that same decomposition not to an image, but to the one dimensional vector of class scores a model outputs.

The mechanism is a discrete wavelet transform, DWT, the same tool used for JPEG2000 style image compression. A DWT decomposes a signal into a coarse approximation and a set of detail coefficients that capture the parts a coarser view would smooth over. FiGKD reshapes a classifier’s logit vector, normally a flat list of one score per class, into a two dimensional matrix so a standard 2D DWT can be applied to it. For CIFAR-100’s hundred classes that becomes a ten by ten grid. For the two hundred classes of TinyImageNet or CUB-200 it becomes ten by twenty. A single pass of the simplest possible wavelet, the Haar wavelet, then yields one low frequency band and three high frequency bands, corresponding to horizontal, vertical, and diagonal detail.

That reshaping step deserves a second look, because a logit vector is not spatial data the way a photograph is. There is no inherent horizontal or vertical direction between class fifty two and class fifty three. The paper is upfront about this. It tested a version that keeps the logits as a flat one dimensional signal and runs a 1D wavelet transform instead, and found the 2D version consistently ahead by an average of roughly half a percentage point across eight different teacher and student pairings on CIFAR-100, table 13 in the paper. The explanation offered is that a 2D transform enforces three separate directional constraints instead of one, giving the student three somewhat independent views of the same detail rather than a single flattened one. It is a pragmatic justification more than a principled one, and it is worth remembering that the spatial layout of a reshaped logit vector is an engineering choice, not a property the classes actually have.

What actually gets copied to the student

Once the teacher’s logits and the student’s logits are both decomposed the same way, FiGKD computes a loss only between the high frequency bands.

The low frequency band, which the paper describes as carrying the dominant class scores, is dropped entirely rather than distilled. The stated reasoning is that ground truth labels already tell the student which class should win, so a coarse approximation of the teacher’s confident prediction is redundant with information the student already has. Only the finer structure among the losing classes, the pattern that reveals which wrong answers the teacher considers almost right, is treated as worth transferring.

That reasoning is intuitive, but the paper does not simply assert it and move on. It backs it up with a small piece of empirical detective work in Section 5.2. The authors ran the same DWT decomposition on the logits produced by a strong high accuracy model, ResNet32x4 at 79.42 percent on CIFAR-100, and a weaker one, ResNet8x4 at 72.50 percent, and checked whether the correct class could still be identified from the high frequency band alone. For the strong model, yes. The peak corresponding to the ground truth class stays visible even after the low frequency content is stripped out. For the weak model, the high frequency band gets diffuse, and the correct class is no longer obviously the winner. Compact student models, in other words, have not yet learned to encode precise class discriminative structure in this detailed part of their own output, which is exactly the gap FiGKD tries to close by handing over the teacher’s version of that structure directly.

“The high frequency component alone is often sufficient to predict the target class.” From the paper’s analysis of strong versus weak models, Section 1

The ablation result that argues against itself

Here is where the paper gets more interesting than a typical distillation write up, and where the abstract undersells its own finding. Table 10 runs a clean four way ablation on both CIFAR-100 and TinyImageNet across several teacher and student pairs. It compares four settings. Distill neither frequency band, meaning a plain baseline. Distill only the low frequency band. Distill only the high frequency band, which is the actual FiGKD method. Distill both bands together.

On the WRN-40-2 to WRN-16-2 pair on CIFAR-100, the plain baseline reaches 76.63 percent. Adding low frequency distillation on top of that baseline does not help. It lands at 76.44 percent, half a point below the version that distills nothing extra at all. High frequency distillation alone reaches 77.06 percent, the best result. Distilling both bands together lands at 76.56 percent, worse than high frequency alone and roughly tied with the plain baseline. The same pattern shows up on TinyImageNet’s ResNet32x4 to ResNet8x4 pair, where the baseline sits at 60.67 percent, low frequency only drops slightly to 60.45 percent, high frequency only climbs to 60.92 percent, and combining both settles at 60.71 percent, again below high frequency alone.

Read that carefully and it says something stronger than distilling high frequency information helps. It says that copying the teacher’s low frequency, dominant class content into the student can make the student worse than giving it no extra distillation signal at all, and that mixing a good signal with a bad one produces a result worse than the good signal by itself. That is not a subtle regularization effect. It is evidence that part of what conventional full logit distillation transfers is actively counterproductive once ground truth labels are already doing that job, and that a student’s limited capacity is genuinely a zero sum resource where wasted signal crowds out useful signal rather than simply doing nothing.

How it holds up across five benchmarks

Empirical claims about a training trick are only as good as how far they generalize, so the paper tests FiGKD broadly. On CIFAR-100 with matching teacher and student architecture families, such as ResNet32x4 teaching ResNet8x4, FiGKD beats the strongest logit based baseline it compares against by an average of 0.66 percentage points and edges out feature based methods that require access to intermediate layers, including CAT-KD and ReviewKD, which is notable since FiGKD only ever touches the final logits. Move to a mismatched setting, where a WRN-40-2 teacher trains a ShuffleNet-V1 student with a completely different architecture, and the gap widens to an average of 1.28 points, with the biggest single jump reaching 1.40 points over the strong MLKD baseline.

TinyImageNet, with twice the classes and more visual variety, shows the same shape of result, a 0.58 point average gain in matched settings and 1.54 points in mismatched settings. The fine grained datasets are where the method earns its keep most clearly. On CUB-200’s two hundred bird species and MIT67’s sixty seven indoor scene categories, both places where classes genuinely look alike, FiGKD beats the strongest baseline by 3.16 and 3.78 percentage points respectively in matched architecture settings, and by 2.35 and 3.08 points when the student architecture differs from the teacher’s.

That upward slope from coarse grained to fine grained tasks lines up neatly with the paper’s core argument. Coarse categories like truck or apple are usually separable using dominant class information alone, so there is less room for a detail focused method to add value. Categories that differ only in subtle features benefit far more from a training signal aimed specifically at the part of the teacher’s output that encodes those subtleties.

One wavelet pass, not several

A natural next question for a wavelet based method is whether going deeper, running the decomposition through two or three levels to extract finer and finer detail, would help further. The paper checked, and the answer is no. Table 11 compares decomposition levels one through three on four teacher and student pairs, and the single level version wins every time. Going from level one to level two on the WRN-40-2 to WRN-16-2 pair drops accuracy from 77.06 to 76.80 percent, and level three drops it further to 76.52 percent. The explanation given is fairly mundane. CIFAR-100’s reshaped logit grid is only ten by ten to begin with. A second decomposition level shrinks the spatial support to three by three, and a third to two by two, points where there simply is not much structure left to extract.

That restraint also pays off in training cost. Table 14 reports the added computational overhead of FiGKD’s default single level setting at roughly 2.0 percent per batch compared to plain knowledge distillation, an addition small enough to ignore in most training budgets. Pushing to two or three decomposition levels raises that overhead to 19.5 and 30.1 percent respectively, for a method that performs worse. It is a rare case in this kind of research where the cheapest setting and the most accurate setting are the same one, rather than a tradeoff the practitioner has to choose between.

The method also turned out to be fairly forgiving of noise. Injecting Gaussian noise directly into the high frequency logit subbands during training, with a standard deviation as large as 0.5, barely moved final accuracy, a drop of well under half a point across the pairs tested in Table 12. The authors read this as evidence that the high frequency band carries a genuinely structured, semantically meaningful signal rather than fragile numerical noise that a small perturbation could wipe out. It is a reasonable interpretation, though it is also possible that a training loss with a modest weight on this term is simply not very sensitive to noise in that term in the first place, and the paper does not fully separate those two explanations.

What this changes for people building compressed models

The practical case for FiGKD rests on three things that matter in real deployment work rather than only in a benchmark table. First, it never touches intermediate feature maps, so a team can distill from any teacher whose logits they can read, including a teacher whose architecture they did not build and cannot instrument, or one served behind an API where only final outputs are visible. Second, the added compute cost is close to nothing, roughly two percent per batch, so it slots into an existing distillation pipeline without materially changing training time. Third, the compression ratios tested are genuinely aggressive. The ResNet50 to MobileNetV2 pairing on TinyImageNet has the teacher carrying over 25 times the parameters and nearly 180 times the floating point operations of the student, and FiGKD still narrows the gap by a meaningful margin over other logit based methods in that setting.

None of that makes this a universal upgrade over every distillation method in every situation. Feature based methods like ReviewKD and CAT-KD still win outright in a handful of the paper’s own tables, particularly on some heterogeneous CIFAR-100 pairs, and they remain the better choice when a team already has full access to teacher internals and is not worried about the coupling that creates. FiGKD’s pitch is narrower and more specific. When you can only see the teacher’s final answer, and your task involves telling apart classes that genuinely look alike, discard the part of that answer your ground truth labels already give you for free, and put your training budget entirely behind the part that does not.

Where the evidence runs thin

A few limitations are worth naming plainly rather than glossing over. The paper is a single author work, and while it is now accepted at Expert Systems with Applications, the version reviewed here is the arXiv preprint, so some numbers could still shift before final publication. The hyperparameter sensitivity analysis in Figure 4 shows that the best balance between the cross entropy weight and the detail loss weight differs between CIFAR-100 and TinyImageNet, and the paper is honest that it did not run a full search on every dataset, which means the reported numbers likely understate what further tuning could squeeze out, in either direction. Every experiment here is image classification with convolutional backbones, VGG, ResNet, WideResNet, MobileNet, and ShuffleNet variants. There is no test against a vision transformer teacher or student, and no test outside image classification, so claims about generality to other architectures or other modalities remain speculative rather than demonstrated. Finally, the noise robustness test uses synthetic Gaussian perturbation, which is a reasonable first check but is not the same as robustness to the kinds of quantization error, adversarial perturbation, or distribution shift a deployed system would actually encounter.

The proposed model in PyTorch

Below is a complete, runnable implementation of the FiGKD loss and a small teacher and student pair, following the paper’s formulation. It uses a fixed Haar wavelet, a single decomposition level, and reshapes each logit vector into as close to a square grid as the number of classes allows, matching the paper’s approach for datasets like CIFAR-100.

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

# --- Reshape a flat logit vector into a near square H x W grid, as in Section 3.3 ---
def factor_grid(num_classes):
    h = int(math.floor(math.sqrt(num_classes)))
    while num_classes % h != 0:
        h -= 1
    w = num_classes // h
    return h, w

# --- A minimal 2D Haar DWT, single level, applied over the last two dims ---
class HaarDWT2D(nn.Module):
    def __init__(self):
        super().__init__()
        # Four Haar filters: LL, LH, HL, HH, each 2x2
        ll = torch.tensor([[0.5, 0.5], [0.5, 0.5]])
        lh = torch.tensor([[0.5, 0.5], [-0.5, -0.5]])
        hl = torch.tensor([[0.5, -0.5], [0.5, -0.5]])
        hh = torch.tensor([[0.5, -0.5], [-0.5, 0.5]])
        filters = torch.stack([ll, lh, hl, hh], dim=0).unsqueeze(1)
        self.register_buffer("filters", filters)

    def forward(self, x):
        # x has shape (batch, 1, H, W) with H and W both even
        out = F.conv2d(x, self.filters, stride=2)
        low = out[:, 0:1]
        high = out[:, 1:4]  # horizontal, vertical, diagonal subbands
        return low, high

class FiGKDLoss(nn.Module):
    def __init__(self, num_classes, alpha=2.0, beta=2.0):
        super().__init__()
        self.h, self.w = factor_grid(num_classes)
        assert self.h % 2 == 0 and self.w % 2 == 0, "pad classes so H and W are even"
        self.dwt = HaarDWT2D()
        self.alpha = alpha
        self.beta = beta
        self.ce = nn.CrossEntropyLoss()

    def forward(self, student_logits, teacher_logits, labels):
        b = student_logits.shape[0]
        s_map = student_logits.view(b, 1, self.h, self.w)
        t_map = teacher_logits.view(b, 1, self.h, self.w)

        _, s_high = self.dwt(s_map)
        with torch.no_grad():
            _, t_high = self.dwt(t_map)

        detail_loss = F.mse_loss(s_high, t_high, reduction="mean")
        ce_loss = self.ce(student_logits, labels)
        total = self.alpha * ce_loss + self.beta * detail_loss
        return total, ce_loss.item(), detail_loss.item()

# --- Smoke test on dummy data, standing in for CIFAR-100 shaped logits ---
if __name__ == "__main__":
    torch.manual_seed(0)
    num_classes = 100
    batch_size = 16

    # Tiny stand in teacher and student heads, enough to prove the loss trains
    teacher_head = nn.Linear(64, num_classes)
    student_head = nn.Linear(32, num_classes)
    optimizer = torch.optim.SGD(student_head.parameters(), lr=0.05, momentum=0.9)

    criterion = FiGKDLoss(num_classes=num_classes, alpha=2.0, beta=2.0)

    dummy_teacher_feats = torch.randn(batch_size, 64)
    dummy_student_feats = torch.randn(batch_size, 32)
    dummy_labels = torch.randint(0, num_classes, (batch_size,))

    losses = []
    for step in range(50):
        teacher_logits = teacher_head(dummy_teacher_feats).detach()
        student_logits = student_head(dummy_student_feats)

        loss, ce_val, detail_val = criterion(student_logits, teacher_logits, dummy_labels)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        losses.append(loss.item())

    assert losses[-1] < losses[0], "loss did not decrease, smoke test failed"
    print(f"Smoke test passed. Loss went from {losses[0]:.4f} to {losses[-1]:.4f}")

The important part to notice is the torch.no_grad() around the teacher’s wavelet pass. The teacher is frozen during distillation, so only the student’s high frequency subbands should receive gradients from the detail loss, exactly as the paper’s formulation implies by treating the teacher’s decomposition as a fixed target.

Conclusion

FiGKD’s core achievement is a small, almost surgical change to a training loss that produces a consistent and in some cases substantial improvement across a genuinely broad test bed, five datasets, ten backbone families, and both matched and mismatched teacher student pairings, without asking for anything more than the teacher’s final output. That combination of simplicity and reach is what makes the paper worth a careful read rather than a skim of its headline numbers.

The deeper contribution, though, is the reframing itself. Most of the last decade of distillation research has approached the teacher’s logit vector as something to split by class identity, target versus non target, or to reweight by temperature, but still fundamentally as a single blob of information to be matched more or less faithfully. Treating that vector as a signal with its own internal frequency structure, and finding that one part of that structure is actively harmful to copy while another part is close to essential, opens a different axis for future work to explore, one that sits orthogonal to almost everything that came before it.

Whether that framing generalizes past image classification is the open question the paper leaves for others to answer. A logit vector over object classes has some rough visual analogy to a small image, which is presumably why reshaping it into a 2D grid and running an off the shelf image wavelet transform on it works reasonably well. It is far less obvious that the same trick transfers cleanly to token level logits in a language model, to bounding box regression outputs in a detector, or to continuous action outputs in a control policy, all places where the notion of a low frequency dominant signal and a high frequency detail signal is much less clearly defined.

The honest limitations are real and worth keeping in view. A single author’s preprint, tuning that clearly still has headroom left on the table by the author’s own admission, and a test suite confined entirely to convolutional image classifiers. None of that erases the strength of the central result, but it does mean the method deserves independent replication before it gets treated as a settled default choice for distillation pipelines.

For anyone currently building a compressed model for an edge device or a real time system, the actionable version of this paper is short. If your student is struggling specifically on classes that look alike, try distilling only the high frequency band of your teacher’s logits and dropping the rest, and check whether that alone beats whatever full logit distillation setup you are running today. Given how little this costs to try, roughly a two percent training time tax and a wavelet transform anyone can implement in an afternoon, it is a cheap experiment with a real chance of a meaningful payoff.

Frequently asked questions

Related reading