I was wondering if we could separate TREAD entirely from REPA and asked claude about it. It's probably a bit too complex though ...
# protocols.py (or within repa.py)
from typing import Optional, Protocol, Tuple
import torch
from torch import Tensor
class TokenAdapter(Protocol):
"""Protocol for adapting token representations between algorithms.
Implementations handle token routing/dropping (e.g., TREAD) by providing:
1. Original token count for spatial calculations
2. Feature alignment between source (denoiser) and target (encoder)
"""
def get_original_num_tokens(self, activations: Tensor) -> int:
"""Return original token count for spatial mapping calculations.
Args:
activations: Denoiser activations [B, num_tokens, D]
Returns:
Original number of tokens before any routing/dropping
"""
...
def align_features(
self,
source: Tensor,
target: Tensor,
) -> Tuple[Tensor, Tensor]:
"""Align source and target features for loss computation.
Args:
source: Denoiser activations [B, num_source, D]
target: Encoder features [B, num_target, D]
Returns:
Aligned (source, target) tensors with matching sequence lengths
"""
...
class IdentityAdapter:
"""Default adapter with no token routing."""
def get_original_num_tokens(self, activations: Tensor) -> int:
return activations.shape[-2]
def align_features(
self,
source: Tensor,
target: Tensor,
) -> Tuple[Tensor, Tensor]:
return source, target
class TREADAdapter:
"""Adapter for TREAD token routing.
TREAD must update `original_num_tokens` and `visible_idx` before each forward pass.
"""
def __init__(self) -> None:
self.original_num_tokens: Optional[int] = None
self.visible_idx: Optional[Tensor] = None
def get_original_num_tokens(self, activations: Tensor) -> int:
if self.original_num_tokens is not None:
return self.original_num_tokens
return activations.shape[-2]
def align_features(
self,
source: Tensor,
target: Tensor,
) -> Tuple[Tensor, Tensor]:
if self.visible_idx is None:
return source, target
# target: [B, num_original, D] -> [B, num_visible, D]
hidden_dim = target.shape[-1]
idx = self.visible_idx.unsqueeze(-1).expand(-1, -1, hidden_dim)
target = target.gather(1, idx)
return source, target
def reset(self) -> None:
"""Reset state after forward pass."""
self.original_num_tokens = None
self.visible_idx = None Originally posted by @photoroman in #13 (comment)
I was wondering if we could separate TREAD entirely from REPA and asked claude about it. It's probably a bit too complex though ...
Originally posted by @photoroman in #13 (comment)