ITADN

Add support for custom conditional adjacency

#64Closedaalmodovares 创建于 2025-04-12
enhancement
A
aalmodovarescommented
### Description Zuko currently supports structured dependencies via an adjacency matrix in MaskedAutoregressiveTransform ([#54](https://github.com/probabilists/zuko/pull/54)). This functionality is valuable for leveraging structural inductive biases. However, the context adjacency—modulating how context conditions each variable—is fixed to an all-ones matrix, implying full context-to-variable conditioning. In [our recent work](https://arxiv.org/abs/2503.15114), we show that structured context adjacencies can be really useful. We propose extending the current implementation to support user-defined context adjacency matrices. If this aligns with Zuko’s design goals, I would be glad to contribute the implementation via a pull request. We have already implemented this internally for our paper, though the code is not yet public. ### Implementation Currently, the adjacency matrix is (D × D), with an implicit (C × D) all-ones matrix added when context is used. We propose accepting an extended (D × [D + C]) matrix, where the last C columns specify context-to-variable dependencies. The original behavior is preserved when the matrix is (D × D). The modification is minimal, localized to MaskedAutoregressiveTransform, and fully backward-compatible. ### Alternatives No alternatives are currently considered.
关闭于 2025-04-13 1 条评论