Transformer forward changes

[ZFTurbo]

Now forward is like that

def forward(self, x, value_residual = None):
    first_values = None
    for attn, ff in self.layers:
        x, next_values = attn(x, value_residual = value_residual)
        first_values = default(first_values, next_values)
        x = ff(x)

    return self.norm(x), first_values

vs before:

def forward(self, x):

        for attn, ff in self.layers:
            x = attn(x) + x
            x = ff(x) + x

        return self.norm(x)

It broke compatibility with old weights. Was it made intentionally? Is it needed for value residual learning and hyper connections?

[lucidrains]

@ZFTurbo yes, it was for both

i think the weights should be fixed to the version of the repository it was trained on

[ZFTurbo]

I decided to change it like that.

 def forward(self, x, value_residual=None):

        first_values = None
        if value_residual is not None:
            for attn, ff in self.layers:
                x, next_values = attn(x, value_residual=value_residual)
                first_values = default(first_values, next_values)
                x = ff(x)
        else:
            # Compatibility with old weights
            for attn, ff in self.layers:
                attn_out, next_values = attn(x, value_residual=None)
                x = attn_out + x
                x = ff(x) + x

        return self.norm(x), first_values

It works normally with old weights. Can you please say your opinion will it broke something in value residual learning and hyper connections training?

[lucidrains]

@ZFTurbo yes that would work, but that's probably not the best way

i bumped a version when i made these two changes, so i think the weights should just be accompanied by a pinned version of the library pip install BS-Roformer==xx.xx.xx