聯邦學習：聯邦異構知識圖譜劃分
1 導引
我們在部落格《聯邦學習：聯邦場景下的多源知識圖譜嵌入》中介紹了聯邦場景下的知識圖譜嵌入，現在讓我們回顧一下其中關於數據部分的細節。在聯邦場景下，\(C\)個知識圖譜\(\left\{\mathcal{G}_c\right\}_{c=1}^C=\left\{\left\{\mathcal{E}_c, \mathcal{R}_c, \mathcal{T}_c\right\}\right\}_{c=1}^C\)位於不同的客戶端上。知識圖譜擁的實體集合\(\mathcal{E}_c\)之間可能會存在重疊，而其關係集合\(\mathcal{R}_c\)和元組集合\(\mathcal{T}_c\)之間則不會重疊[1]。我們聯繫一下現實場景看這是合理的，比如在不同客戶端對應不同銀行的情況下，由於不同銀行都有著自己的業務流程，所以關係集合不重疊。
接下來我們來看具體在實驗環節怎麼去劃分聯邦異構知識圖譜數據。
2 聯邦異構知識圖譜劃分
我們在部落格《分散式機器學習：PageRank演算法的並行化實現（PySpark）》中所說，分散式圖數據的劃分可分為點劃分和邊劃分兩種，邊劃分是對圖中某些邊進行分裂，這使得不同的worker的點不同，但可能存有相同的邊拷貝。而點劃分是對圖中某些點進行分裂，使得不同的worker的邊不同，可能存有相同的點拷貝。不過知識圖譜的情況要簡化得多，因為知識圖譜的圖數據本身就是按一條條的邊(元組)\((h,r,t)\)來存儲的，直接對元組進行劃分其實就等價於了點劃分的方式。
由於我們這裡的本地知識圖譜\(\{\mathcal{E}_c, \mathcal{R}_c, \mathcal{T}_c\}\)中每個知識圖譜的關係\(\mathcal{R}_c\)（即邊的種類）是不同的，我們在劃分元組之前我們需要先對關係進行劃分，然後針對關係劃分的結果來劃分元組。 待元組劃分到本地後，還需要將原有的實體和關係的索引映射到本地索引。最後，再在本地進行訓練/驗證/測試集的拆分。整體數據劃分流程圖如下：
2.1 劃分關係
我們選擇隨機地將關係\(\mathcal{R}\)不重疊地劃分到不同的client上：
random.shuffle(triples)
# triples為元祖集合，大小為 (n_triples, 3)
# 每各元組按(h,t,r)順序存儲
triples = np.concatenate(triples)
# 先根據邊的類型edge_type（即關係類型)將不同的edge_type映射到不同的client_id
edge_types = list(set(triples[:, 2]))
random.shuffle(edge_types)
edge_type_to_cid = {}
n_edge_types_per_client = len(edge_types)//n_clients
for id, edge_type in enumerate(edge_types):
c_id = id // n_edge_types_per_client
if c_id < n_clients - 1:
edge_type_to_cid[edge_type] = c_id
else:
edge_type_to_cid[edge_type] = n_clients - 1
2.2 確定元組劃分
在關係的劃分確定之後，我們可以根據每個元組\((h,r,t)\)中\(r\)的劃分情況來決定該元組的劃分情況。程式碼如下：
# 然後根據edge_type到client_id的映射情況，來將元組triples劃分到不同的client
c_id_triples = [[] for i in range(n_clients)]
for triple in triples:
edge_type = triple[2]
c_id = edge_type_to_cid[edge_type]
c_id_triples[c_id].append(triple.reshape(1, -1))
2.3 索引映射
劃分好元組之後，子圖就確定了，接下來我們還需要將子圖的實體和關係的索引進行重新編號，如下圖所示：
對於具體的局部索引如何安排，我們採用隨機選擇的方式。程式碼如下：
# mapping global indices to local indices
c_id_triples_ori = [[] for i in range(n_clients)]
for c_id in range(n_clients):
triples = np.concatenate(c_id_triples[c_id])
c_id_triples_ori[c_id] = triples
edge_index = triples[:, :2]
edge_type = triples[:, 2]
# map entity indices to local entity indices
index_mapping = {}
entities = list(set(edge_index.flatten()))
random.shuffle(entities)
for index, entity in enumerate(entities):
index_mapping[entity] = index
f = lambda x: index_mapping[x]
f = np.vectorize(f)
client_entity_local_index = f(edge_index)
# map edge indices to local entity indices
index_mapping = {}
edges = copy.deepcopy(list(set((edge_type))))
random.shuffle(edges)
for index, edge in enumerate(edges):
index_mapping[edge] = index
f = lambda x: index_mapping[x]
f = np.vectorize(f)
client_edge_local_index = f(edge_type)
c_id_triples[c_id] = np.concatenate([client_entity_local_index, \
client_edge_local_index.reshape(-1, 1)], axis=1)
2.4 訓練/驗證/測試集拆分
最後，還需要在本地劃分訓練集、驗證集和測試集。如下面的程式碼展示了按照0.8/0.1/0.1對本地的元組進行拆分。數據集劃分完畢之後，則訓練/驗證/測試集對應的實體(edge_index)和關係類型(edge_type)就都確立了：
# split train, valid, test dataset
for c_id in range(n_clients):
n_triples = c_id_triples[c_id].shape[0]
n_train = int(n_triples * 0.8)
n_val = int((n_triples - n_train) * 0.5)
n_test = n_triples - n_train - n_val
mod_to_slice = {"train": slice(0, n_train), \
"valid": slice(n_train, n_train+n_val), "test": slice(-n_test, n_triples)}
for mode in ["train", "valid", "test"]:
client_data[c_id][mode]["edge_index_ori"] = c_id_triples_ori[c_id][mod_to_slice[mode], : 2].T
client_data[c_id][mode]["edge_index"] = c_id_triples[c_id][mod_to_slice[mode], : 2].T
client_data[c_id][mode]["edge_type_ori"] = c_id_triples_ori[c_id][mod_to_slice[mode], 2]
client_data[c_id][mode]["edge_type"] = c_id_triples[c_id][mod_to_slice[mode], 2]
3 關於異構性的分析和解決
根據我們前面的定義，在聯邦場景下不同客戶端的知識圖譜滿足實體重疊，因此在進行聯邦訓練的過程中最簡單的方式就是對重疊實體的embeddings進行平均。但是我們知道，知識圖譜可能本身就具有一定的異構性，因為其中的某個實體可能會擁有著不同的關係路徑[2]，如下圖所示：
在聯邦場景下這種異構性則更加明顯，因為如我們前面所說，同一個實體在不同的client的關係路徑肯定不同，如果只採用本地嵌入的方法，那麼不同的client會映射到不同的嵌入空間。此時，如果對來自不同嵌入空間embeddings直接進行聚合，就會丟失掉許多有用的語義資訊。
如上圖所示[3]，知識圖譜School中的元組表示Bob和Jack的學業資訊，Amazon.com知識圖譜中則表示他們的購物資訊。對於Bob和Jack實體而言，在不同的知識圖譜中他們擁有不同的關係，導致了他們的語義資訊在不同知識圖譜中的差異。
如何對聯邦場景下知識圖譜的異構性進行解決，成為一個必須要考慮的問題，目前文獻[3]已採用對比學習對其進行了一定程度的解決，大家感興趣的可以去閱讀一下。
參考
[1] Chen M, Zhang W, Yuan Z, et al. Fede: Embedding knowledge graphs in federated setting[C]//The 10th International Joint Conference on Knowledge Graphs. 2021: 80-88.
[2] Li Z, Liu H, Zhang Z, et al. Learning knowledge graph embedding with heterogeneous relation attention networks[J]. IEEE Transactions on Neural Networks and Learning Systems, 2021.
[3] Chen M, Zhang W, Yuan Z, et al. Federated knowledge graph completion via embedding-contrastive learning[J]. Knowledge-Based Systems, 2022, 252: 109459.
