Google Titans: 테스트 타임에 기억하는 법을 배우는 Neural Long-Term Memory 아키텍처

# Titans 아키텍처의 Surprise 기반 메모리 업데이트 핵심 로직 (개념 코드)
import torch
import torch.nn as nn

class NeuralMemory(nn.Module):
    def __init__(self, dim, memory_lr=0.01):
        super().__init__()
        # 장기 메모리 = 작은 MLP (key-value 연상 기억)
        self.memory_mlp = nn.Sequential(
            nn.Linear(dim, dim * 2),
            nn.SiLU(),
            nn.Linear(dim * 2, dim)
        )
        self.memory_lr = memory_lr  # 메모리 업데이트 속도

    def compute_surprise(self, query, target):
        """Surprise = 예측 오차의 크기 (gradient norm)"""
        pred = self.memory_mlp(query)
        loss = nn.functional.mse_loss(pred, target)
        grad = torch.autograd.grad(loss, self.memory_mlp.parameters())
        surprise = sum(g.norm() for g in grad)  # 놀라울수록 강하게 기억
        return surprise, loss

    def update_memory(self, key, value, forget_gate):
        """놀라운 정보를 메모리에 기록 (test-time update)"""
        surprise, loss = self.compute_surprise(key, value)
        # Forgetting Gate: 오래된 기억 decay + 새 정보 기록
        effective_lr = self.memory_lr * surprise.item() * forget_gate
        for param in self.memory_mlp.parameters():
            if param.grad is not None:
                param.data -= effective_lr * param.grad

    def recall(self, query):
        """쿼리로 장기 메모리에서 정보 검색"""
        return self.memory_mlp(query)

# MAC (Memory as Context) 변형 사용 예시
# memory_output = neural_memory.recall(current_query)
# context = torch.cat([memory_output, short_term_kv_cache], dim=1)
# output = attention(query, context)  # 장기 + 단기 메모리 통합