Spaces:
Paused
Paused
File size: 5,293 Bytes
20857b0 | 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 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 | """Context modelling and spatial prediction for neural tokens.
Instead of storing absolute token values, predict each token from its
spatial neighbours (left, top, top-left) and store only the residual.
Left neighbour
Top neighbour
Previous frame (temporal)
Residuals concentrate around zero == much lower entropy.
"""
def _zigzag(v: int) -> int:
return (v << 1) ^ (v >> 31)
def _unzigzag(z: int) -> int:
return (z >> 1) ^ -(z & 1)
def context_predict_2d(tokens_2d: list[list[int]]) -> list[int]:
"""Predict tokens using left + top neighbours in a 2D grid.
tokens_2d: list of rows, each row is a list of tokens.
Returns zigzag-encoded residuals.
Prediction: P(x,y) = (left + top) // 2
Residual: R(x,y) = token(x,y) - P(x,y)
First row: top unavailable, use only left neighbour.
First col: left unavailable, use only top neighbour.
[0,0]: stored as-is.
"""
if not tokens_2d or not tokens_2d[0]:
return []
rows = len(tokens_2d)
cols = len(tokens_2d[0])
residuals: list[int] = []
for r in range(rows):
for c in range(cols):
tok = tokens_2d[r][c]
if r == 0 and c == 0:
pred = 0
elif r == 0:
pred = tokens_2d[r][c - 1]
elif c == 0:
pred = tokens_2d[r - 1][c]
else:
pred = (tokens_2d[r][c - 1] + tokens_2d[r - 1][c]) // 2
residual = tok - pred
residuals.append(_zigzag(residual))
return residuals
def context_unpredict_2d(residuals_zigzag: list[int], rows: int, cols: int) -> list[list[int]]:
"""Reverse context_predict_2d: residuals -> original tokens."""
tokens_2d: list[list[int]] = [[0] * cols for _ in range(rows)]
idx = 0
for r in range(rows):
for c in range(cols):
res = _unzigzag(residuals_zigzag[idx])
if r == 0 and c == 0:
pred = 0
elif r == 0:
pred = tokens_2d[r][c - 1]
elif c == 0:
pred = tokens_2d[r - 1][c]
else:
pred = (tokens_2d[r][c - 1] + tokens_2d[r - 1][c]) // 2
tokens_2d[r][c] = pred + res
idx += 1
return tokens_2d
def context_predict_temporal(frame_tokens: list[list[int]]) -> list[int]:
"""Predict tokens from previous frame + spatial neighbours.
Prediction: P(t,x,y) = token(t-1, x, y) + spatial_correction
where spatial_correction = (left + top - top_left) // 3
This is a simple learned-adjacent predictor:
- Temporal prediction handles stationary / slow-moving content
- Spatial prediction handles edges / gradients
Returns zigzag residuals in row-major order.
"""
if not frame_tokens:
return []
rows = len(frame_tokens[0])
cols = len(frame_tokens[0][0])
residuals: list[int] = []
prev_frame: list[list[int]] = [[0] * cols for _ in range(rows)]
for t, frame in enumerate(frame_tokens):
for r in range(rows):
for c in range(cols):
tok = frame[r][c]
temporal_pred = prev_frame[r][c] if t > 0 else 0
if r == 0 and c == 0:
spatial_correction = 0
elif r == 0:
spatial_correction = (frame[r][c - 1] - prev_frame[r][c - 1]) // 2
elif c == 0:
spatial_correction = (frame[r - 1][c] - prev_frame[r - 1][c]) // 2
else:
left_delta = frame[r][c - 1] - prev_frame[r][c - 1]
top_delta = frame[r - 1][c] - prev_frame[r - 1][c]
spatial_correction = (left_delta + top_delta) // 2
pred = temporal_pred + spatial_correction
residual = tok - pred
residuals.append(_zigzag(residual))
prev_frame = [row[:] for row in frame]
return residuals
def context_unpredict_temporal(residuals_zigzag: list[int], rows: int, cols: int,
n_frames: int) -> list[list[list[int]]]:
"""Reverse context_predict_temporal."""
frames: list[list[list[int]]] = []
prev_frame: list[list[int]] = [[0] * cols for _ in range(rows)]
idx = 0
for _ in range(n_frames):
frame: list[list[int]] = [[0] * cols for _ in range(rows)]
for r in range(rows):
for c in range(cols):
res = _unzigzag(residuals_zigzag[idx])
temporal_pred = prev_frame[r][c]
if r == 0 and c == 0:
spatial_correction = 0
elif r == 0:
spatial_correction = (frame[r][c - 1] - prev_frame[r][c - 1]) // 2
elif c == 0:
spatial_correction = (frame[r - 1][c] - prev_frame[r - 1][c]) // 2
else:
left_delta = frame[r][c - 1] - prev_frame[r][c - 1]
top_delta = frame[r - 1][c] - prev_frame[r - 1][c]
spatial_correction = (left_delta + top_delta) // 2
pred = temporal_pred + spatial_correction
frame[r][c] = pred + res
idx += 1
frames.append(frame)
prev_frame = [row[:] for row in frame]
return frames
|