File size: 9,546 Bytes
d4fff7c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
"""
Wave Processor — Meaning Resolution via Interference.

Instead of softmax attention, ambiguity is resolved through wave
interference patterns. Constructive interference amplifies the
"correct" interpretation; destructive interference suppresses
alternatives.

This is inspired by how holograms work — information is encoded
in interference patterns, not in individual beams.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from typing import Tuple, Optional

from .particle import ParticleState
from ..utils.math_utils import complex_wave, phase_coherence, resonance_detect


class WaveProcessor(nn.Module):
    """Process information via wave interference.
    
    Key insight: when multiple particles are "in phase" (semantically
    aligned), their waves constructively interfere, amplifying their
    shared meaning. When "out of phase," they destructively interfere,
    suppressing noise.
    
    This replaces the softmax(Q·K^T)·V computation in Transformers.
    """
    
    def __init__(
        self,
        d_semantic: int = 256,
        n_wave_heads: int = 8,
        d_wave: int = 32,
    ):
        super().__init__()
        self.d_semantic = d_semantic
        self.n_heads = n_wave_heads
        self.d_wave = d_wave
        
        # Wave generation — particles emit waves based on their state
        self.wave_generator = nn.Sequential(
            nn.Linear(d_semantic, n_wave_heads * d_wave * 2),  # amplitude + phase
            nn.GELU(),
        )
        
        # Wave reception — particles decode received waves
        self.wave_receiver = nn.Sequential(
            nn.Linear(n_wave_heads * d_wave, d_semantic),
            nn.GELU(),
        )
        
        # Resonance detector — finds standing wave patterns
        self.resonance_net = nn.Sequential(
            nn.Linear(d_wave, 64),
            nn.GELU(),
            nn.Linear(64, 1),
            nn.Sigmoid(),
        )
        
        # Interference decoder — extracts meaning from interference
        self.interference_decoder = nn.Sequential(
            nn.Linear(d_wave * 2, d_wave),
            nn.GELU(),
            nn.Linear(d_wave, d_wave),
        )
    
    def forward(self, particles: ParticleState) -> Tuple[torch.Tensor, dict]:
        """Process particles via wave interference.
        
        Returns:
            output: Processed semantic representation [batch, seq, d_semantic]
            diagnostics: Wave processing diagnostics
        """
        batch, seq_len, _ = particles.semantic.shape
        device = particles.semantic.device
        
        # === Step 1: Generate waves from particles ===
        wave_params = self.wave_generator(particles.semantic)
        wave_params = wave_params.reshape(batch, seq_len, self.n_heads, self.d_wave, 2)
        
        amplitudes = torch.sigmoid(wave_params[..., 0])  # [batch, seq, heads, d_wave]
        phases = wave_params[..., 1] * np.pi  # [batch, seq, heads, d_wave]
        
        # Modulate amplitude by particle's own amplitude
        amplitudes = amplitudes * particles.amplitude.unsqueeze(-1).unsqueeze(-1)
        
        # === Step 2: Create complex waves ===
        waves = amplitudes * torch.exp(1j * phases)  # [batch, seq, heads, d_wave]
        
        # === Step 3: Compute interference patterns ===
        # For each head, compute the total wave field
        # This is the sum of all particle waves at each point
        
        # Total field at each position = sum of all other particles' waves
        # received at this position (with distance-based phase shift)
        
        # Simplified: interference = sum of waves weighted by phase alignment
        # Phase alignment between particles i and j = cos(phase_i - phase_j)
        
        phase_i = phases.unsqueeze(3)  # [batch, seq, 1, heads, d_wave]
        phase_j = phases.unsqueeze(2)  # [batch, 1, seq, heads, d_wave]
        phase_diff = phase_i - phase_j  # [batch, seq, seq, heads, d_wave]
        
        # Interference factor: constructive (+1) or destructive (-1)
        interference = torch.cos(phase_diff)  # [batch, seq, seq, heads, d_wave]
        
        # Amplitude-weighted interference
        amp_i = amplitudes.unsqueeze(3)
        amp_j = amplitudes.unsqueeze(2)
        weighted_interference = interference * amp_i * amp_j
        
        # Sum over source particles → received interference at each position
        received = weighted_interference.sum(dim=2)  # [batch, seq, heads, d_wave]
        
        # === Step 4: Detect resonances ===
        # Resonances are standing wave patterns — stable meanings
        resonance_scores = self.resonance_net(received)  # [batch, seq, heads, d_wave, 1]
        resonance_scores = resonance_scores.squeeze(-1)
        
        # Weight received waves by resonance
        resonant_waves = received * resonance_scores.unsqueeze(-1)
        
        # === Step 5: Decode interference patterns ===
        # Split into real and imaginary parts for processing
        received_real = resonant_waves.real
        received_imag = resonant_waves.imag
        
        # Decode meaning from interference
        decoded = self.interference_decoder(
            torch.cat([received_real, received_imag], dim=-1)
        )  # [batch, seq, heads, d_wave]
        
        # Flatten heads
        decoded = decoded.reshape(batch, seq_len, self.n_heads * self.d_wave)
        
        # Final projection
        output = self.wave_receiver(decoded)
        
        # === Diagnostics ===
        coherence = phase_coherence(particles.phase)
        
        diagnostics = {
            'coherence': coherence.mean().item(),
            'mean_amplitude': amplitudes.mean().item(),
            'resonance_strength': resonance_scores.mean().item(),
        }
        
        return output, diagnostics


class WaveAttention(nn.Module):
    """Wave-based alternative to standard attention.
    
    Instead of: softmax(Q·K^T/√d)·V
    We compute: interference_pattern(waves)·V
    
    Key differences:
    1. No quadratic matrix — waves propagate linearly
    2. Phase-based selection — not softmax
    3. Constructive/destructive interference — not weighted average
    """
    
    def __init__(self, d_model: int = 256, n_heads: int = 8):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_head = d_model // n_heads
        
        # Wave generation (replaces Q, K projections)
        self.emit_proj = nn.Linear(d_model, n_heads * self.d_head * 2)  # amp + phase
        
        # Value projection (same as Transformer — values are values)
        self.v_proj = nn.Linear(d_model, d_model)
        
        # Output projection
        self.out_proj = nn.Linear(d_model, d_model)
        
        # Learned wavelength (how far waves travel)
        self.wavelength = nn.Parameter(torch.ones(n_heads) * 2.0)
    
    def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        """Wave-based attention.
        
        Args:
            x: [batch, seq, d_model]
            mask: [batch, seq, seq] optional attention mask
            
        Returns:
            output: [batch, seq, d_model]
        """
        batch, seq, _ = x.shape
        
        # Generate waves
        wave_params = self.emit_proj(x)
        wave_params = wave_params.reshape(batch, seq, self.n_heads, self.d_head, 2)
        
        amplitudes = torch.sigmoid(wave_params[..., 0])
        phases = wave_params[..., 1] * np.pi
        
        # Values
        V = self.v_proj(x).reshape(batch, seq, self.n_heads, self.d_head)
        
        # Compute interference between all pairs
        # Phase difference
        p_i = phases.unsqueeze(2)  # [batch, seq_i, 1, heads, d_head]
        p_j = phases.unsqueeze(1)  # [batch, 1, seq_j, heads, d_head]
        
        # Distance-based phase shift (waves travel and phase-shift)
        positions = torch.arange(seq, device=x.device, dtype=torch.float32)
        dist = (positions.unsqueeze(0) - positions.unsqueeze(1)).abs()  # [seq, seq]
        
        # Phase shift = 2π * distance / wavelength
        wavelength = self.wavelength.view(1, 1, 1, self.n_heads, 1)
        phase_shift = 2 * np.pi * dist.unsqueeze(-1).unsqueeze(-1) / wavelength
        
        # Total phase difference (emitter phase + travel phase - receiver phase)
        total_phase_diff = p_j + phase_shift - p_i
        
        # Interference factor
        interference = torch.cos(total_phase_diff)  # [batch, seq, seq, heads, d_head]
        
        # Amplitude-weighted
        amp_i = amplitudes.unsqueeze(2)
        amp_j = amplitudes.unsqueeze(1)
        weights = (interference * amp_i * amp_j).sum(dim=-1)  # [batch, seq, seq, heads]
        weights = weights / (self.d_head ** 0.5)
        
        if mask is not None:
            weights = weights.masked_fill(mask.unsqueeze(-1) == 0, float('-inf'))
        
        # Normalize (like softmax but based on wave physics)
        weights = F.softmax(weights, dim=2)
        
        # Apply to values
        V_expanded = V.transpose(1, 2)  # [batch, heads, seq, d_head]
        weights_t = weights.permute(0, 3, 1, 2)  # [batch, heads, seq_i, seq_j]
        output = torch.bmm(weights_t, V_expanded)  # [batch, heads, seq, d_head]
        
        # Reshape and project
        output = output.transpose(1, 2).reshape(batch, seq, self.d_model)
        return self.out_proj(output)