File size: 3,935 Bytes
fe31b61 | 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 | """Strong independent-factor controls sharing the same numerical evidence."""
from __future__ import annotations
import torch
from torch.nn.functional import one_hot
from .factorial import hypotheses
def fit(x,h,sizes,kind):
d=x.shape[1];eye=torch.eye(d,device=x.device,dtype=x.dtype);parts=[]
if kind=='factorized_ridge':
z=torch.cat((x,x.new_ones((len(x),1))),-1)
penalty=eye.new_zeros((d+1,d+1));penalty.diagonal().fill_(8.);penalty[-1,-1]=.01
inverse=torch.linalg.inv(z.T@z+penalty)
for i,size in enumerate(sizes):
label=h[:,i];count=torch.bincount(label,minlength=size).double()
sums=x.new_zeros((size,d)).scatter_add_(0,label[:,None].expand_as(x),x)
mean=(sums+x.mean(0))/(count[:,None]+1)
residual=x-mean[label]
scatter=x.new_zeros((size,d,d)).scatter_add_(0,label[:,None,None].expand(-1,d,d),residual[:,:,None]*residual[:,None,:])
pool=scatter.sum(0)/max(1,len(x)-size)
if kind=='factorized_nb':
var=(scatter.diagonal(dim1=-2,dim2=-1)+4*pool.diagonal()[None])/(count[:,None]+4)
covariance=torch.diag_embed(var)+.05*eye
elif kind=='factorized_lda':
covariance=(.5*pool+.5*torch.diag_embed(pool.diagonal())+.05*eye)[None].expand(size,-1,-1)
elif kind=='factorized_kernel':
overall=(x-x.mean(0)).T@(x-x.mean(0))/max(1,len(x)-1)
covariance=(.5*overall+.5*eye)*.7**2
elif kind=='factorized_ridge':
# Class-balanced targets impose the same uniform ontology prior.
target=one_hot(label,size).double()/(count[None]+1)
weights=inverse@z.T@target
covariance=eye
else:raise ValueError(kind)
parts.append(dict(mean=mean,cov=covariance,count=count,
centers=x,labels=label,weights=weights if kind=='factorized_ridge' else None,size=size))
return dict(parts=parts,sizes=tuple(sizes),kind=kind,grid=hypotheses(sizes,x.device),neural_parameters=0)
def posterior(model,x,observed):
ix=list(observed);n=len(x);logs=[]
for part in model['parts']:
size=part['size']
if not ix:
p=x.new_full((n,size),1/size)
elif model['kind']=='factorized_kernel':
cov=part['cov'][ix][:,ix];chol=torch.linalg.cholesky(cov)
means=torch.cat((part['centers'],part['mean']),0)
delta=x[:,None,ix]-means[None,:,ix]
white=torch.linalg.solve_triangular(chol,delta.reshape(-1,len(ix)).T,upper=False)
loglike=-.5*white.square().sum(0).reshape(n,-1)
result=[]
for a in range(size):
components=torch.cat((torch.where(part['labels']==a)[0],part['labels'].new_tensor([len(part['centers'])+a])))
weight=x.new_ones(len(components));weight[-1]=.5
result.append(torch.logsumexp(loglike[:,components]+weight.log()[None],-1)-(part['count'][a]+.5).log())
p=torch.stack(result,-1).softmax(-1)
elif model['kind']=='factorized_ridge':
# Missing coordinates contribute no invented observed value; the full
# data diagnostic comparison uses every coordinate. Acquisition evaluates
# only world models, rather than claiming these are missing-data trained.
scores=x[:,ix]@part['weights'][ix]+part['weights'][-1]
p=scores.clamp_min(1e-4);p/=p.sum(-1,keepdim=True)
else:
cov=part['cov'][:,ix][:,:,ix];chol=torch.linalg.cholesky(cov)
delta=(x[:,None,ix]-part['mean'][None,:,ix]).permute(1,2,0)
white=torch.linalg.solve_triangular(chol,delta,upper=False)
mahal=white.square().sum(1).T
logdet=chol.diagonal(dim1=-2,dim2=-1).log().sum(-1)
p=(-.5*mahal-logdet[None]).softmax(-1)
logs.append(p.clamp_min(1e-300).log()[:,model['grid'][:,len(logs)]])
return torch.stack(logs).sum(0).softmax(-1)
|