Skip to content

v-code01/oftfold

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

6 Commits
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

oftfold

PEFT's OFT (orthogonal finetuning) adapter rotates a layer's weight by a learned orthogonal matrix. For a Conv2d the forward instead rotates the input: it unfolds the input into overlapping patches with stride 1, rotates each patch, and folds them back with F.fold. F.fold is the adjoint of the unfold, not its inverse, so it sums the overlaps. For any kernel_size > 1 the composition fold(unfold(x)) scales each pixel by the number of patches that cover it, so it is not x. The rotation is the identity at initialisation, yet a freshly created OFT-Conv2d adapter already changes the model output, and merge/merge_and_unload, which rotates the flattened filter directly, disagrees with the forward by ~100 percent. Both are silent.

The code

peft, src/peft/tuners/oft/layer.py, 0.19.1 and current main (byte-identical):

class OFTRotationModule(nn.Module):
    def _unfold(self, x):                       # stride 1, overlapping patches (lines 189-213)
        x_unfolded = x.unfold(2, kh, 1).unfold(3, kw, 1)
        ...
    def _fold(self, x_unfolded, orig_shape):    # F.fold sums the overlaps (lines 215-244)
        x_folded = F.fold(..., output_size=(H, W), kernel_size=(kh, kw), stride=(1, 1))
        ...
    def forward(self, x):                       # unfold, rotate, fold (lines 246-285)
        if len(orig_shape) == 4:
            x = self._unfold(x)
        x_rotated_reshaped = einsum("...rk,rkc->...rc", x_reshaped, orth_rotate)
        if len(orig_shape) == 4:
            x_rotated = self._fold(x_rotated, orig_shape)

class Conv2d(...):
    def merge(self, ...):                       # rotates the flattened filter, the correct map (821-876)
        oft_mat = self.get_delta_weight(active_adapter)
        orig_weights = torch.mm(oft_mat, orig_weights.to(oft_mat.dtype))

F.fold is the transpose of the unfold, so fold(unfold(x))[i] = x[i] * (patches covering i). Interior pixels of a k-wide kernel are covered k times, so the composition scales the signal rather than reconstructing it. With the identity rotation (oft_R = 0 at init gives Cayley(0) = I) the Conv2d forward reduces to base_conv(fold(unfold(x))), which is not base_conv(x) for kernel_size > 1. The merge path never unfolds; it rotates the flattened filter directly, so it and the forward compute different functions.

What it does

Driving the real OFT layer (fp64):

root cause |fold(unfold(x)) - x|:
  k=1: 0.0
  k=2: 6.7
  k=3: 12.5
identity-init OFT adapter forward vs base (should be 0 = no-op):
  Conv2d k=1     : 0.0
  Conv2d k=3     : 7.089   <- not a no-op
  Linear control : 0.0
merge vs forward (should be 0):
  Conv2d k=1     : 4.4e-16
  Conv2d k=3     : 6.410   <- merge disagrees with forward

A zero-initialised adapter, whose rotation is exactly the identity, must leave the output unchanged. On a 3x3 Conv2d it shifts the output by 7.089 before any training. Once the rotation is non-trivial, the merged weight and the forward differ by 6.410 on an output of scale about 5, so merge_and_unload produces a different model than the one that was trained. The base cases isolate the fault: a 1x1 Conv2d has no patch overlap and an nn.Linear never unfolds, and both are exact to machine precision. The gap grows with the kernel because the overlap count grows with the kernel.

Scope

The trigger is automatic on the flagship OFT use: OFT and its orthogonal-finetuning line target the convolutional layers of diffusion UNets, which are 3x3, so the default OFTConfig on a Conv2d reaches this path with no warning. The forward is wrong from initialisation and the merged weight does not match it, both silently. The fix is to invert the patch map rather than take its adjoint, for example dividing the folded result by the per-pixel overlap count (fold(unfold(ones))), or to rotate the weight in the forward as the merge path already does.

Layout

  • excerpt.py: the unfold, the fold, the forward, and the merge quoted with the flags that name the fault (Apache-2.0).
  • fold.py: the fold/unfold overlap-sum modelled on a 1D signal, so fold(unfold(x)) equals x scaled by the coverage count and is the identity only for k = 1.
  • consequence.py: the real OFT helpers' fold(unfold(x)) gap, the identity-init adapter versus the base output for a 1x1 and a 3x3 Conv2d and an nn.Linear, and the merge-versus-forward gap.
  • test_oftfold.py: the model folds to the identity only for k = 1 and otherwise scales by the coverage; the real fold(unfold) gap grows with the kernel; the identity-init adapter is a no-op only for the 1x1 conv and the Linear; and merge disagrees with the forward for the 3x3 conv.

Reproduce

python fold.py
python consequence.py
python test_oftfold.py

The unfold, fold, forward, and merge are quoted from current main; the tensors are produced by the real OFT layer. The fix is to make the Conv2d forward invert the patch map, so an identity rotation is a no-op and the forward matches the merged weight.

About

PEFT OFT Conv2d forward folds unfolded patches with F.fold (the adjoint, not the inverse), so for kernel_size>1 a freshly-initialized identity-rotation adapter already changes the output and merge_and_unload diverges from the forward

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages