Rein++ Fine-Tunes Billion-Parameter Vision Models Cheaply

The two walls that stop foundation models at segmentation

Vision foundation models like CLIP, DINOv2, and InternVL learned general visual representations from enormous unlabeled or weakly labeled collections, and they transfer beautifully to many tasks. Semantic segmentation, the job of labeling every pixel, should be an easy customer. It is not, and the paper from Zhixiang Wei and colleagues at the University of Science and Technology of China names two walls.

The first is a scale mismatch. A model like DINOv2 was pre trained on the web scale LVD-142M collection, while a segmentation set such as Cityscapes or GTAV is orders of magnitude smaller. When you take a backbone carrying hundreds of millions to billions of parameters and fully fine tune it on that small set, you overwrite its broad knowledge and it generalizes worse, not better. The model memorizes the little dataset and loses the wide view.

The second wall is distribution shift. Real deployment is messy. A segmentation model trained on clean daytime streets meets rain, fog, night, and snow, none of which were well represented in pre training. Labeling data for every one of those conditions is expensive and slow. So the practical need is to adapt to new visual domains without labels, which is exactly where models trained on one tidy setting fall down.

Older work handled both walls with modest backbones like ResNet and small vision transformers, through domain generalization methods and unsupervised adaptation methods. What nobody had properly done was bring the billion parameter foundation models into this arena efficiently. That gap is the opening Rein++ walks through.

The baseline that should not work so well

Before proposing anything, the authors ran a test that sets the tone for the whole paper. Take a foundation model, freeze it completely, and train only a segmentation decoder on top. No clever adaptation, nothing touched inside the backbone. This frozen backbone approach ought to be a weak baseline. Instead it is formidable.

A frozen DINOv2-Giant with a learnable decoder generalizes to unseen domains well enough to beat the previous state of the art method, HRDA, by roughly 10 points of mIoU. Read that again. A model that learned nothing about segmentation datasets during pre training, and whose backbone was never adjusted for the task, outperforms a specialized method that was trained end to end. The representations from web scale pre training already carry generalization that trained specialists on small backbones cannot match.

That result reframes the problem. The question is no longer how to teach a backbone to segment. It is how to gently steer a backbone that already knows a great deal, without disturbing what makes it strong. Freezing is too blunt, because a model pre trained on generic proxy tasks still lacks the fine grained, task specific granularity that crowded urban scenes demand. The goal becomes a light touch that adds segmentation skill while leaving the general knowledge intact.

Rein-G, steering with learnable tokens

The generalization component, named Rein-G, is where the light touch lives. Rather than modifying the backbone’s weights, it inserts a small module between the frozen layers that refines the feature map each layer produces before passing it on. The backbone stays frozen. Only the module trains. The name is a deliberate pun, since rein is what you use to steer a horse without rebuilding it.

Concretely, for the features a layer produces, Rein-G computes a refinement and adds it back before the next layer sees them. The features flow through the frozen stack, nudged at every step.

The interesting part is how the refinement is computed. Rein-G holds a set of learnable tokens, and the paper’s central idea is that these tokens attach to instances, meaning distinct objects in the scene. For each layer it measures how much each image patch resembles each token with a dot product, which produces a similarity map that behaves like attention.

Because each token connects to a different instance, the refinement can adjust the features of one object differently from another within the same image. That per instance flexibility is what lets a generic backbone start behaving like a segmentation model that cares about individual things. The refinement is then read off from the similarity map and a set of token derived values, split across several heads so the module can capture complementary cues in parallel, and passed through a small MLP with a GELU nonlinearity for extra expressiveness.

Two efficiency tricks keep this cheap. The tokens are built as a product of two thin matrices, a low rank form, so a long token sequence costs few parameters.

And the MLP weights are shared across all layers, so adding depth does not multiply the parameter count. The result is a module that fine tunes less than 1 percent of the backbone yet beats full fine tuning. The paper reports the full backbone at about 304 million trainable parameters against Rein-G at about 2.5 million, and the smaller footprint wins.

The counterintuitive lesson is that on a small dataset, training fewer parameters is not a compromise for the sake of memory. It is often the thing that makes the model generalize better. A plain reading of the Rein++ analysis

Why fewer parameters generalize better

This is worth dwelling on, because it inverts a common instinct. The authors trace the training loss and the test performance together as trainable parameters grow from zero for a frozen model, to 2.5 million for Rein-G, to 304 million for full tuning. Training loss falls the whole way, as more parameters always fit the training data better. But test performance on unseen domains peaks at Rein-G and then declines. Full tuning records the lowest training loss and mediocre generalization, which is the textbook signature of overfitting.

In ordinary tasks you would catch this with early stopping on a validation set. Domain generalization removes that safety net, because the whole point is that the test distribution is unknown, so you cannot build a valid checkpoint for it. Combine an unknown target with a small fine tuning set and the pressure to overfit is strong. Keeping the trainable parameter count tiny is a structural defense against it. Rein-G generalizes better partly because it fits the backbone well and partly because it simply cannot overfit the small dataset the way a fully trained backbone can.

Rein-A, adapting without any labels

Strong generalization still leaves value on the table, because in practice you often have piles of unlabeled images from the exact target domain you care about. Rein-A is the adaptation half that puts those to use. It keeps the backbone and Rein-G doing the heavy lifting and adapts the model to unlabeled target images through a teacher and student self training loop, where a slowly averaged teacher produces pseudo labels that supervise the student.

The design choices here are mostly about stability, and they are more subtle than they first appear. Modern foundation model segmentation heads use mask classification, in the style of Mask2Former, rather than classifying each pixel independently. That decoder normally relies on bipartite matching to assign predicted masks to ground truth objects. Under adaptation the ground truth is replaced by noisy pseudo labels, and the matching becomes unstable, causing what the authors call query drift and class collapse early in training. A naive port of an existing per pixel adaptation method into this mask setting simply fails to converge.

Rein-A works around this with a nice observation. When the class count is small and the teacher updates slowly, the student’s queries stay aligned with the teacher’s queries by index throughout training. So it can drop the fragile bipartite matching entirely and supervise each student query directly from the teacher query of the same index. That single change removes the main source of instability.

On top of that it adds two complementary losses, one at the semantic logit level using a class mixing procedure across domains, and one at the instance mask level using masked target views. And it borrows segmentation priors from the Segment Anything Model through a semantic transfer module, which lends the class agnostic but crisp mask boundaries of SAM to sharpen the target predictions. The combination is what turns a generalized model into an adapted one without ever touching a target label.

What the numbers show

The headline setting is synthetic to real, training on the game rendered GTAV data and testing on real streets. Here the generalization model reaches 71.9 percent mIoU on Cityscapes, and turning on adaptation with unlabeled target images lifts that to 78.2 percent. Getting a 6 point boost from images you never labeled is the practical payoff of the adaptation half.

A few of these deserve comment. Paired with a Radio-Giant backbone, Rein-G averages 68.8 percent mIoU across three real datasets, which is about 23 points above ResNet101 based methods and 3.1 points above the foundation model driven TQDM. Against the strong adaptation baseline MIC, Rein-A is not a modest step. It gains 8.0 points on the hard Cityscapes to ACDC transfer and a striking 11.3 points on Cityscapes to Dark Zurich, exactly the night and adverse weather conditions where earlier methods struggled most. The advantage is largest on the semantically complex classes like person, rider, and the various vehicle types, while the easy classes like road and sky were already near saturation for everyone.

Breadth backs up the headline. The framework was tested across five foundation model families, namely CLIP, EVA02, DINOv2, Radio, and InternVL-C, and three scales up to 6 billion parameters. In every case Rein-G beat both the frozen and the fully tuned versions of the same backbone. On DINOv2-Giant it topped full tuning by 6.4, 3.9, and 2.4 points across three settings. That a single recipe helps every backbone it touches is a sign the idea is structural, not tuned to one model.

The efficiency that makes billions practical

Numbers on a leaderboard are only half the story. The other half is whether you can run this at all on the hardware most labs have. Here the frozen design pays off twice.

On memory, fully fine tuning a 6 billion parameter InternVL-C is simply infeasible on a standard 40 GB A100, failing even at a batch size of one because it wants 96 GB. Adaptation of such a model would normally demand over 80 GB. Rein++ brings that down to roughly a third, into the range of a single accessible GPU, because it only ever computes gradients for a tiny module. The paper reports Rein-A needing about 0.11 GB for its trainable parameters against 4.35 GB for a full DINOv2-Giant backbone.

On storage, the win compounds when you serve many domains. Adapting to a new scene means saving a small set of tokens, not a fresh copy of the whole backbone. You can keep one frozen giant on disk and a shelf of lightweight token sets, one per condition, rather than duplicating billions of weights for each. For anyone running segmentation across many environments, that is the difference between a feasible deployment and a storage nightmare.

The honest limits

None of this arrives without caveats, and the paper is fairly candid. The whole approach is downstream of the foundation model’s quality. Rein-G steers what the backbone already sees, so a domain the pre trained model represents poorly will still be hard, since there is only so much a thin module can add on top of weak features. The method rides the giant and cannot exceed its reach.

The adaptation half also carries the usual fragility of self training. It leans on teacher pseudo labels that can be noisy, and the authors themselves show that the mask classification decoder is delicate under that noise, which is why they had to engineer around bipartite matching. Their index based alignment depends on conditions that held in their experiments, a small class count and a slow teacher update, and it is reasonable to ask how it behaves when those do not hold. The ablations confirm the pieces are load bearing, since removing the mix branch costs 5.1 points and dropping Rein-G entirely costs more than 5 points, so the stability is earned rather than free.

There is also a scope note worth stating plainly. The evidence lives in driving and urban scene benchmarks, the GTAV, Cityscapes, ACDC, Dark Zurich, and BDD100K family. That is a well studied but narrow slice of the visual world. The recipe is general in principle, but its behavior on medical imaging, remote sensing, or other far domains is not shown here and should not be assumed.

Why the frozen giant idea travels

Step back and Rein++ is part of a broader shift in how the field treats large pre trained models. The old reflex was to fine tune everything for each new task. The emerging view, which parameter efficient methods like LoRA started in language models and Rein-G extends to dense vision, is that the pre trained weights are a fixed asset to be steered, not clay to be reshaped. Freeze the expensive knowledge, train a cheap controller, and you get better generalization and far lower cost at the same time.

The instance token idea has legs beyond street scenes. Any dense prediction task where a frozen backbone gives good features but lacks task specific sharpness could borrow the same pattern of light per instance refinement between frozen layers. And the adaptation lessons, especially the fragility of mask classification decoders under pseudo labels and the fix of index based query alignment, are useful to anyone building self training systems on modern segmentation heads. Those are the transferable parts a practitioner in an adjacent area should take away.

Conclusion

Rein++ makes a clean case that the right way to use a vision foundation model for a hard dense task is not to retrain it but to steer it. The paper’s own baseline sets up the argument, since a frozen giant already beats trained specialists on unseen domains, and everything after that is about adding segmentation skill without disturbing the general knowledge that made the giant strong.

The generalization half does this with learnable tokens that attach to instances and refine features between frozen layers, kept cheap by a low rank form and shared weights. The payoff is not only efficiency but better generalization, because a tiny trainable footprint cannot overfit a small dataset the way a fully tuned backbone does. That inversion, where fewer parameters raise test accuracy, is the conceptual heart of the work.

The adaptation half then turns unlabeled target images into real gains, closing the gap to night and adverse weather with a self training loop carefully redesigned for the delicate mask classification decoders that modern foundation models use. The results are broad rather than narrow, spanning five model families and scales up to 6 billion parameters, and the efficiency figures bring billion parameter adaptation onto a single accessible GPU.

The limits keep it grounded. Performance is capped by the backbone, self training is inherently noisy, and the evidence lives in driving scenes rather than the whole visual world. Those are the seams a follow up should work on. For now Rein++ offers a durable template. As foundation models keep growing past the point where anyone can afford to fine tune them fully, the ability to freeze the giant and train a sliver stops being a nice option and starts being the only practical one.

The listing below reconstructs the core of Rein-G, namely the multi head token to feature similarity, the low rank learnable tokens, the shared MLP with a GELU nonlinearity, and the residual refinement inserted between frozen backbone layers. It uses small dummy tensors so it runs anywhere and ends with a smoke test. Wrap a real frozen DINOv2 or CLIP backbone to reproduce the full method.

# rein_g.py
# Minimal, runnable reconstruction of the Rein-G refinement module.
# The backbone is a stand in so the file runs without a real VFM.

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


class LowRankTokens(nn.Module):
    """Learnable tokens built as A x B, a low rank form.

    m tokens of width c cost only r(m + c) parameters, with r far
    smaller than c, which keeps the module light.
    """

    def __init__(self, m, c, r=16):
        super().__init__()
        self.A = nn.Parameter(torch.randn(m, r) * 0.02)
        self.B = nn.Parameter(torch.randn(r, c) * 0.02)

    def forward(self):
        return self.A @ self.B                 # [m, c]


class ReinG(nn.Module):
    """One shared Rein-G refiner applied at every frozen layer.

    Computes a per instance residual for the patch features using a
    multi head dot product between patch features and learnable tokens,
    then a shared GELU MLP. Weights are shared across layers.
    """

    def __init__(self, c, m=100, r=16, heads=4):
        super().__init__()
        assert c % heads == 0
        self.c, self.h, self.dh = c, heads, c // heads
        self.tokens = LowRankTokens(m, c, r)
        self.to_value = nn.Linear(c, c)         # token values
        # shared GELU MLP, squeeze to c/2 then back to c
        self.mlp = nn.Sequential(
            nn.Linear(c, c // 2),
            nn.GELU(),
            nn.Linear(c // 2, c),
        )

    def _split(self, x, n):
        # [*, c] -> [heads, *, dh]
        return x.reshape(*x.shape[:-1], self.h, self.dh).movedim(-2, 0)

    def forward(self, f):
        """f is [n, c] patch features for one layer. Returns refined f."""
        n, c = f.shape
        T = self.tokens()                       # [m, c]
        V = self.to_value(T)                    # [m, c]

        fh = self._split(f, n)                  # [h, n, dh]
        Th = self._split(T, T.shape[0])         # [h, m, dh]
        Vh = self._split(V, V.shape[0])         # [h, m, dh]

        # dot product similarity, scaled, softmax over tokens
        sim = torch.matmul(fh, Th.transpose(-1, -2))    # [h, n, m]
        attn = F.softmax(sim / (self.dh ** 0.5), dim=-1)
        upd = torch.matmul(attn, Vh)            # [h, n, dh]

        # merge heads back to [n, c]
        upd = upd.movedim(0, -2).reshape(n, c)
        delta = self.mlp(upd + f)               # shared GELU MLP
        return f + delta                       # residual refinement


class FrozenBackboneWithRein(nn.Module):
    """Stand in frozen stack with one shared Rein-G between layers."""

    def __init__(self, c=64, depth=6, m=100, r=16, heads=4):
        super().__init__()
        # frozen layers, weights never updated
        self.layers = nn.ModuleList(nn.Linear(c, c) for _ in range(depth))
        for p in self.layers.parameters():
            p.requires_grad = False
        self.rein = ReinG(c, m, r, heads)       # the only trainable part

    def forward(self, f):
        for layer in self.layers:
            f = self.rein(f)                    # refine then forward
            f = layer(f)
        return f


def count_trainable(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


def smoke_test():
    """Runs on dummy data, no vision foundation model needed."""
    torch.manual_seed(0)
    n, c = 256, 64                              # 256 patches, width 64
    model = FrozenBackboneWithRein(c=c, depth=6)

    f = torch.randn(n, c)
    out = model(f)

    total = sum(p.numel() for p in model.parameters())
    train = count_trainable(model)
    assert out.shape == (n, c)
    assert torch.isfinite(out).all()
    print("output shape:", tuple(out.shape))
    print(f"trainable {train} of {total} params "
          f"({100 * train / total:.1f}% trained, rest frozen)")
    print("smoke test passed")


if __name__ == "__main__":
    smoke_test()

Frequently asked questions

You can also open the source article through its IEEE TPAMI listing, which links the official version and the supplementary material.