File size: 5,146 Bytes
8f89e30 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | """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)
# Eq. (33): aggregate over space before contracting each query.
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)
|