HeteroAKD Bridges CNN and Transformer Segmentation Models

Analysis by the aitrendblend editorial team · Pillar: Knowledge distillation and model compression · Source paper published 2025

The problem with borrowing a teacher from a different family

Knowledge distillation for segmentation usually assumes the teacher and the student came from the same architectural gene pool. Train a big ResNet, shrink it into a small ResNet, transfer the soft labels or the intermediate features, done. That assumption made sense while most segmenters were convolutional. It stopped making sense once transformer based segmenters such as SegFormer became genuinely competitive, because now the best available teacher for your compact model might not share a single design principle with it.

The authors point to a concrete number to justify the effort. Transferring knowledge from a DeepLabV3 model built on ResNet101 to a SegFormer style Mix Transformer B1 student produced a 3.66 point mIoU jump on ADE20K under their method. That single heterogeneous pairing beat the improvement that a prior state of the art method, Af DCD, achieved when distilling into a same family CNN student. In other words, picking a teacher from outside your architecture family is not a compromise. Done right, it can be the stronger move.

The catch is that CNN and transformer models look at data through different inductive biases. Convolutions build up local patterns layer by layer. Self attention can relate distant pixels immediately. Forcing a student to copy a teacher’s raw intermediate activations when the two networks process information so differently is a bit like asking someone to copy another person’s handwriting by tracing their thoughts instead of their pen strokes. It rarely works cleanly, and the paper backs that intuition with data rather than hand waving.

What centered kernel alignment reveals about the mismatch

To measure the gap rather than assume it, the researchers ran a centered kernel alignment analysis, comparing feature similarity between ResNet101 and Mix Transformer B4 layer by layer on five hundred Cityscapes samples. Homogeneous pairs of networks tend to learn similar representations at layers sitting in similar relative positions, which is exactly why classic feature imitation distillation works reasonably well between two CNNs. Heterogeneous pairs behave differently. CNN and transformer features only line up meaningfully in the shallow layers, then drift apart as depth increases.

That single result explains why so many existing feature based distillation methods stumble on cross architecture pairs. They were built around an assumption, layer position implies representational similarity, that simply does not hold once the backbone families diverge. Any fix that just projects student features into the teacher’s dimensionality and calls it aligned is treating a symptom, not the underlying cause.

The teacher is not always the smarter one in the room

The second finding is arguably the more interesting one, and it upends a habit baked into almost every distillation paper, the assumption that the teacher’s prediction is always the better target. The authors compared per class IoU scores between a transformer teacher and a CNN student, and between a CNN teacher and a transformer student, on real Cityscapes samples. For classes like truck and bus, the supposedly weaker student sometimes beat its own teacher by several points.

That makes intuitive sense once you sit with it. A convolutional network’s local receptive fields might carve out compact vehicle boundaries more cleanly, while a transformer’s global attention might do better with classes that depend on long range context, such as road or sky. Neither architecture is universally superior, each one is simply better tuned to certain visual patterns. A distillation method that blindly makes the student imitate every teacher prediction, pixel by pixel, will happily teach the student to unlearn something it already had right.

How HeteroAKD actually works

HeteroAKD answers both problems with a small pipeline rather than one clever trick. Instead of matching raw features, it maps both networks into a shared prediction space. Instead of blindly copying the teacher, it builds a custom target for every single pixel, informed by the ground truth label itself.

Step one, a shared logits space

Rather than fighting the CNN and transformer feature mismatch directly, the method sidesteps it. Both the teacher’s intermediate feature map and the student’s intermediate feature map are pushed through a lightweight projector, a single one by one convolution followed by batch normalization and a ReLU, that turns each one into a categorical logit map with the same shape as the final segmentation output.

Because both maps now live in the same class probability space, comparing them stops being an apples to oranges problem. The projector is thrown away at inference time, so it adds zero cost to the deployed student model. This single design choice is what lets the framework work symmetrically, meaning a CNN can teach a transformer or a transformer can teach a CNN using the exact same machinery.

Step two, deciding whose knowledge to trust

This is the part that answers the truck and bus problem from earlier. For every pixel, the framework measures how close each model’s prediction sits to the actual ground truth label using a cross entropy score.

A lower score means the model’s prediction is closer to the label, which the authors treat as a proxy for reliability. That score is computed separately for the teacher and the student, then turned into a weighting factor that literally decides how much of the final target comes from each side.

The student’s own weight is simply one minus that value. Whichever model is closer to the truth at that pixel earns more influence over the target the student is asked to chase.

The team calls this the knowledge mixing mechanism, and it produces a hybrid target that is neither pure teacher nor pure student, it is whichever one earned the vote of confidence at that particular pixel. One practical wrinkle worth flagging, feeding a freshly initialized student’s raw predictions into this mix would be pointless, so the student is fully trained under label supervision first, then distillation begins from that warmed up checkpoint.

Step three, focusing effort where it counts

Not every pixel deserves equal attention during distillation. A pixel the student already nails does not need more pressure, while a pixel the student is still missing, and that the hybrid target actually gets right, is exactly where learning should concentrate. The knowledge evaluation mechanism formalizes that intuition.

The indicator function only fires when the student’s error at that pixel is worse than the hybrid target’s error, meaning there is genuine room to improve. That importance score then feeds a softmax style reweighting across classes.

Those weights scale the actual distillation loss the student optimizes against the hybrid target.

Putting the three pieces into one objective

The final training signal combines the ordinary segmentation cross entropy loss, a standard logits based KD loss borrowed from the original Hinton, Vinyals, and Dean formulation, and the new heterogeneous loss described above.

Both projection heads used to align dimensions are discarded once training finishes, and the features pulled from the very last layer of each backbone are what feed the heterogeneous loss. Nothing about the deployed student’s architecture or runtime cost changes because of any of this.

How it performed across three benchmarks

The team tested HeteroAKD against five established segmentation distillation baselines, SKD, IFVD, CWD, CIRKD, and Af DCD, reimplemented on the same two codebases those methods originally shipped with, so the comparison is not stacked by mismatched training tricks.

Cityscapes results

On Cityscapes, using four backbone families and three segmentation heads, HeteroAKD beat the untrained baseline student in every single pairing, with the biggest single jump reaching 3.37 points of mIoU and an average gain across all pairings of 2.04 points. A few standout numbers from the paper’s main table follow.

Worth pausing on one detail buried in the table. In the SegFormer MiT B4 to DeepLabV3 ResNet18 pairing, the trained student ends up 0.46 points ahead of its own teacher. A compressed CNN student, taught by a much larger transformer, walks away better than the transformer that taught it. That result only makes sense once you accept the paper’s earlier point, the two architectures know different things, and mixing that knowledge intelligently can beat either one alone.

Pascal VOC and ADE20K results

To check the method was not tuned to one dataset’s quirks, the authors repeated the comparison on Pascal VOC and the far more class dense ADE20K benchmark. HeteroAKD again finished on top in every heterogeneous pairing tested, with a maximum mIoU margin of 2.10 points and an average margin of 0.93 points over the next best competing method.

The ADE20K jump on the SegFormer to DeepLabV3 pairing, 3.66 points, is the same figure the introduction leans on to argue that heterogeneous distillation deserves attention in its own right.

What the ablation study actually proves

It is easy for a paper to bundle several new components together and claim the whole package works without showing which piece is pulling its weight. The authors do the opposite here, stripping the framework down piece by piece on Cityscapes.

Two things stand out. First, neither the ordinary logits loss nor the new heterogeneous loss alone gets close to what the combination achieves, so the paper’s claim that intermediate feature learning and output level learning are complementary holds up under the numbers. Second, pulling out either the mixing mechanism or the evaluation mechanism costs real performance, an average of roughly 0.35 points and 0.49 points respectively across the two directions. Neither piece is decorative.

Where the method still has room to grow

The authors also ran a same architecture control, comparing HeteroAKD against Af DCD on strictly homogeneous teacher student pairs, CNN to CNN and transformer to transformer. HeteroAKD still won both, gaining an average of 1.61 points over the baseline. But one comparison in that table deserves attention. Distilling a DeepLabV3 ResNet18 student from a same family DeepLabV3 ResNet101 teacher produced a bigger jump, 2.51 points, than distilling the same student from the heterogeneous SegFormer MiT B4 teacher, 1.89 points.

The paper states this plainly in its own conclusion rather than burying it, noting that in certain cases the efficiency of knowledge distillation from a heterogeneous teacher can still trail what a homogeneous teacher delivers. That is a useful piece of honesty for anyone planning to adopt this approach. Heterogeneous distillation is a genuinely useful new tool, not a strictly superior replacement for matching architecture families when a strong same family teacher happens to already be available.

Complete PyTorch implementation

The snippet below is an independent, runnable reimplementation of the core HeteroAKD loss stack, the feature projector, the knowledge mixing mechanism, the knowledge evaluation mechanism, and the combined training objective, built directly from the equations in the paper. It is written to be readable rather than production tuned, and it ends with a smoke test on random tensors so you can confirm it runs before wiring it into a real segmentation pipeline.

# heteroakd.py
# Independent PyTorch reimplementation of the HeteroAKD loss components
# Based on the equations in Huang, Hu, Zhang, Chen, and Gao, arXiv:2504.07691

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


class LogitsProjector(nn.Module):
    """Projects a backbone feature map into a per pixel class logit map.

    This matches Equation 4 in the paper, a single one by one
    convolution followed by batch normalization and a ReLU, used to
    push both teacher and student features into a shared aligned
    logits space before any distillation loss is computed.
    """

    def __init__(self, in_channels, num_classes):
        super().__init__()
        self.proj = nn.Sequential(
            nn.Conv2d(in_channels, num_classes, kernel_size=1),
            nn.BatchNorm2d(num_classes),
            nn.ReLU(inplace=True),
        )

    def forward(self, feature_map):
        return self.proj(feature_map)


def pixel_reliability(logits, one_hot_labels, eps=1e-6):
    """Computes Equation 5, the pixel level knowledge reliability score.

    A lower value means the prediction is closer to the ground truth
    label at that pixel and channel, so it is treated as more
    trustworthy knowledge.
    """
    probs = torch.sigmoid(logits).clamp(eps, 1.0 - eps)
    ce = -(
        one_hot_labels * torch.log(probs)
        + (1.0 - one_hot_labels) * torch.log(1.0 - probs)
    )
    return ce


def knowledge_mixing(teacher_logits, student_logits, one_hot_labels, eps=1e-6):
    """Equations 5 through 7, the teacher student knowledge mixing mechanism.

    Returns the hybrid target Z_hat, plus the teacher weight map so it
    can be inspected or logged during training.
    """
    h_teacher = pixel_reliability(teacher_logits, one_hot_labels, eps)
    h_student = pixel_reliability(student_logits, one_hot_labels, eps)

    teacher_weight = 1.0 - (h_teacher / (h_teacher + h_student + eps))
    hybrid_logits = teacher_weight * teacher_logits + (1.0 - teacher_weight) * student_logits
    return hybrid_logits, teacher_weight


def knowledge_evaluation_weights(student_logits, hybrid_logits, one_hot_labels, eps=1e-6):
    """Equations 8 and 9, the teacher student knowledge evaluation mechanism.

    Produces a per class weight map that emphasizes pixels where the
    student is still behind the hybrid target and there is genuine
    room for the distillation signal to help.
    """
    h_student = pixel_reliability(student_logits, one_hot_labels, eps)
    h_hybrid = pixel_reliability(hybrid_logits, one_hot_labels, eps)

    gap = h_student - h_hybrid
    needs_help = (gap > 0).float()
    delta_h = needs_help * gap

    scores = h_student + delta_h
    weights = F.softmax(scores, dim=1)
    return weights


def hakd_loss(student_logits, hybrid_logits, weights, temperature=1.0):
    """Equation 10, the heterogeneous architecture knowledge distillation loss."""
    teacher_soft = torch.sigmoid(hybrid_logits / temperature)
    student_log_soft = torch.log(torch.sigmoid(student_logits / temperature).clamp(1e-6, 1.0))
    per_pixel = -(teacher_soft * student_log_soft) * weights
    return per_pixel.mean()


def kd_loss(student_logits, teacher_logits, temperature=1.0):
    """Equation 1, the standard logits based KD loss using KL divergence."""
    student_log_soft = F.log_softmax(student_logits / temperature, dim=1)
    teacher_soft = F.softmax(teacher_logits / temperature, dim=1)
    return F.kl_div(student_log_soft, teacher_soft, reduction="batchmean")


class HeteroAKDLoss(nn.Module):
    """Equation 11, the full training objective.

    Combines the task cross entropy loss, the standard KD loss, and
    the new heterogeneous architecture distillation loss.
    """

    def __init__(self, num_classes, lambda_kd=0.1, lambda_hakd=1.0, temperature=1.0):
        super().__init__()
        self.num_classes = num_classes
        self.lambda_kd = lambda_kd
        self.lambda_hakd = lambda_hakd
        self.temperature = temperature
        self.task_loss = nn.CrossEntropyLoss()

    def forward(self, student_seg_logits, labels, student_proj_logits, teacher_proj_logits):
        task = self.task_loss(student_seg_logits, labels)

        kd = kd_loss(student_seg_logits, teacher_proj_logits, self.temperature)

        one_hot = F.one_hot(labels.clamp(0, self.num_classes - 1), self.num_classes)
        one_hot = one_hot.permute(0, 3, 1, 2).float()

        hybrid_logits, teacher_weight = knowledge_mixing(
            teacher_proj_logits, student_proj_logits, one_hot
        )
        weights = knowledge_evaluation_weights(student_proj_logits, hybrid_logits, one_hot)
        hakd = hakd_loss(student_proj_logits, hybrid_logits, weights, self.temperature)

        total = task + self.lambda_kd * kd + self.lambda_hakd * hakd
        return total, {
            "task": task.item(),
            "kd": kd.item(),
            "hakd": hakd.item(),
            "teacher_weight_mean": teacher_weight.mean().item(),
        }


def training_step(student_backbone, student_head, student_projector,
                  teacher_backbone, teacher_projector, criterion,
                  optimizer, images, labels):
    """One optimization step. The teacher stays frozen throughout."""
    student_backbone.train()
    student_head.train()
    teacher_backbone.eval()

    with torch.no_grad():
        teacher_features = teacher_backbone(images)
        teacher_proj_logits = teacher_projector(teacher_features)

    student_features = student_backbone(images)
    student_seg_logits = student_head(student_features)
    student_proj_logits = student_projector(student_features)

    student_seg_logits = F.interpolate(
        student_seg_logits, size=labels.shape[-2:], mode="bilinear", align_corners=False
    )
    teacher_proj_logits = F.interpolate(
        teacher_proj_logits, size=labels.shape[-2:], mode="bilinear", align_corners=False
    )
    student_proj_logits = F.interpolate(
        student_proj_logits, size=labels.shape[-2:], mode="bilinear", align_corners=False
    )

    loss, logs = criterion(student_seg_logits, labels, student_proj_logits, teacher_proj_logits)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    return loss.item(), logs


@torch.no_grad()
def evaluate_miou(student_backbone, student_head, images, labels, num_classes):
    """A minimal mean intersection over union evaluator for sanity checks."""
    student_backbone.eval()
    student_head.eval()

    features = student_backbone(images)
    logits = student_head(features)
    logits = F.interpolate(logits, size=labels.shape[-2:], mode="bilinear", align_corners=False)
    preds = logits.argmax(dim=1)

    ious = []
    for cls in range(num_classes):
        pred_mask = preds == cls
        label_mask = labels == cls
        intersection = (pred_mask & label_mask).sum().item()
        union = (pred_mask | label_mask).sum().item()
        if union > 0:
            ious.append(intersection / union)
    return sum(ious) / max(len(ious), 1)


class TinyBackbone(nn.Module):
    """A stand in backbone, replace with a real ResNet or MiT encoder."""

    def __init__(self, out_channels=64):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(3, 32, 3, padding=1), nn.ReLU(inplace=True),
            nn.Conv2d(32, out_channels, 3, padding=1), nn.ReLU(inplace=True),
        )

    def forward(self, x):
        return self.net(x)


class TinyHead(nn.Module):
    def __init__(self, in_channels, num_classes):
        super().__init__()
        self.head = nn.Conv2d(in_channels, num_classes, kernel_size=1)

    def forward(self, x):
        return self.head(x)


if __name__ == "__main__":
    # Smoke test on random dummy data, confirms the full pipeline runs
    # end to end before you connect real backbones and a real dataloader.
    torch.manual_seed(0)
    num_classes = 5
    batch, height, width = 2, 64, 64

    images = torch.randn(batch, 3, height, width)
    labels = torch.randint(0, num_classes, (batch, height, width))

    student_backbone = TinyBackbone(out_channels=32)
    student_head = TinyHead(32, num_classes)
    student_projector = LogitsProjector(32, num_classes)

    teacher_backbone = TinyBackbone(out_channels=48)
    teacher_projector = LogitsProjector(48, num_classes)

    criterion = HeteroAKDLoss(num_classes=num_classes, lambda_kd=0.1, lambda_hakd=1.0, temperature=1.0)
    optimizer = torch.optim.SGD(
        list(student_backbone.parameters())
        + list(student_head.parameters())
        + list(student_projector.parameters()),
        lr=0.01,
    )

    loss_value, logs = training_step(
        student_backbone, student_head, student_projector,
        teacher_backbone, teacher_projector, criterion,
        optimizer, images, labels,
    )
    miou = evaluate_miou(student_backbone, student_head, images, labels, num_classes)

    print("loss", loss_value)
    print("component logs", logs)
    print("sanity mIoU on random data", miou)

Run that file directly and it should print a loss value, the breakdown across the task, KD, and HAKD components, plus a sanity mIoU on random data that will look meaningless because the inputs are random noise, which is exactly what a smoke test is supposed to confirm before you point it at a real dataset.

Honest limitations

Conclusion

HeteroAKD’s real contribution is not a new architecture or a bigger model. It is a working answer to a question that most segmentation distillation papers had been quietly avoiding, what do you do when the best available teacher was not built the same way as your student. By pushing both networks into a shared logits space and then letting ground truth labels arbitrate whose prediction to trust at every single pixel, the method turns architectural mismatch from a liability into something closer to a second, independent set of eyes on the same image.

The conceptual shift matters more than any single benchmark number. Most of the field has treated the teacher as an oracle, something the student should imitate as faithfully as possible. This paper treats the teacher as one fallible source of information among two, weighed against the same ground truth the student is ultimately trying to match. That framing, a trust score computed against labels rather than blind imitation, is portable well beyond semantic segmentation.

Whether that portability holds is worth watching closely. The core mechanism, project into a comparable output space, score each source’s reliability against ground truth, mix and reweight accordingly, does not obviously depend on segmentation specific machinery. Classification, depth estimation, and object detection all share the same basic distillation setup, a teacher, a student, and a loss that ties them together, so a version of this idea applied to those tasks would not be a stretch.

The limitations the authors flag themselves deserve equal weight to the wins. A homogeneous teacher sometimes still wins the efficiency race, hyperparameters do not transfer cleanly across backbone pairings, and every result here depends on having labeled data available during distillation, which rules out purely self supervised pipelines. None of that undercuts the core idea, it just means HeteroAKD is a strong new tool for a specific and increasingly common situation, not a universal replacement for choosing a matched teacher when one happens to be available.

As segmentation backbones keep diversifying, ConvNeXt variants, hybrid CNN transformer designs, the occasional state space model, the practical question of which teacher to pick for a compact deployed student is only going to come up more often. A framework that treats that mismatch as a feature to exploit rather than a problem to route around is worth keeping on the shortlist.

Frequently asked questions

Related reading