Online Federation For Mixtures of Proprietary Agents with Black-Box Encoders

Codebase for the paper "Online Federation For Mixtures of Proprietary Agents with Black-Box Encoders"

Codebase Structure

The main code of the algorithm is located in fedmoe. You can re-generate experiment results and artifacts by running the codes in experiments.

Run Experiments

The experimental scripts are all located in the experiments directory. Three agent types are considered: Echo State Network (ESN), Random Feature Network (RFN), and Pre-trained Transformer. You can find the experimental scripts for each dataset under the model experiment directory. Provided README files include the exact commands needed to run each experiment. There are hyperparameter tuning scripts as a part of each experiment. Don't forget to install and activate your virtual environment as described here.

Example Running Script

As mentioned before, provided README files in the experiments directory include the exact commands needed to run each experiment, in addition to the steps required for hyper-parameter tuning and transformer pre-training. Note that the slurm scripts to run the codes on an HPC environment are provided as well as scripts needed to run the code on your local machine. Here is an example of running one of the experiments on your local machine if you already know the value of hyper-parameters. This command runs the ESN experiments. Details such as the number of clients and the dataset name are specified in config.yaml, and its address should be specified as CONFIG_PATH. For more details, please see the ESN experiment documentation.

python -m experiments.esn_experiments.run_esn_experiment \
        --config_path ${CONFIG_PATH} \
        --result_dir ${RUN_OUTPUT_DIR} \
        --hidden_dim ${HIDDEN_DIM} \
        --alpha ${ALPHA_VALUE} \
        --gamma ${GAMMA_VALUE} \
        --sigma ${SIGMA_VALUE} \
        --K ${K_VALUE} \
        --eta ${ETA_VALUE} \
        --client_T ${CLIENT_T_VALUE} \
        --game_sync_freq ${GAME_SYNC} \
        --game_T ${GAME_T_VALUE} \
        --random_seed ${SEED} \
        > ${RUN_OUTPUT_FILE} 2>&1 &

To run other experiments, please see RFN experiment documentation, and Transformer model experiment documentation.

fedmoe Components

• client_manager manages a group of clients or experts and acts as a medium between the clients and the server.
• server implements the main time loop of the algorithm as described in the paper (See Algorithm 1 in the Appendix B).

• datasets: Contains scripts related to generating, loading, processing, and managing the datasets.
• clients: Contains implementations of various client(agents) types, each with specific models and functionalities. These clients represent the agents or experts in our setup.
• models: Includes different model architectures used in the experiments. Each model is designed to integrate with the FedMOE framework and can be tailored for specific tasks.
• game: Implements the core logic of the Nash game, including expectation estimations and computation of matrices.
• state: The state of the clients and the game is saved and managed throughout the time-series algorithm.
• tests: Contains unit tests to ensure the reliability and correctness of the code. These tests cover various components to prevent bugs and maintain code quality.
• utils: Scripts for various tasks, including data preprocessing, visualization of plots using the experiment artifacts (i.e., JSON files), and other reusable code snippets.

Developing

Installing dependencies

The easiest way to set up the development environment:

# Install uv (if not already installed)
curl -LsSf https://astral.sh/uv/install.sh | sh

# Install all dependencies
uv sync --group dev

Note that the --group flag installs optional dependency groups required for the full development workflow. See the pyproject.toml file for details about each group.

If you need to update dependencies, you should change the requirements in pyproject.toml and then update the uv.lock using the command uv lock

Format checks and coding guidelines

To invoke pre-commit hooks, you can install the pre-commit hooks to be run locally. Activate your environment and run:

pre-commit install

To run the checks on all the files:

pre-commit run --all-files

Citation

If you find our work (code or the paper) helpful, please use the citation below.

@article{yang2025onlinefederationmixturesproprietary,
      title={Online Federation For Mixtures of Proprietary Agents with Black-Box Encoders},
      author={Xuwei Yang and Fatemeh Tavakoli and David B. Emerson and Anastasis Kratsios},
      year={2025},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2505.00216},
}