| """Independent implementation of LiNO, arXiv:2606.03262, Eqs. (5)-(44). |
| |
| Only the efficient scattering variant for one-dimensional fields is implemented. |
| Unspecified numerical choices are documented in results/reproduction_spec.md. |
| """ |
|
|
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
|
|
| class AdaptiveMirror(nn.Module): |
| def __init__(self, width, eps): |
| super().__init__() |
| self.direction = nn.Linear(width, width) |
| self.eps = eps |
|
|
| def forward(self, field): |
| normal = F.normalize(self.direction(field), dim=-1, eps=self.eps) |
| component = (field * normal).sum(-1, keepdim=True) |
| return field - 2 * component * normal |
|
|
|
|
| class AdaptiveRefraction(nn.Module): |
| def __init__(self, width, rate_range, eps): |
| super().__init__() |
| self.direction = nn.Linear(width, width) |
| self.contrast = nn.Parameter(torch.zeros(())) |
| self.rate_range, self.eps = rate_range, eps |
|
|
| def forward(self, field): |
| normal = F.normalize(self.direction(field), dim=-1, eps=self.eps) |
| normal_component = (field * normal).sum(-1, keepdim=True) * normal |
| return field + self.rate_range * self.contrast.tanh() * normal_component |
|
|
|
|
| class EfficientScattering(nn.Module): |
| def __init__(self, width, eps): |
| super().__init__() |
| self.query = nn.Linear(width, width, bias=False) |
| self.key = nn.Linear(width, width, bias=False) |
| self.value = nn.Linear(width, width, bias=False) |
| self.coordinates = nn.Sequential(nn.Linear(1, width), nn.GELU(), nn.Linear(width, width)) |
| self.diffusion = nn.Conv1d(width, width, 3, padding=1, groups=width, bias=False) |
| self.local_projection = nn.Linear(width, width, bias=False) |
| self.local_logit = nn.Parameter(torch.tensor(-1.0)) |
| self.log_strength = nn.Parameter(torch.tensor(-2.0)) |
| self.scale, self.eps = width ** -0.5, eps |
|
|
| def forward(self, field, grid): |
| position = self.coordinates(grid) |
| query = F.elu(self.query(field) * self.scale + position) + 1 |
| key = F.elu(self.key(field) * self.scale + position) + 1 |
| value = self.value(field) |
| |
| summary = key.transpose(1, 2) @ value |
| numerator = query @ summary |
| denominator = (query * key.sum(1, keepdim=True)).sum(-1, keepdim=True) |
| global_average = numerator / (denominator + self.eps) |
| local_field = self.diffusion(field.transpose(1, 2)).transpose(1, 2) |
| local_average = self.local_projection(local_field) |
| beta = self.local_logit.sigmoid() |
| average = (1 - beta) * global_average + beta * local_average |
| return self.log_strength.exp() * (average - field) |
|
|
|
|
| class LightBlock(nn.Module): |
| def __init__(self, width, rate_range, eps): |
| super().__init__() |
| self.before_light = nn.LayerNorm(width) |
| self.before_ffn = nn.LayerNorm(width) |
| self.mirror = AdaptiveMirror(width, eps) |
| self.refraction = AdaptiveRefraction(width, rate_range, eps) |
| self.scattering = EfficientScattering(width, eps) |
| self.gate = nn.Sequential(nn.Linear(width, 2 * width), nn.GELU(), nn.Linear(2 * width, 3)) |
| self.mix_projection = nn.Linear(width, width) |
| self.ffn = nn.Sequential(nn.Linear(width, 2 * width), nn.GELU(), nn.Linear(2 * width, width)) |
|
|
| def forward(self, field, grid): |
| normalized = self.before_light(field) |
| weights = self.gate(normalized.mean(1)).softmax(-1) |
| branches = (self.mirror(normalized), self.refraction(normalized), self.scattering(normalized, grid)) |
| mixture = sum(weights[:, i, None, None] * branch for i, branch in enumerate(branches)) |
| field = field + self.mix_projection(mixture) |
| return field + self.ffn(self.before_ffn(field)) |
|
|
|
|
| class LiNO(nn.Module): |
| """[batch, space, 2] -> [batch, space, 1], normalized field and x input.""" |
|
|
| def __init__(self, in_channels=2, out_channels=1, width=128, depth=8, |
| scattering="efficient", rate_range=0.25, eps=1e-6): |
| super().__init__() |
| if scattering != "efficient": |
| raise ValueError("This experiment requires efficient scattering") |
| self.architecture = dict(in_channels=in_channels, out_channels=out_channels, |
| width=width, depth=depth, scattering=scattering, |
| rate_range=rate_range, eps=eps) |
| self.lift = nn.Linear(in_channels, width) |
| self.blocks = nn.ModuleList([LightBlock(width, rate_range, eps) for _ in range(depth)]) |
| self.project = nn.Linear(width, out_channels) |
|
|
| def forward(self, inputs): |
| if inputs.ndim != 3 or inputs.shape[-1] != self.architecture["in_channels"]: |
| raise ValueError(f"Expected [B,N,{self.architecture['in_channels']}], got {inputs.shape}") |
| grid = torch.linspace(0, 1, inputs.shape[1], device=inputs.device, dtype=inputs.dtype).view(1, -1, 1) |
| field = self.lift(inputs) |
| for block in self.blocks: |
| field = block(field, grid) |
| return self.project(field) |
|
|
