Views
No views yet
1# x: batch x seq_len x dim
2# router_global: map x -> logits over N_units
3# router_local[u]: map x -> logits over E_per_unit
4
5units_top = topk(router_global(x), U_g) # indices y scores
6outputs = []
7for u in units_top:
8 experts_top = topk(router_local[u](x), k) # per-token or per-position
9 # compute experts in parallel (sparse)
10 expert_outs = [expert_u_e(x) for e in experts_top]
11 # collaborative fusion (weighted sum) + unit adapter
12 weights = softmax(scores_for_experts)
13 fused = sum(w * out for w,out in zip(weights, expert_outs))
14 fused = unit_adapter[u](fused) # small shared layer per unit
15 outputs.append(fused * unit_selection_weight[u]) # combine units
16y = sum(outputs) # final aggregated output
171# %% [markdown]
2# CULO — Demo Notebook (Research-style)
3#
4# Purpose: concise, research-oriented PyTorch notebook implementing the CULO layer
5# (Colaboración de Unidades Locales de Expertos). Includes minimal experiments with
6# synthetic data, diagnostics (unit usage, cooperation loss), and plots. CPU-first.
7#
8# Usage: open with Jupyter / VSCode (Jupytext-friendly). Run sequentially.
9
10# %%
11# Imports and device
12import math
13import torch
14import torch.nn as nn
15import torch.nn.functional as F
16import matplotlib.pyplot as plt
17import random
18from typing import Tuple
19
20# Device: CPU (per user choice). If you want GPU, change device to 'cuda' manually.
21DEVICE = torch.device('cpu')
22print('Device:', DEVICE)
23
24# %% [markdown]
25# Model definition — terse, research-oriented comments.
26# CULO: hierarchical MoE with $N_{units}$ units, each containing $E$ experts.
27# Routing is token-wise. Global router selects up to $U_g$ units; local router selects up to $k$ experts per unit.
28
29# %%
30class Expert(nn.Module):
31 """Simple FFN expert."""
32 def __init__(self, d_model: int, d_hidden: int):
33 super().__init__()
34 self.net = nn.Sequential(
35 nn.Linear(d_model, d_hidden),
36 nn.GELU(),
37 nn.Linear(d_hidden, d_model)
38 )
39 def forward(self, x):
40 return self.net(x)
41
42class CULOLayer(nn.Module):
43 """Compact CULO implementation focused on clarity for research prototyping.
44
45 Key elements:
46 - global_router: token -> N_units logits
47 - local_routers: per-unit token -> E logits
48 - experts: per-unit ModuleList of Expert
49 - unit_adapters: small shared linear per unit
50 - unit_scalars: learned scalar per unit
51 """
52 def __init__(self, d_model: int, d_hidden: int, N_units: int = 8, E_per_unit: int = 4,
53 k:int = 1, U_g:int = 1, coop_coef: float = 0.01):
54 super().__init__()
55 assert 1 <= U_g <= N_units
56 self.d_model = d_model
57 self.N_units = N_units
58 self.E_per_unit = E_per_unit
59 self.k = k
60 self.U_g = U_g
61 self.coop_coef = coop_coef
62
63 self.global_router = nn.Linear(d_model, N_units)
64 self.local_routers = nn.ModuleList([nn.Linear(d_model, E_per_unit) for _ in range(N_units)])
65 self.experts = nn.ModuleList([nn.ModuleList([Expert(d_model, d_hidden) for _ in range(E_per_unit)]) for _ in range(N_units)])
66 self.unit_adapters = nn.ModuleList([nn.Sequential(nn.Linear(d_model, d_model), nn.GELU()) for _ in range(N_units)])
67 self.unit_scalars = nn.Parameter(torch.ones(N_units))
68
69 def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, dict]:
70 # x: (B, S, D)
71 B, S, D = x.shape
72 tokens = x.reshape(B*S, D) # (T, D)
73 T = tokens.shape[0]
74
75 # Global routing
76 g_logits = self.global_router(tokens) # (T, N)
77 g_soft = F.softmax(g_logits, dim=-1)
78 top_vals, top_idx = torch.topk(g_soft, self.U_g, dim=-1) # (T, U_g)
79 mask_units = torch.zeros_like(g_soft)
80 mask_units.scatter_(1, top_idx, top_vals)
81 unit_weights = mask_units / (mask_units.sum(dim=-1, keepdim=True) + 1e-9)
82 usage_per_unit = unit_weights.mean(dim=0) # (N,)
83
84 # Pre-allocate
85 fused_per_unit = tokens.new_zeros((T, self.N_units, D))
86 coop_loss = 0.0
87
88 # Per-unit processing
89 for u in range(self.N_units):
90 l_logits = self.local_routers[u](tokens) # (T, E)
91 l_soft = F.softmax(l_logits, dim=-1)
92 top_e_vals, top_e_idx = torch.topk(l_soft, self.k, dim=-1)
93 mask_experts = torch.zeros_like(l_soft)
94 mask_experts.scatter_(1, top_e_idx, top_e_vals)
95 expert_weights = mask_experts / (mask_experts.sum(dim=-1, keepdim=True) + 1e-9)
96
97 # compute all expert outputs (small-scale demo)
98 expert_outs = []
99 for e in range(self.E_per_unit):
100 out = self.experts[u][e](tokens) # (T, D)
101 expert_outs.append(out.unsqueeze(1))
102 expert_outs = torch.cat(expert_outs, dim=1) # (T, E, D)
103
104 fused = (expert_weights.unsqueeze(-1) * expert_outs).sum(dim=1) # (T, D)
105 adapted = self.unit_adapters[u](fused) # (T, D)
106 scaled = adapted * self.unit_scalars[u] * unit_weights[:, u].unsqueeze(-1)
107 fused_per_unit[:, u, :] = scaled
108
109 # lightweight cooperation proxy: sample tokens and compute variance among selected experts
110 # This proxy is cheap and works for prototyping; replace with more rigorous metric if needed.
111 sample_n = min(32, T)
112 if sample_n > 0:
113 idxs = torch.randperm(T)[:sample_n]
114 topk_for_sample = top_e_idx[idxs] # (sample_n, k)
115 # build mask for selected experts
116 mask_s = torch.zeros((sample_n, self.E_per_unit), device=tokens.device)
117 mask_s.scatter_(1, topk_for_sample, 1.0)
118 mask_s = mask_s.unsqueeze(-1) # (sample_n, E, 1)
119 sampled_expert_outs = expert_outs[idxs] * mask_s # (sample_n, E, D)
120 denom = mask_s.sum(dim=1) + 1e-9
121 mean_selected = sampled_expert_outs.sum(dim=1) / denom
122 diffs = sampled_expert_outs - mean_selected.unsqueeze(1)
123 sq = (diffs ** 2).sum(dim=-1)
124 coop_loss = coop_loss + sq.mean()
125
126 out_tokens = fused_per_unit.sum(dim=1) # (T, D)
127 out = out_tokens.view(B, S, D)
128
129 diagnostics = {
130 'usage_per_unit': usage_per_unit.detach().cpu(),
131 'coop_loss': coop_loss * (self.coop_coef / max(1, self.N_units))
132 }
133
134 return out, diagnostics
135
136# %%
137class DemoModel(nn.Module):
138 """Embedding -> CULO -> residual -> classifier (toy language modeling head)."""
139 def __init__(self, vocab: int = 200, d_model: int = 64, d_hidden: int = 128,
140 N_units: int = 8, E_per_unit: int = 4, k: int = 1, U_g: int = 1):
141 super().__init__()
142 self.embed = nn.Embedding(vocab, d_model)
143 self.culo = CULOLayer(d_model=d_model, d_hidden=d_hidden, N_units=N_units, E_per_unit=E_per_unit, k=k, U_g=U_g)
144 self.norm = nn.LayerNorm(d_model)
145 self.out_proj = nn.Linear(d_model, vocab)
146
147 def forward(self, tokens):
148 x = self.embed(tokens)
149 res, diag = self.culo(x)
150 x = self.norm(x + res)
151 logits = self.out_proj(x)
152 return logits, diag
153
154# %% [markdown]
155# Small reproducible experiment (synthetic data). Hyperparameters set for CPU-friendly run.
156
157# %%
158def run_experiment(steps=50, B=8, S=16, vocab=200,
159 d_model=64, d_hidden=128, N_units=8, E_per_unit=4, k=1, U_g=1,
160 lr=1e-3, print_every=10, seed=42):
161 torch.manual_seed(seed); random.seed(seed)
162 model = DemoModel(vocab=vocab, d_model=d_model, d_hidden=d_hidden, N_units=N_units, E_per_unit=E_per_unit, k=k, U_g=U_g).to(DEVICE)
163 opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4)
164 loss_fn = nn.CrossEntropyLoss()
165
166 history = {'loss': [], 'total_loss': [], 'usage_var': [], 'coop': []}
167
168 for step in range(steps):
169 tokens = torch.randint(0, vocab, (B, S), device=DEVICE)
170 logits, diag = model(tokens)
171 target = torch.randint(0, vocab, (B, S), device=DEVICE)
172 loss = loss_fn(logits.view(-1, vocab), target.view(-1))
173
174 usage = diag['usage_per_unit'].to(DEVICE)
175 l_balance = usage.var()
176 l_coop = diag['coop_loss'].to(DEVICE)
177 total_loss = loss + 0.1 * l_balance + 0.01 * l_coop
178
179 opt.zero_grad(); total_loss.backward(); opt.step()
180
181 history['loss'].append(loss.item()); history['total_loss'].append(total_loss.item())
182 history['usage_var'].append(usage.var().item()); history['coop'].append(l_coop.item())
183
184 if (step + 1) % print_every == 0 or step == 0 or step == steps-1:
185 print(f"step {step+1}/{steps}: loss={loss.item():.4f}, total={total_loss.item():.4f}, usage_var={usage.var().item():.6f}")
186
187 return model, history
188
189# %% [markdown]
190# Run the experiment (small, CPU-friendly). Adjust `steps` for longer runs.
191
192# %%
193if __name__ == '__main__':
194 model, history = run_experiment(steps=60, B=8, S=16, vocab=200, d_model=64, d_hidden=128, N_units=8, E_per_unit=4, k=1, U_g=1)
195
196 # Plot diagnostics
197 fig, axs = plt.subplots(2, 2, figsize=(10, 6))
198 axs = axs.flatten()
199 axs[0].plot(history['loss'])
200 axs[0].set_title('Task loss')
201 axs[1].plot(history['total_loss'])
202 axs[1].set_title('Total loss (with regularizers)')
203 axs[2].plot(history['usage_var'])
204 axs[2].set_title('Unit usage variance')
205 axs[3].plot(history['coop'])
206 axs[3].set_title('Cooperation proxy')
207 plt.tight_layout()
208 plt.show()
209
210# %% [markdown]
211# Notes (research style):
212# - This prototype uses dense computation for all experts (computes every expert's output) and masks
213# via top-k selection. Replace with bucketed / sparse kernels for production efficiency.
214# - The cooperation proxy is a heuristic; for papers consider stronger metrics (e.g., representation-distance on held-out tokens, mutual information, or task-specific agreement metrics).
215# - Hyperparameters to sweep: N_units, E_per_unit, k, U_g, coop_coef, and regularization weights.
216# - For reproducible experiments on GPU, set DEVICE = torch.device('cuda') and ensure deterministic flags as needed.
217
218# End of notebook